Update semantic segmentation transforms to use OTX's instead of torchvision

src/otx/core/data/dataset/segmentation.py

@@ -213,11 +213,10 @@ def _get_item_impl(self, index: int) -> SegDataEntity | None:
                 image_color_channel=self.image_color_channel,
                 ignored_labels=ignored_labels,
             ),
-            gt_seg_map=tv_tensors.Mask(
-                mask,
-            ),
+            masks=tv_tensors.Mask(mask[None]),
         )
-        return self._apply_transforms(entity)
+        transformed_entity = self._apply_transforms(entity)
+        return transformed_entity.wrap(masks=transformed_entity.masks[0]) if transformed_entity else None
 
     @property
     def collate_fn(self) -> Callable:

src/otx/core/data/entity/segmentation.py

@@ -24,15 +24,15 @@
 class SegDataEntity(OTXDataEntity):
     """Data entity for segmentation task.
 
-    :param gt_seg_map: mask annotations
+    :param mask: mask annotations
     """
 
     @property
     def task(self) -> OTXTaskType:
         """OTX Task type definition."""
         return OTXTaskType.SEMANTIC_SEGMENTATION
 
-    gt_seg_map: tv_tensors.Mask
+    masks: tv_tensors.Mask
 
 
 @dataclass
@@ -66,7 +66,7 @@ def collate_fn(
             batch_size=batch_data.batch_size,
             images=batch_data.images,
             imgs_info=batch_data.imgs_info,
-            masks=[entity.gt_seg_map for entity in entities],
+            masks=[entity.masks for entity in entities],
         )
 
     def pin_memory(self) -> SegBatchDataEntity:

src/otx/core/data/transform_libs/mmseg.py

@@ -37,7 +37,7 @@ def transform(self, results: dict) -> dict:
             msg = "__otx__ key should be passed from the previous pipeline (LoadImageFromFile)"
             raise RuntimeError(msg)
         if isinstance(otx_data_entity, SegDataEntity):
-            gt_masks = otx_data_entity.gt_seg_map.numpy()
+            gt_masks = otx_data_entity.masks.numpy()
             results["gt_seg_map"] = gt_masks
             # we need this to properly handle seg maps during transforms
             results["seg_fields"] = ["gt_seg_map"]
@@ -69,7 +69,7 @@ def transform(self, results: dict) -> SegDataEntity:
         return SegDataEntity(
             image=image,
             img_info=image_info,
-            gt_seg_map=masks,
+            masks=masks,
         )
 
 

src/otx/recipe/_base_/data/instance_segmentation.yaml

@@ -16,6 +16,7 @@ train_subset:
   num_workers: 2
   sampler:
     class_path: torch.utils.data.RandomSampler
+
 val_subset:
   subset_name: val
   transform_lib_type: TORCHVISION
@@ -26,6 +27,7 @@ val_subset:
   num_workers: 2
   sampler:
     class_path: torch.utils.data.RandomSampler
+
 test_subset:
   subset_name: test
   transform_lib_type: TORCHVISION

src/otx/recipe/_base_/data/rotated_detection.yaml

@@ -16,6 +16,7 @@ train_subset:
   num_workers: 2
   sampler:
     class_path: torch.utils.data.RandomSampler
+
 val_subset:
   subset_name: val
   transform_lib_type: TORCHVISION
@@ -26,6 +27,7 @@ val_subset:
   num_workers: 2
   sampler:
     class_path: torch.utils.data.RandomSampler
+
 test_subset:
   subset_name: test
   transform_lib_type: TORCHVISION

src/otx/recipe/_base_/data/semantic_segmentation.yaml

@@ -13,22 +13,23 @@ train_subset:
   transform_lib_type: TORCHVISION
   to_tv_image: true
   transforms:
-    - class_path: torchvision.transforms.v2.RandomResizedCrop
+    - class_path: otx.core.data.transform_libs.torchvision.RandomResizedCrop
       init_args:
-        size:
+        scale:
           - 512
           - 512
-        scale:
+        crop_ratio_range:
           - 0.2
           - 1.0
-        ratio:
+        aspect_ratio_range:
           - 0.5
           - 2.0
