✅ [Add] tests, increase test coverage
Browse files- requirements.txt +1 -0
- tests/test_model/test_module.py +0 -7
- tests/test_model/test_yolo.py +11 -1
- tests/{test_utils → test_tools}/test_data_augmentation.py +0 -0
- tests/{test_utils → test_tools}/test_loss_functions.py +0 -1
- tests/test_tools/test_solver.py +114 -0
- tests/test_utils/test_bounding_box_utils.py +163 -0
- tests/{test_tools → test_utils}/test_module_utils.py +37 -1
- yolo/config/config.py +0 -3
- yolo/tools/solver.py +1 -0
- yolo/utils/bounding_box_utils.py +3 -3
- yolo/utils/module_utils.py +1 -2
requirements.txt
CHANGED
@@ -5,6 +5,7 @@ loguru
|
|
5 |
numpy
|
6 |
opencv-python
|
7 |
Pillow
|
|
|
8 |
pytest
|
9 |
pyyaml
|
10 |
requests
|
|
|
5 |
numpy
|
6 |
opencv-python
|
7 |
Pillow
|
8 |
+
pycocotools
|
9 |
pytest
|
10 |
pyyaml
|
11 |
requests
|
tests/test_model/test_module.py
CHANGED
@@ -43,13 +43,6 @@ def test_adown():
|
|
43 |
assert out.shape == (1, OUT_CHANNELS, 32, 32)
|
44 |
|
45 |
|
46 |
-
def test_adown():
|
47 |
-
adown = ADown(IN_CHANNELS, OUT_CHANNELS)
|
48 |
-
x = torch.randn(1, IN_CHANNELS, 64, 64)
|
49 |
-
out = adown(x)
|
50 |
-
assert out.shape == (1, OUT_CHANNELS, 32, 32)
|
51 |
-
|
52 |
-
|
53 |
def test_cblinear():
|
54 |
cblinear = CBLinear(IN_CHANNELS, [5, 5])
|
55 |
x = torch.randn(1, IN_CHANNELS, 64, 64)
|
|
|
43 |
assert out.shape == (1, OUT_CHANNELS, 32, 32)
|
44 |
|
45 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
46 |
def test_cblinear():
|
47 |
cblinear = CBLinear(IN_CHANNELS, [5, 5])
|
48 |
x = torch.randn(1, IN_CHANNELS, 64, 64)
|
tests/test_model/test_yolo.py
CHANGED
@@ -16,7 +16,7 @@ config_path = "../../yolo/config"
|
|
16 |
config_name = "config"
|
17 |
|
18 |
|
19 |
-
def
|
20 |
with initialize(config_path=config_path, version_base=None):
|
21 |
cfg: Config = compose(config_name=config_name)
|
22 |
|
@@ -26,6 +26,16 @@ def test_build_model():
|
|
26 |
assert len(model.model) == 39
|
27 |
|
28 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
29 |
@pytest.fixture
|
30 |
def cfg() -> Config:
|
31 |
with initialize(config_path="../../yolo/config", version_base=None):
|
|
|
16 |
config_name = "config"
|
17 |
|
18 |
|
19 |
+
def test_build_model_v9c():
|
20 |
with initialize(config_path=config_path, version_base=None):
|
21 |
cfg: Config = compose(config_name=config_name)
|
22 |
|
|
|
26 |
assert len(model.model) == 39
|
27 |
|
28 |
|
29 |
+
def test_build_model_v9m():
|
30 |
+
with initialize(config_path=config_path, version_base=None):
|
31 |
+
cfg: Config = compose(config_name=config_name, overrides=[f"model=v9-m"])
|
32 |
+
|
33 |
+
OmegaConf.set_struct(cfg.model, False)
|
34 |
+
cfg.weight = None
|
35 |
+
model = YOLO(cfg.model)
|
36 |
+
assert len(model.model) == 39
|
37 |
+
|
38 |
+
|
39 |
@pytest.fixture
|
40 |
def cfg() -> Config:
|
41 |
with initialize(config_path="../../yolo/config", version_base=None):
|
tests/{test_utils → test_tools}/test_data_augmentation.py
RENAMED
File without changes
|
tests/{test_utils → test_tools}/test_loss_functions.py
RENAMED
@@ -31,7 +31,6 @@ def model(cfg: Config):
|
|
31 |
@pytest.fixture
|
32 |
def vec2box(cfg: Config, model):
|
33 |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
34 |
-
print(device)
|
35 |
return Vec2Box(model, cfg.image_size, device)
|
36 |
|
37 |
|
|
|
31 |
@pytest.fixture
|
32 |
def vec2box(cfg: Config, model):
|
33 |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
34 |
return Vec2Box(model, cfg.image_size, device)
|
35 |
|
36 |
|
tests/test_tools/test_solver.py
ADDED
@@ -0,0 +1,114 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import sys
|
2 |
+
from pathlib import Path
|
3 |
+
from unittest.mock import MagicMock, patch
|
4 |
+
|
5 |
+
import pytest
|
6 |
+
import torch
|
7 |
+
from hydra import compose, initialize
|
8 |
+
|
9 |
+
project_root = Path(__file__).resolve().parent.parent.parent
|
10 |
+
sys.path.append(str(project_root))
|
11 |
+
|
12 |
+
from yolo.config.config import (
|
13 |
+
Config,
|
14 |
+
DataConfig,
|
15 |
+
LossConfig,
|
16 |
+
TrainConfig,
|
17 |
+
ValidationConfig,
|
18 |
+
)
|
19 |
+
from yolo.model.yolo import YOLO, create_model
|
20 |
+
from yolo.tools.data_loader import create_dataloader
|
21 |
+
from yolo.tools.loss_functions import create_loss_function
|
22 |
+
from yolo.tools.solver import ( # Adjust the import to your module
|
23 |
+
ModelTester,
|
24 |
+
ModelTrainer,
|
25 |
+
ModelValidator,
|
26 |
+
)
|
27 |
+
from yolo.utils.bounding_box_utils import Vec2Box
|
28 |
+
from yolo.utils.logging_utils import ProgressLogger
|
29 |
+
from yolo.utils.model_utils import (
|
30 |
+
ExponentialMovingAverage,
|
31 |
+
create_optimizer,
|
32 |
+
create_scheduler,
|
33 |
+
)
|
34 |
+
|
35 |
+
|
36 |
+
@pytest.fixture
|
37 |
+
def cfg() -> Config:
|
38 |
+
with initialize(config_path="../../yolo/config", version_base=None):
|
39 |
+
cfg: Config = compose(config_name="config")
|
40 |
+
cfg.weight = None
|
41 |
+
return cfg
|
42 |
+
|
43 |
+
|
44 |
+
@pytest.fixture
|
45 |
+
def cfg_validaion() -> Config:
|
46 |
+
with initialize(config_path="../../yolo/config", version_base=None):
|
47 |
+
cfg: Config = compose(config_name="config", overrides=["task=validation"])
|
48 |
+
cfg.weight = None
|
49 |
+
return cfg
|
50 |
+
|
51 |
+
|
52 |
+
@pytest.fixture
|
53 |
+
def cfg_inference() -> Config:
|
54 |
+
with initialize(config_path="../../yolo/config", version_base=None):
|
55 |
+
cfg: Config = compose(config_name="config", overrides=["task=inference"])
|
56 |
+
cfg.weight = None
|
57 |
+
return cfg
|
58 |
+
|
59 |
+
|
60 |
+
@pytest.fixture
|
61 |
+
def device() -> torch.device:
|
62 |
+
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
63 |
+
return device
|
64 |
+
|
65 |
+
|
66 |
+
@pytest.fixture
|
67 |
+
def model(cfg: Config, device) -> YOLO:
|
68 |
+
model = create_model(cfg.model, weight_path=None)
|
69 |
+
return model.to(device)
|
70 |
+
|
71 |
+
|
72 |
+
@pytest.fixture
|
73 |
+
def vec2box(cfg: Config, model: YOLO, device) -> Vec2Box:
|
74 |
+
model = create_model(cfg.model, weight_path=None).to(device)
|
75 |
+
vec2box = Vec2Box(model, cfg.image_size, device)
|
76 |
+
return vec2box
|
77 |
+
|
78 |
+
|
79 |
+
@pytest.fixture
|
80 |
+
def progress_logger(cfg: Config):
|
81 |
+
progress_logger = ProgressLogger(cfg, exp_name=cfg.name)
|
82 |
+
return progress_logger
|
83 |
+
|
84 |
+
|
85 |
+
def test_model_trainer_initialization(cfg: Config, model: YOLO, vec2box: Vec2Box, progress_logger, device):
|
86 |
+
trainer = ModelTrainer(cfg, model, vec2box, progress_logger, device, use_ddp=False)
|
87 |
+
assert trainer.model == model
|
88 |
+
assert trainer.device == device
|
89 |
+
assert trainer.optimizer is not None
|
90 |
+
assert trainer.scheduler is not None
|
91 |
+
assert trainer.loss_fn is not None
|
92 |
+
assert trainer.progress == progress_logger
|
93 |
+
|
94 |
+
|
95 |
+
# def test_model_trainer_train_one_batch(config, model, vec2box, progress_logger, device):
|
96 |
+
# trainer = ModelTrainer(config, model, vec2box, progress_logger, device, use_ddp=False)
|
97 |
+
# images = torch.rand(1, 3, 224, 224)
|
98 |
+
# targets = torch.rand(1, 5)
|
99 |
+
# loss_item = trainer.train_one_batch(images, targets)
|
100 |
+
# assert isinstance(loss_item, dict)
|
101 |
+
|
102 |
+
|
103 |
+
def test_model_validator_initialization(cfg_validaion: Config, model: YOLO, vec2box: Vec2Box, progress_logger, device):
|
104 |
+
validator = ModelValidator(cfg_validaion.task, model, vec2box, progress_logger, device)
|
105 |
+
assert validator.model == model
|
106 |
+
assert validator.device == device
|
107 |
+
assert validator.progress == progress_logger
|
108 |
+
|
109 |
+
|
110 |
+
def test_model_tester_initialization(cfg_inference: Config, model: YOLO, vec2box: Vec2Box, progress_logger, device):
|
111 |
+
tester = ModelTester(cfg_inference, model, vec2box, progress_logger, device)
|
112 |
+
assert tester.model == model
|
113 |
+
assert tester.device == device
|
114 |
+
assert tester.progress == progress_logger
|
tests/test_utils/test_bounding_box_utils.py
ADDED
@@ -0,0 +1,163 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import sys
|
2 |
+
from pathlib import Path
|
3 |
+
|
4 |
+
import pytest
|
5 |
+
import torch
|
6 |
+
from hydra import compose, initialize
|
7 |
+
from torch import Tensor, allclose, float32, isclose, nn, tensor
|
8 |
+
|
9 |
+
project_root = Path(__file__).resolve().parent.parent.parent
|
10 |
+
sys.path.append(str(project_root))
|
11 |
+
from yolo import Config, NMSConfig, create_model
|
12 |
+
from yolo.utils.bounding_box_utils import (
|
13 |
+
Vec2Box,
|
14 |
+
bbox_nms,
|
15 |
+
calculate_iou,
|
16 |
+
calculate_map,
|
17 |
+
generate_anchors,
|
18 |
+
transform_bbox,
|
19 |
+
)
|
20 |
+
|
21 |
+
EPS = 1e-4
|
22 |
+
|
23 |
+
|
24 |
+
@pytest.fixture
|
25 |
+
def dummy_bboxes():
|
26 |
+
bbox1 = tensor([[50, 80, 150, 140], [30, 20, 100, 80]], dtype=float32)
|
27 |
+
bbox2 = tensor([[90, 70, 160, 160], [40, 40, 90, 120]], dtype=float32)
|
28 |
+
return bbox1, bbox2
|
29 |
+
|
30 |
+
|
31 |
+
def test_calculate_iou_2d(dummy_bboxes):
|
32 |
+
bbox1, bbox2 = dummy_bboxes
|
33 |
+
iou = calculate_iou(bbox1, bbox2)
|
34 |
+
expected_iou = tensor([[0.4138, 0.1905], [0.0096, 0.3226]])
|
35 |
+
assert iou.shape == (2, 2)
|
36 |
+
assert allclose(iou, expected_iou, atol=EPS)
|
37 |
+
|
38 |
+
|
39 |
+
def test_calculate_iou_3d(dummy_bboxes):
|
40 |
+
bbox1, bbox2 = dummy_bboxes
|
41 |
+
iou = calculate_iou(bbox1[None], bbox2[None])
|
42 |
+
expected_iou = tensor([[0.4138, 0.1905], [0.0096, 0.3226]])
|
43 |
+
assert iou.shape == (1, 2, 2)
|
44 |
+
assert allclose(iou, expected_iou, atol=EPS)
|
45 |
+
|
46 |
+
|
47 |
+
def test_calculate_diou(dummy_bboxes):
|
48 |
+
bbox1, bbox2 = dummy_bboxes
|
49 |
+
iou = calculate_iou(bbox1, bbox2, "diou")
|
50 |
+
expected_diou = tensor([[0.3816, 0.0943], [-0.2048, 0.2622]])
|
51 |
+
|
52 |
+
assert iou.shape == (2, 2)
|
53 |
+
assert allclose(iou, expected_diou, atol=EPS)
|
54 |
+
|
55 |
+
|
56 |
+
def test_calculate_ciou(dummy_bboxes):
|
57 |
+
bbox1, bbox2 = dummy_bboxes
|
58 |
+
iou = calculate_iou(bbox1, bbox2, metrics="ciou")
|
59 |
+
# TODO: check result!
|
60 |
+
expected_ciou = tensor([[0.3769, 0.0853], [-0.2050, 0.2602]])
|
61 |
+
assert iou.shape == (2, 2)
|
62 |
+
assert allclose(iou, expected_ciou, atol=EPS)
|
63 |
+
|
64 |
+
bbox1 = tensor([[50, 80, 150, 140], [30, 20, 100, 80]], dtype=float32)
|
65 |
+
bbox2 = tensor([[90, 70, 160, 160], [40, 40, 90, 120]], dtype=float32)
|
66 |
+
|
67 |
+
|
68 |
+
def test_transform_bbox_xywh_to_Any(dummy_bboxes):
|
69 |
+
bbox1, _ = dummy_bboxes
|
70 |
+
transformed_bbox = transform_bbox(bbox1, "xywh -> xyxy")
|
71 |
+
expected_bbox = tensor([[50.0, 80.0, 200.0, 220.0], [30.0, 20.0, 130.0, 100.0]])
|
72 |
+
assert allclose(transformed_bbox, expected_bbox)
|
73 |
+
|
74 |
+
|
75 |
+
def test_transform_bbox_xycwh_to_Any(dummy_bboxes):
|
76 |
+
bbox1, bbox2 = dummy_bboxes
|
77 |
+
transformed_bbox = transform_bbox(bbox1, "xycwh -> xycwh")
|
78 |
+
assert allclose(transformed_bbox, bbox1)
|
79 |
+
|
80 |
+
transformed_bbox = transform_bbox(bbox2, "xyxy -> xywh")
|
81 |
+
expected_bbox = tensor([[90.0, 70.0, 70.0, 90.0], [40.0, 40.0, 50.0, 80.0]])
|
82 |
+
assert allclose(transformed_bbox, expected_bbox)
|
83 |
+
|
84 |
+
|
85 |
+
def test_transform_bbox_xyxy_to_Any(dummy_bboxes):
|
86 |
+
bbox1, bbox2 = dummy_bboxes
|
87 |
+
transformed_bbox = transform_bbox(bbox1, "xyxy -> xyxy")
|
88 |
+
assert allclose(transformed_bbox, bbox1)
|
89 |
+
|
90 |
+
transformed_bbox = transform_bbox(bbox2, "xyxy -> xycwh")
|
91 |
+
expected_bbox = tensor([[125.0, 115.0, 70.0, 90.0], [65.0, 80.0, 50.0, 80.0]])
|
92 |
+
assert allclose(transformed_bbox, expected_bbox)
|
93 |
+
|
94 |
+
|
95 |
+
def test_transform_bbox_invalid_format(dummy_bboxes):
|
96 |
+
bbox, _ = dummy_bboxes
|
97 |
+
|
98 |
+
# Test invalid input format
|
99 |
+
with pytest.raises(ValueError, match="Invalid input or output format"):
|
100 |
+
transform_bbox(bbox, "invalid->xyxy")
|
101 |
+
|
102 |
+
# Test invalid output format
|
103 |
+
with pytest.raises(ValueError, match="Invalid input or output format"):
|
104 |
+
transform_bbox(bbox, "xywh->invalid")
|
105 |
+
|
106 |
+
|
107 |
+
def test_generate_anchors():
|
108 |
+
image_size = [256, 256]
|
109 |
+
strides = [8, 16, 32]
|
110 |
+
anchors, scalers = generate_anchors(image_size, strides)
|
111 |
+
assert anchors.shape[0] == scalers.shape[0]
|
112 |
+
assert anchors.shape[1] == 2
|
113 |
+
|
114 |
+
|
115 |
+
def test_vec2box_autoanchor():
|
116 |
+
with initialize(config_path="../../yolo/config", version_base=None):
|
117 |
+
cfg: Config = compose(config_name="config", overrides=["model=v9-m"])
|
118 |
+
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
119 |
+
model = create_model(cfg.model, weight_path=None).to(device)
|
120 |
+
vec2box = Vec2Box(model, cfg.image_size, device)
|
121 |
+
assert vec2box.strides == [8, 16, 32]
|
122 |
+
|
123 |
+
vec2box.update((320, 640))
|
124 |
+
assert vec2box.anchor_grid.shape == (4200, 2)
|
125 |
+
assert vec2box.scaler.shape == tuple([4200])
|
126 |
+
|
127 |
+
|
128 |
+
def test_bbox_nms():
|
129 |
+
cls_dist = tensor(
|
130 |
+
[[[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
|
131 |
+
)
|
132 |
+
bbox = tensor(
|
133 |
+
[[[50, 50, 100, 100], [60, 60, 110, 110]], [[40, 40, 90, 90], [70, 70, 120, 120]]], # Example bounding boxes
|
134 |
+
dtype=float32,
|
135 |
+
)
|
136 |
+
nms_cfg = NMSConfig(min_confidence=0.5, min_iou=0.5)
|
137 |
+
|
138 |
+
expected_output = [
|
139 |
+
tensor(
|
140 |
+
[
|
141 |
+
[1.0000, 50.0000, 50.0000, 100.0000, 100.0000, 0.6682],
|
142 |
+
[0.0000, 60.0000, 60.0000, 110.0000, 110.0000, 0.6457],
|
143 |
+
]
|
144 |
+
)
|
145 |
+
]
|
146 |
+
|
147 |
+
output = bbox_nms(cls_dist, bbox, nms_cfg)
|
148 |
+
|
149 |
+
for out, exp in zip(output, expected_output):
|
150 |
+
assert allclose(out, exp, atol=1e-4), f"Output: {out} Expected: {exp}"
|
151 |
+
|
152 |
+
|
153 |
+
def test_calculate_map():
|
154 |
+
predictions = tensor([[0, 60, 60, 160, 160, 0.5], [0, 40, 40, 120, 120, 0.5]]) # [class, x1, y1, x2, y2]
|
155 |
+
ground_truths = tensor([[0, 50, 50, 150, 150], [0, 30, 30, 100, 100]]) # [class, x1, y1, x2, y2]
|
156 |
+
|
157 |
+
mean_ap, first_ap = calculate_map(predictions, ground_truths)
|
158 |
+
|
159 |
+
expected_mean_ap = tensor(0.2)
|
160 |
+
expected_first_ap = tensor(0.5)
|
161 |
+
|
162 |
+
assert isclose(mean_ap, expected_mean_ap, atol=1e-5), f"Mean AP mismatch: {mean_ap} != {expected_mean_ap}"
|
163 |
+
assert isclose(first_ap, expected_first_ap, atol=1e-5), f"First AP mismatch: {first_ap} != {expected_first_ap}"
|
tests/{test_tools → test_utils}/test_module_utils.py
RENAMED
@@ -1,3 +1,4 @@
|
|
|
|
1 |
import sys
|
2 |
from pathlib import Path
|
3 |
|
@@ -6,7 +7,11 @@ from torch import nn
|
|
6 |
|
7 |
project_root = Path(__file__).resolve().parent.parent.parent
|
8 |
sys.path.append(str(project_root))
|
9 |
-
from yolo.utils.module_utils import
|
|
|
|
|
|
|
|
|
10 |
|
11 |
|
12 |
@pytest.mark.parametrize(
|
@@ -35,3 +40,34 @@ def test_get_activation(activation_name, expected_type):
|
|
35 |
def test_get_activation_invalid():
|
36 |
with pytest.raises(ValueError):
|
37 |
create_activation_function("unsupported_activation")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import re
|
2 |
import sys
|
3 |
from pathlib import Path
|
4 |
|
|
|
7 |
|
8 |
project_root = Path(__file__).resolve().parent.parent.parent
|
9 |
sys.path.append(str(project_root))
|
10 |
+
from yolo.utils.module_utils import (
|
11 |
+
auto_pad,
|
12 |
+
create_activation_function,
|
13 |
+
divide_into_chunks,
|
14 |
+
)
|
15 |
|
16 |
|
17 |
@pytest.mark.parametrize(
|
|
|
40 |
def test_get_activation_invalid():
|
41 |
with pytest.raises(ValueError):
|
42 |
create_activation_function("unsupported_activation")
|
43 |
+
|
44 |
+
|
45 |
+
def test_divide_into_chunks():
|
46 |
+
input_list = [0, 1, 2, 3, 4, 5]
|
47 |
+
chunk_num = 2
|
48 |
+
expected_output = [[0, 1, 2], [3, 4, 5]]
|
49 |
+
assert divide_into_chunks(input_list, chunk_num) == expected_output
|
50 |
+
|
51 |
+
|
52 |
+
def test_divide_into_chunks_non_divisible_length():
|
53 |
+
input_list = [0, 1, 2, 3, 4, 5]
|
54 |
+
chunk_num = 4
|
55 |
+
with pytest.raises(
|
56 |
+
ValueError,
|
57 |
+
match=re.escape("The length of the input list (6) must be exactly divisible by the number of chunks (4)."),
|
58 |
+
):
|
59 |
+
divide_into_chunks(input_list, chunk_num)
|
60 |
+
|
61 |
+
|
62 |
+
def test_divide_into_chunks_single_chunk():
|
63 |
+
input_list = [0, 1, 2, 3, 4, 5]
|
64 |
+
chunk_num = 1
|
65 |
+
expected_output = [[0, 1, 2, 3, 4, 5]]
|
66 |
+
assert divide_into_chunks(input_list, chunk_num) == expected_output
|
67 |
+
|
68 |
+
|
69 |
+
def test_divide_into_chunks_equal_chunks():
|
70 |
+
input_list = [0, 1, 2, 3, 4, 5, 6, 7]
|
71 |
+
chunk_num = 4
|
72 |
+
expected_output = [[0, 1], [2, 3], [4, 5], [6, 7]]
|
73 |
+
assert divide_into_chunks(input_list, chunk_num) == expected_output
|
yolo/config/config.py
CHANGED
@@ -163,9 +163,6 @@ class YOLOLayer(nn.Module):
|
|
163 |
layer_type: str
|
164 |
usable: bool
|
165 |
|
166 |
-
def __post_init__(self):
|
167 |
-
super().__init__()
|
168 |
-
|
169 |
|
170 |
IDX_TO_ID = [
|
171 |
1,
|
|
|
163 |
layer_type: str
|
164 |
usable: bool
|
165 |
|
|
|
|
|
|
|
166 |
|
167 |
IDX_TO_ID = [
|
168 |
1,
|
yolo/tools/solver.py
CHANGED
@@ -122,6 +122,7 @@ class ModelTrainer:
|
|
122 |
self.progress.finish_one_epoch(epoch_loss, epoch)
|
123 |
|
124 |
self.validator.solve(self.validation_dataloader, epoch_idx=epoch)
|
|
|
125 |
self.progress.finish_train()
|
126 |
|
127 |
|
|
|
122 |
self.progress.finish_one_epoch(epoch_loss, epoch)
|
123 |
|
124 |
self.validator.solve(self.validation_dataloader, epoch_idx=epoch)
|
125 |
+
# TODO: save model if result are better than before
|
126 |
self.progress.finish_train()
|
127 |
|
128 |
|
yolo/utils/bounding_box_utils.py
CHANGED
@@ -46,7 +46,7 @@ def calculate_iou(bbox1, bbox2, metrics="iou") -> Tensor:
|
|
46 |
# Calculate IoU
|
47 |
iou = intersection_area / (union_area + EPS)
|
48 |
if metrics == "iou":
|
49 |
-
return iou
|
50 |
|
51 |
# Calculate centroid distance
|
52 |
cx1 = (bbox1[..., 2] + bbox1[..., 0]) / 2
|
@@ -62,7 +62,7 @@ def calculate_iou(bbox1, bbox2, metrics="iou") -> Tensor:
|
|
62 |
|
63 |
diou = iou - (cent_dis / diag_dis)
|
64 |
if metrics == "diou":
|
65 |
-
return diou
|
66 |
|
67 |
# Compute aspect ratio penalty term
|
68 |
arctan = torch.atan((bbox1[..., 2] - bbox1[..., 0]) / (bbox1[..., 3] - bbox1[..., 1] + EPS)) - torch.atan(
|
@@ -268,7 +268,7 @@ class Vec2Box:
|
|
268 |
def __init__(self, model: YOLO, image_size, device):
|
269 |
self.device = device
|
270 |
|
271 |
-
if hasattr(model, "strides"):
|
272 |
logger.info(f"🈶 Found stride of model {model.strides}")
|
273 |
self.strides = model.strides
|
274 |
else:
|
|
|
46 |
# Calculate IoU
|
47 |
iou = intersection_area / (union_area + EPS)
|
48 |
if metrics == "iou":
|
49 |
+
return iou.to(dtype)
|
50 |
|
51 |
# Calculate centroid distance
|
52 |
cx1 = (bbox1[..., 2] + bbox1[..., 0]) / 2
|
|
|
62 |
|
63 |
diou = iou - (cent_dis / diag_dis)
|
64 |
if metrics == "diou":
|
65 |
+
return diou.to(dtype)
|
66 |
|
67 |
# Compute aspect ratio penalty term
|
68 |
arctan = torch.atan((bbox1[..., 2] - bbox1[..., 0]) / (bbox1[..., 3] - bbox1[..., 1] + EPS)) - torch.atan(
|
|
|
268 |
def __init__(self, model: YOLO, image_size, device):
|
269 |
self.device = device
|
270 |
|
271 |
+
if hasattr(model, "strides") and getattr(model, "strides"):
|
272 |
logger.info(f"🈶 Found stride of model {model.strides}")
|
273 |
self.strides = model.strides
|
274 |
else:
|
yolo/utils/module_utils.py
CHANGED
@@ -68,8 +68,7 @@ def divide_into_chunks(input_list, chunk_num):
|
|
68 |
|
69 |
if list_size % chunk_num != 0:
|
70 |
raise ValueError(
|
71 |
-
f"The length of the input list ({list_size}) must be exactly
|
72 |
-
divisible by the number of chunks ({chunk_num})."
|
73 |
)
|
74 |
|
75 |
chunk_size = list_size // chunk_num
|
|
|
68 |
|
69 |
if list_size % chunk_num != 0:
|
70 |
raise ValueError(
|
71 |
+
f"The length of the input list ({list_size}) must be exactly divisible by the number of chunks ({chunk_num})."
|
|
|
72 |
)
|
73 |
|
74 |
chunk_size = list_size // chunk_num
|