🔀 [Merge] branch 'main' into TEST

tests/test_utils/test_bounding_box_utils.py

@@ -138,9 +138,9 @@ def test_anc2box_autoanchor(inference_v7_cfg: Config):
     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]),
+        torch.Size([1, 1, 80, 40, 2]),
+        torch.Size([1, 1, 40, 20, 2]),
+        torch.Size([1, 1, 20, 10, 2]),
     ]
     assert anc2box.anchor_scale.shape == torch.Size([3, 1, 3, 1, 1, 2])
 

yolo/tools/solver.py

@@ -45,7 +45,8 @@ class ValidateModel(BaseModel):
 
     def validation_step(self, batch, batch_idx):
         batch_size, images, targets, rev_tensor, img_paths = batch
-        predicts = self.post_process(self(images), image_size=images.shape[2:])
+        H, W = images.shape[2:]
+        predicts = self.post_process(self(images), image_size=[W, H])
         batch_metrics = self.metric(
             [to_metrics_format(predict) for predict in predicts], [to_metrics_format(target) for target in targets]
         )

yolo/utils/bounding_box_utils.py

@@ -122,7 +122,7 @@ def generate_anchors(image_size: List[int], strides: List[int]):
         all_anchors [HW x 2]:
         all_scalers [HW]: The index of the best targets for each anchors
     """
-    H, W = image_size
+    W, H = image_size
     anchors = []
     scaler = []
     for stride in strides:
@@ -312,17 +312,18 @@ class Vec2Box:
         self.anchor_grid, self.scaler = anchor_grid.to(device), scaler.to(device)
 
     def create_auto_anchor(self, model: YOLO, image_size):
-        dummy_input = torch.zeros(1, 3, *image_size).to(self.device)
+        W, H = image_size
+        dummy_input = torch.zeros(1, 3, H, W).to(self.device)
         dummy_output = model(dummy_input)
         strides = []
         for predict_head in dummy_output["Main"]:
             _, _, *anchor_num = predict_head[2].shape
-            strides.append(image_size[1] // anchor_num[1])
+            strides.append(W // anchor_num[1])
         return strides
 
     def update(self, image_size):
         """
-        image_size: H, W
+        image_size: W, H
         """
         if self.image_size == image_size:
             return
@@ -365,12 +366,13 @@ class Anc2Box:
         self.class_num = model.num_classes
 
     def create_auto_anchor(self, model: YOLO, image_size):
-        dummy_input = torch.zeros(1, 3, *image_size).to(self.device)
+        W, H = image_size
+        dummy_input = torch.zeros(1, 3, H, W).to(self.device)
         dummy_output = model(dummy_input)
         strides = []
         for predict_head in dummy_output["Main"]:
             _, _, *anchor_num = predict_head.shape
-            strides.append(image_size[1] // anchor_num[1])
+            strides.append(W // anchor_num[1])
         return strides
 
     def generate_anchors(self, image_size: List[int]):
@@ -383,7 +385,7 @@ class Anc2Box:
         return anchor_grids
 
     def update(self, image_size):
-        self.anchor_grid = self.generate_anchors(image_size)
+        self.anchor_grids = self.generate_anchors(image_size)
 
     def __call__(self, predicts: List[Tensor]):
         preds_box, preds_cls, preds_cnf = [], [], []