-        antialias: true
+        transform_mask: true
     - class_path: otx.core.data.transform_libs.torchvision.PhotoMetricDistortion
+    - class_path: otx.core.data.transform_libs.torchvision.RandomFlip
       init_args:
+        prob: 0.5
         is_numpy_to_tvtensor: true
-    - class_path: torchvision.transforms.v2.RandomHorizontalFlip
     - class_path: torchvision.transforms.v2.ToDtype
       init_args:
         dtype: ${as_torch_dtype:torch.float32}
@@ -38,18 +39,21 @@ train_subset:
         std: [58.395, 57.12, 57.375]
   sampler:
     class_path: torch.utils.data.RandomSampler
+
 val_subset:
   subset_name: val
   batch_size: 8
   num_workers: 4
   transform_lib_type: TORCHVISION
   to_tv_image: true
   transforms:
-    - class_path: torchvision.transforms.v2.Resize
+    - class_path: otx.core.data.transform_libs.torchvision.Resize
       init_args:
-        size:
+        scale:
           - 512
           - 512
+        transform_mask: true
+        is_numpy_to_tvtensor: true
     - class_path: torchvision.transforms.v2.ToDtype
       init_args:
         dtype: ${as_torch_dtype:torch.float32}
@@ -59,18 +63,21 @@ val_subset:
         std: [58.395, 57.12, 57.375]
   sampler:
     class_path: torch.utils.data.RandomSampler
+
 test_subset:
   subset_name: test
   num_workers: 4
   batch_size: 8
   transform_lib_type: TORCHVISION
   to_tv_image: true
   transforms:
-    - class_path: torchvision.transforms.v2.Resize
+    - class_path: otx.core.data.transform_libs.torchvision.Resize
       init_args:
-        size:
+        scale:
           - 512
           - 512
+        transform_mask: true
+        is_numpy_to_tvtensor: true
     - class_path: torchvision.transforms.v2.ToDtype
       init_args:
         dtype: ${as_torch_dtype:torch.float32}

src/otx/recipe/classification/h_label_cls/deit_tiny.yaml

@@ -44,7 +44,6 @@ overrides:
         - class_path: otx.core.data.transform_libs.torchvision.RandomResizedCrop
           init_args:
             scale: 224
-            backend: cv2
         - class_path: otx.core.data.transform_libs.torchvision.RandomFlip
           init_args:
             prob: 0.5

src/otx/recipe/classification/h_label_cls/efficientnet_b0.yaml

@@ -43,7 +43,6 @@ overrides:
         - class_path: otx.core.data.transform_libs.torchvision.EfficientNetRandomCrop
           init_args:
             scale: 224
-            backend: cv2
         - class_path: otx.core.data.transform_libs.torchvision.RandomFlip
           init_args:
             prob: 0.5

src/otx/recipe/classification/h_label_cls/efficientnet_v2.yaml

@@ -43,7 +43,6 @@ overrides:
         - class_path: otx.core.data.transform_libs.torchvision.EfficientNetRandomCrop
           init_args:
             scale: 224
-            backend: cv2
         - class_path: otx.core.data.transform_libs.torchvision.RandomFlip
           init_args:
             prob: 0.5

src/otx/recipe/classification/h_label_cls/mobilenet_v3_large.yaml

@@ -49,7 +49,6 @@ overrides:
         - class_path: otx.core.data.transform_libs.torchvision.RandomResizedCrop
           init_args:
             scale: 224
-            backend: cv2
         - class_path: otx.core.data.transform_libs.torchvision.RandomFlip
           init_args:
             prob: 0.5

src/otx/recipe/classification/h_label_cls/tv_efficientnet_b3.yaml

@@ -45,7 +45,6 @@ overrides:
         - class_path: otx.core.data.transform_libs.torchvision.EfficientNetRandomCrop
           init_args:
             scale: 224
-            backend: cv2
         - class_path: otx.core.data.transform_libs.torchvision.RandomFlip
           init_args:
             prob: 0.5

src/otx/recipe/classification/h_label_cls/tv_efficientnet_v2_l.yaml

