🐛 [Fix] bugs in dynamic shape in training

yolo/tools/data_augmentation.py

@@ -9,10 +9,11 @@ from torchvision.transforms import functional as TF
 class AugmentationComposer:
     """Composes several transforms together."""
 
-    def __init__(self, transforms, image_size: int = [640, 640]):
+    def __init__(self, transforms, image_size: int = [640, 640], base_size: int = 640):
         self.transforms = transforms
         # TODO: handle List of image_size [640, 640]
         self.pad_resize = PadAndResize(image_size)
+        self.base_size = base_size
 
         for transform in self.transforms:
             if hasattr(transform, "set_parent"):
@@ -122,7 +123,7 @@ class Mosaic:
 
         assert self.parent is not None, "Parent is not set. Mosaic cannot retrieve image size."
 
-        img_sz = self.parent.image_size[0]  # Assuming `image_size` is defined in parent
+        img_sz = self.parent.base_size  # Assuming `image_size` is defined in parent
         more_data = self.parent.get_more_data(3)  # get 3 more images randomly
 
         data = [(image, boxes)] + more_data

yolo/tools/data_loader.py

@@ -30,11 +30,11 @@ class YoloDataset(Dataset):
         self.image_size = data_cfg.image_size
         phase_name = dataset_cfg.get(phase, phase)
         self.batch_size = data_cfg.batch_size
-        self.dynamic_shape = getattr(data_cfg, "dynamic_shape", True)
+        self.dynamic_shape = getattr(data_cfg, "dynamic_shape", False)
         self.base_size = mean(self.image_size)
 
         transforms = [eval(aug)(prob) for aug, prob in augment_cfg.items()]
-        self.transform = AugmentationComposer(transforms, self.image_size)
+        self.transform = AugmentationComposer(transforms, self.image_size, self.base_size)
         self.transform.get_more_data = self.get_more_data
         self.img_paths, self.bboxes, self.ratios = tensorlize(self.load_data(Path(dataset_cfg.path), phase_name))
 
@@ -53,20 +53,21 @@ class YoloDataset(Dataset):
 
         if not cache_path.exists():
             logger.info(f":factory: Generating {phase_name} cache")
-            data = self.filter_data(dataset_path, phase_name)
+            data = self.filter_data(dataset_path, phase_name, self.dynamic_shape)
             torch.save(data, cache_path)
         else:
             data = torch.load(cache_path, weights_only=False)
             logger.info(f":package: Loaded {phase_name} cache")
         return data
 
-    def filter_data(self, dataset_path: Path, phase_name: str) -> list:
+    def filter_data(self, dataset_path: Path, phase_name: str, sort_image: bool = False) -> list:
         """
         Filters and collects dataset information by pairing images with their corresponding labels.
 
         Parameters:
             images_path (Path): Path to the directory containing image files.
             labels_path (str): Path to the directory containing label files.
+            sort_image (bool): If True, sorts the dataset by the width-to-height ratio of images in descending order.
 
         Returns:
             list: A list of tuples, each containing the path to an image file and its associated segmentation as a tensor.
@@ -105,8 +106,11 @@ class YoloDataset(Dataset):
             labels = self.load_valid_labels(image_id, image_seg_annotations)
 
             img_path = images_path / image_name
-            with Image.open(img_path) as img:
-                width, height = img.size
+            if sort_image:
+                with Image.open(img_path) as img:
+                    width, height = img.size
+            else:
+                width, height = 0, 1
             data.append((img_path, labels, width / height))
             valid_inputs += 1
 

yolo/tools/solver.py

@@ -85,6 +85,7 @@ class TrainModel(ValidateModel):
 
     def on_train_epoch_start(self):
         self.trainer.optimizers[0].next_epoch(ceil(len(self.train_loader) / self.trainer.world_size))
+        self.vec2box.update(self.cfg.image_size)
 
     def training_step(self, batch, batch_idx):
         lr_dict = self.trainer.optimizers[0].next_batch()