File size: 6,320 Bytes
88e45b9 a54ff08 88e45b9 4b46de4 88e45b9 4b46de4 88e45b9 a54ff08 88e45b9 4b46de4 88e45b9 b038f54 88e45b9 b038f54 88e45b9 b038f54 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 |
import sys
from pathlib import Path
import pytest
import torch
from hydra import compose, initialize
from torch import allclose, float32, isclose, tensor
project_root = Path(__file__).resolve().parent.parent.parent
sys.path.append(str(project_root))
from yolo import Config, NMSConfig, create_model
from yolo.config.config import AnchorConfig
from yolo.utils.bounding_box_utils import (
Anc2Box,
Vec2Box,
bbox_nms,
calculate_iou,
calculate_map,
generate_anchors,
transform_bbox,
)
EPS = 1e-4
@pytest.fixture
def dummy_bboxes():
bbox1 = tensor([[50, 80, 150, 140], [30, 20, 100, 80]], dtype=float32)
bbox2 = tensor([[90, 70, 160, 160], [40, 40, 90, 120]], dtype=float32)
return bbox1, bbox2
def test_calculate_iou_2d(dummy_bboxes):
bbox1, bbox2 = dummy_bboxes
iou = calculate_iou(bbox1, bbox2)
expected_iou = tensor([[0.4138, 0.1905], [0.0096, 0.3226]])
assert iou.shape == (2, 2)
assert allclose(iou, expected_iou, atol=EPS)
def test_calculate_iou_3d(dummy_bboxes):
bbox1, bbox2 = dummy_bboxes
iou = calculate_iou(bbox1[None], bbox2[None])
expected_iou = tensor([[0.4138, 0.1905], [0.0096, 0.3226]])
assert iou.shape == (1, 2, 2)
assert allclose(iou, expected_iou, atol=EPS)
def test_calculate_diou(dummy_bboxes):
bbox1, bbox2 = dummy_bboxes
iou = calculate_iou(bbox1, bbox2, "diou")
expected_diou = tensor([[0.3816, 0.0943], [-0.2048, 0.2622]])
assert iou.shape == (2, 2)
assert allclose(iou, expected_diou, atol=EPS)
def test_calculate_ciou(dummy_bboxes):
bbox1, bbox2 = dummy_bboxes
iou = calculate_iou(bbox1, bbox2, metrics="ciou")
# TODO: check result!
expected_ciou = tensor([[0.3769, 0.0853], [-0.2050, 0.2602]])
assert iou.shape == (2, 2)
assert allclose(iou, expected_ciou, atol=EPS)
bbox1 = tensor([[50, 80, 150, 140], [30, 20, 100, 80]], dtype=float32)
bbox2 = tensor([[90, 70, 160, 160], [40, 40, 90, 120]], dtype=float32)
def test_transform_bbox_xywh_to_Any(dummy_bboxes):
bbox1, _ = dummy_bboxes
transformed_bbox = transform_bbox(bbox1, "xywh -> xyxy")
expected_bbox = tensor([[50.0, 80.0, 200.0, 220.0], [30.0, 20.0, 130.0, 100.0]])
assert allclose(transformed_bbox, expected_bbox)
def test_transform_bbox_xycwh_to_Any(dummy_bboxes):
bbox1, bbox2 = dummy_bboxes
transformed_bbox = transform_bbox(bbox1, "xycwh -> xycwh")
assert allclose(transformed_bbox, bbox1)
transformed_bbox = transform_bbox(bbox2, "xyxy -> xywh")
expected_bbox = tensor([[90.0, 70.0, 70.0, 90.0], [40.0, 40.0, 50.0, 80.0]])
assert allclose(transformed_bbox, expected_bbox)
def test_transform_bbox_xyxy_to_Any(dummy_bboxes):
bbox1, bbox2 = dummy_bboxes
transformed_bbox = transform_bbox(bbox1, "xyxy -> xyxy")
assert allclose(transformed_bbox, bbox1)
transformed_bbox = transform_bbox(bbox2, "xyxy -> xycwh")
expected_bbox = tensor([[125.0, 115.0, 70.0, 90.0], [65.0, 80.0, 50.0, 80.0]])
assert allclose(transformed_bbox, expected_bbox)
def test_transform_bbox_invalid_format(dummy_bboxes):
bbox, _ = dummy_bboxes
# Test invalid input format
with pytest.raises(ValueError, match="Invalid input or output format"):
transform_bbox(bbox, "invalid->xyxy")
# Test invalid output format
with pytest.raises(ValueError, match="Invalid input or output format"):
transform_bbox(bbox, "xywh->invalid")
def test_generate_anchors():
image_size = [256, 256]
strides = [8, 16, 32]
anchors, scalers = generate_anchors(image_size, strides)
assert anchors.shape[0] == scalers.shape[0]
assert anchors.shape[1] == 2
def test_vec2box_autoanchor():
with initialize(config_path="../../yolo/config", version_base=None):
cfg: Config = compose(config_name="config", overrides=["model=v9-m"])
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = create_model(cfg.model, weight_path=None).to(device)
vec2box = Vec2Box(model, cfg.model.anchor, cfg.image_size, device)
assert vec2box.strides == [8, 16, 32]
vec2box.update((320, 640))
assert vec2box.anchor_grid.shape == (4200, 2)
assert vec2box.scaler.shape == tuple([4200])
def test_anc2box_autoanchor(inference_v7_cfg: Config):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = create_model(inference_v7_cfg.model, weight_path=None).to(device)
anchor_cfg: AnchorConfig = inference_v7_cfg.model.anchor.copy()
del anchor_cfg.strides
anc2box = Anc2Box(model, anchor_cfg, inference_v7_cfg.image_size, device)
assert anc2box.strides == [8, 16, 32]
anc2box.update((320, 640))
anchor_grids_shape = [anchor_grid.shape for anchor_grid in anc2box.anchor_grids]
assert anchor_grids_shape == [
torch.Size([1, 1, 80, 80, 2]),
torch.Size([1, 1, 40, 40, 2]),
torch.Size([1, 1, 20, 20, 2]),
]
assert anc2box.anchor_scale.shape == torch.Size([3, 1, 3, 1, 1, 2])
def test_bbox_nms():
cls_dist = tensor(
[[[0.1, 0.7, 0.2], [0.6, 0.3, 0.1]], [[0.4, 0.4, 0.2], [0.5, 0.4, 0.1]]] # Example class distribution
)
bbox = tensor(
[[[50, 50, 100, 100], [60, 60, 110, 110]], [[40, 40, 90, 90], [70, 70, 120, 120]]], # Example bounding boxes
dtype=float32,
)
nms_cfg = NMSConfig(min_confidence=0.5, min_iou=0.5)
expected_output = [
tensor(
[
[1.0000, 50.0000, 50.0000, 100.0000, 100.0000, 0.6682],
[0.0000, 60.0000, 60.0000, 110.0000, 110.0000, 0.6457],
]
)
]
output = bbox_nms(cls_dist, bbox, nms_cfg)
for out, exp in zip(output, expected_output):
assert allclose(out, exp, atol=1e-4), f"Output: {out} Expected: {exp}"
def test_calculate_map():
predictions = tensor([[0, 60, 60, 160, 160, 0.5], [0, 40, 40, 120, 120, 0.5]]) # [class, x1, y1, x2, y2]
ground_truths = tensor([[0, 50, 50, 150, 150], [0, 30, 30, 100, 100]]) # [class, x1, y1, x2, y2]
mAP = calculate_map(predictions, ground_truths)
expected_ap50 = tensor(0.5)
expected_ap50_95 = tensor(0.2)
assert isclose(mAP["mAP.5"], expected_ap50, atol=1e-5), f"AP50 mismatch"
assert isclose(mAP["mAP.5:.95"], expected_ap50_95, atol=1e-5), f"Mean AP mismatch"
|