@@ -45,7 +45,6 @@ overrides:
         - class_path: otx.core.data.transform_libs.torchvision.EfficientNetRandomCrop
           init_args:
             scale: 224
-            backend: cv2
         - class_path: otx.core.data.transform_libs.torchvision.RandomFlip
           init_args:
             prob: 0.5

src/otx/recipe/classification/h_label_cls/tv_mobilenet_v3_small.yaml

@@ -45,7 +45,6 @@ overrides:
         - class_path: otx.core.data.transform_libs.torchvision.RandomResizedCrop
           init_args:
             scale: 224
-            backend: cv2
         - class_path: otx.core.data.transform_libs.torchvision.RandomFlip
           init_args:
             prob: 0.5

src/otx/recipe/classification/multi_class_cls/deit_tiny.yaml

@@ -43,7 +43,6 @@ overrides:
         - class_path: otx.core.data.transform_libs.torchvision.RandomResizedCrop
           init_args:
             scale: 224
-            backend: cv2
         - class_path: otx.core.data.transform_libs.torchvision.RandomFlip
           init_args:
             prob: 0.5

src/otx/recipe/classification/multi_class_cls/dino_v2.yaml

@@ -50,7 +50,6 @@ overrides:
         - class_path: otx.core.data.transform_libs.torchvision.RandomResizedCrop
           init_args:
             scale: 224
-            backend: cv2
             is_numpy_to_tvtensor: true
       sampler:
         class_path: otx.algo.samplers.balanced_sampler.BalancedSampler

src/otx/recipe/classification/multi_class_cls/efficientnet_b0.yaml

@@ -44,7 +44,6 @@ overrides:
         - class_path: otx.core.data.transform_libs.torchvision.EfficientNetRandomCrop
           init_args:
             scale: 224
-            backend: cv2
         - class_path: otx.core.data.transform_libs.torchvision.RandomFlip
           init_args:
             prob: 0.5

src/otx/recipe/classification/multi_class_cls/efficientnet_v2.yaml

@@ -43,7 +43,6 @@ overrides:
         - class_path: otx.core.data.transform_libs.torchvision.EfficientNetRandomCrop
           init_args:
             scale: 224
-            backend: cv2
         - class_path: otx.core.data.transform_libs.torchvision.RandomFlip
           init_args:
             prob: 0.5

src/otx/recipe/classification/multi_class_cls/mobilenet_v3_large.yaml

@@ -48,7 +48,6 @@ overrides:
         - class_path: otx.core.data.transform_libs.torchvision.RandomResizedCrop
           init_args:
             scale: 224
-            backend: cv2
         - class_path: otx.core.data.transform_libs.torchvision.RandomFlip
           init_args:
             prob: 0.5

src/otx/recipe/classification/multi_label_cls/deit_tiny.yaml

@@ -49,7 +49,6 @@ overrides:
         - class_path: otx.core.data.transform_libs.torchvision.RandomResizedCrop
           init_args:
             scale: 224
-            backend: cv2
         - class_path: otx.core.data.transform_libs.torchvision.RandomFlip
           init_args:
             prob: 0.5

src/otx/recipe/classification/multi_label_cls/efficientnet_b0.yaml

@@ -46,7 +46,6 @@ overrides:
         - class_path: otx.core.data.transform_libs.torchvision.EfficientNetRandomCrop
           init_args:
             scale: 224
-            backend: cv2
         - class_path: otx.core.data.transform_libs.torchvision.RandomFlip
           init_args:
             prob: 0.5

src/otx/recipe/classification/multi_label_cls/efficientnet_v2.yaml

@@ -49,7 +49,6 @@ overrides:
         - class_path: otx.core.data.transform_libs.torchvision.EfficientNetRandomCrop
           init_args:
             scale: 224
-            backend: cv2
         - class_path: otx.core.data.transform_libs.torchvision.RandomFlip
           init_args:
             prob: 0.5

src/otx/recipe/classification/multi_label_cls/mobilenet_v3_large.yaml

@@ -50,7 +50,6 @@ overrides:
         - class_path: otx.core.data.transform_libs.torchvision.RandomResizedCrop
           init_args:
             scale: 224
