| import torch |
| from test_model import tiny_config |
|
|
| from objectmodel_v1.losses import ObjectModelCriterion |
|
|
|
|
| def reference_dense_targets(criterion, logits, targets, level_index): |
| batch, _, height, width = logits.shape |
| target_logits = torch.zeros_like(logits) |
| target_boxes = torch.zeros(batch, 4, height, width, device=logits.device) |
| positive = torch.zeros(batch, height, width, device=logits.device) |
| for batch_index, target in enumerate(targets): |
| if target["labels"].numel() == 0: |
| continue |
| centers = target["boxes"][:, :2] |
| grid = (centers * torch.tensor([width, height], device=centers.device)).long() |
| grid[:, 0].clamp_(0, width - 1) |
| grid[:, 1].clamp_(0, height - 1) |
| areas = target["boxes"][:, 2] * target["boxes"][:, 3] |
| target_levels = torch.where(areas < 0.02, 0, torch.where(areas < 0.15, 1, 2)) |
| for target_index in torch.where(target_levels == level_index)[0]: |
| center_x, center_y = grid[target_index] |
| label = target["labels"][target_index] |
| candidates = [] |
| for dy in (-1, 0, 1): |
| for dx in (-1, 0, 1): |
| x = int((center_x + dx).clamp(0, width - 1)) |
| y = int((center_y + dy).clamp(0, height - 1)) |
| candidates.append((dx * dx + dy * dy, x, y)) |
| for _, x, y in sorted(candidates)[: criterion.dense_topk]: |
| target_logits[batch_index, label, y, x] = 1.0 |
| target_boxes[batch_index, :, y, x] = target["boxes"][target_index] |
| positive[batch_index, y, x] = 1.0 |
| return target_logits, target_boxes, positive |
|
|
|
|
| def test_dense_targets_match_reference_with_boundaries_and_collisions(): |
| criterion = ObjectModelCriterion(tiny_config()) |
| logits = torch.randn(2, 5, 8, 8) |
| targets = [ |
| { |
| "boxes": torch.tensor( |
| [ |
| [0.01, 0.01, 0.10, 0.10], |
| [0.02, 0.02, 0.11, 0.11], |
| [0.99, 0.99, 0.50, 0.50], |
| ] |
| ), |
| "labels": torch.tensor([1, 2, 3]), |
| }, |
| {"boxes": torch.empty(0, 4), "labels": torch.empty(0, dtype=torch.long)}, |
| ] |
|
|
| for level_index in range(3): |
| expected = reference_dense_targets(criterion, logits, targets, level_index) |
| actual = criterion._dense_targets(logits, targets, level_index) |
| for expected_tensor, actual_tensor in zip(expected, actual, strict=True): |
| assert torch.equal(expected_tensor, actual_tensor) |
|
|
|
|
| def test_dense_targets_match_reference_for_production_topk(): |
| config = tiny_config() |
| config["loss"]["dense_topk"] = 5 |
| criterion = ObjectModelCriterion(config) |
| logits = torch.randn(2, 5, 12, 10) |
| targets = [ |
| {"boxes": torch.rand(9, 4), "labels": torch.randint(0, 5, (9,))}, |
| {"boxes": torch.rand(4, 4), "labels": torch.randint(0, 5, (4,))}, |
| ] |
|
|
| for level_index in range(3): |
| expected = reference_dense_targets(criterion, logits, targets, level_index) |
| actual = criterion._dense_targets(logits, targets, level_index) |
| for expected_tensor, actual_tensor in zip(expected, actual, strict=True): |
| assert torch.equal(expected_tensor, actual_tensor) |
|
|
