henry000 commited on
Commit
88e45b9
·
1 Parent(s): 7a45a10

✅ [Add] tests, increase test coverage

Browse files
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 test_build_model():
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 auto_pad, create_activation_function
 
 
 
 
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