-            backend: cv2
         - class_path: otx.core.data.transform_libs.torchvision.RandomFlip
           init_args:
             prob: 0.5

src/otx/recipe/classification/multi_label_cls/tv_efficientnet_b3.yaml

@@ -45,7 +45,6 @@ overrides:
         - class_path: otx.core.data.transform_libs.torchvision.EfficientNetRandomCrop
           init_args:
             scale: 224
-            backend: cv2
         - class_path: otx.core.data.transform_libs.torchvision.RandomFlip
           init_args:
             prob: 0.5

src/otx/recipe/classification/multi_label_cls/tv_efficientnet_v2_l.yaml

@@ -49,7 +49,6 @@ overrides:
         - class_path: otx.core.data.transform_libs.torchvision.EfficientNetRandomCrop
           init_args:
             scale: 224
-            backend: cv2
         - class_path: otx.core.data.transform_libs.torchvision.RandomFlip
           init_args:
             prob: 0.5

src/otx/recipe/classification/multi_label_cls/tv_mobilenet_v3_small.yaml

@@ -45,7 +45,6 @@ overrides:
         - class_path: otx.core.data.transform_libs.torchvision.RandomResizedCrop
           init_args:
             scale: 224
-            backend: cv2
         - class_path: otx.core.data.transform_libs.torchvision.RandomFlip
           init_args:
             prob: 0.5

src/otx/recipe/semantic_segmentation/dino_v2.yaml

@@ -42,22 +42,23 @@ overrides:
   data:
     train_subset:
       transforms:
-        - class_path: torchvision.transforms.v2.RandomResizedCrop
+        - class_path: otx.core.data.transform_libs.torchvision.RandomResizedCrop
           init_args:
-            size:
+            scale:
               - 560
               - 560
-            scale:
+            crop_ratio_range:
               - 0.2
               - 1.0
-            ratio:
+            aspect_ratio_range:
               - 0.5
               - 2.0
-            antialias: true
+            transform_mask: true
         - class_path: otx.core.data.transform_libs.torchvision.PhotoMetricDistortion
+        - class_path: otx.core.data.transform_libs.torchvision.RandomFlip
           init_args:
+            prob: 0.5
             is_numpy_to_tvtensor: true
-        - class_path: torchvision.transforms.v2.RandomHorizontalFlip
         - class_path: torchvision.transforms.v2.ToDtype
           init_args:
             dtype: ${as_torch_dtype:torch.float32}
@@ -68,11 +69,13 @@ overrides:
 
     val_subset:
       transforms:
-        - class_path: torchvision.transforms.v2.Resize
+        - class_path: otx.core.data.transform_libs.torchvision.Resize
           init_args:
-            size:
+            scale:
               - 560
               - 560
+            transform_mask: true
+            is_numpy_to_tvtensor: true
         - class_path: torchvision.transforms.v2.ToDtype
           init_args:
             dtype: ${as_torch_dtype:torch.float32}
@@ -83,11 +86,13 @@ overrides:
 
     test_subset:
       transforms:
-        - class_path: torchvision.transforms.v2.Resize
+        - class_path: otx.core.data.transform_libs.torchvision.Resize
           init_args:
-            size:
+            scale:
               - 560
               - 560
+            transform_mask: true
+            is_numpy_to_tvtensor: true
         - class_path: torchvision.transforms.v2.ToDtype
           init_args:
             dtype: ${as_torch_dtype:torch.float32}

tests/unit/core/data/test_dataset.py

@@ -102,8 +102,8 @@ def test_ignore_index(self, fxt_mock_dm_subset):
         # The mask is np.eye(10) with label_id = 0,
         # so that the diagonal is filled with zero
         # and others are filled with ignore_index.
-        gt_seg_map = next(iter(dataset)).gt_seg_map
-        assert gt_seg_map.sum() == (10 * 10 - 10) * 100
+        masks = next(iter(dataset)).masks
+        assert masks.sum() == (10 * 10 - 10) * 100
 
     def test_overflown_ignore_index(self, fxt_mock_dm_subset):
         dataset = OTXSegmentationDataset(