Refactor Object Detection Transforms (#835)

* refactor object detection transform

* add a protocol for ltdetr transform args

Author: yutong-xiang-97

src/lightly_train/_commands/benchmark_task.py

@@ -59,10 +59,10 @@
     _yolo_to_xyxy,
 )
 from lightly_train._task_models.task_model import TaskModel
-from lightly_train._transforms.object_detection_transform import (
-    ObjectDetectionCollateFunction,
-    ObjectDetectionTransform,
-    ObjectDetectionTransformArgs,
+from lightly_train._transforms.ltdetr_transforms.object_detection import (
+    LTDETRObjectDetectionCollateFunction,
+    LTDETRObjectDetectionTransform,
+    LTDETRObjectDetectionTransformArgs,
 )
 from lightly_train._transforms.transform import NormalizeArgs, ResizeArgs
 from lightly_train.types import (
@@ -341,7 +341,7 @@ def _make_val_bbox_params() -> BboxParams:
     )
 
 
-class _BenchmarkValTransformArgs(ObjectDetectionTransformArgs):
+class _BenchmarkValTransformArgs(LTDETRObjectDetectionTransformArgs):
     """Val transform args that mirror the training validation pipeline.
 
     All augmentations are disabled. Only resize and normalize (plus bbox params)
@@ -419,7 +419,7 @@ def _create_val_dataloader(
     val_dataset_args = data_args.get_val_args()
     dataset_cls = val_dataset_args.get_dataset_cls()
     image_info = list(val_dataset_args.list_image_info())
-    transform = ObjectDetectionTransform(transform_args=transform_args)
+    transform = LTDETRObjectDetectionTransform(transform_args=transform_args)
     dataset = dataset_cls(
         dataset_args=val_dataset_args,
         image_info=image_info,
@@ -435,7 +435,7 @@ def _create_val_dataloader(
         shuffle=False,
         num_workers=num_workers,
         drop_last=True,
-        collate_fn=ObjectDetectionCollateFunction(
+        collate_fn=LTDETRObjectDetectionCollateFunction(
             split="val",
             transform_args=transform_args,
         ),

src/lightly_train/_data/object_detection_dataset.py

@@ -23,8 +23,8 @@
 
 from lightly_train._data import file_helpers
 from lightly_train._data.task_dataset import TaskDataset
-from lightly_train._transforms.object_detection_transform import (
-    ObjectDetectionCollateFunction,
+from lightly_train._transforms.ltdetr_transforms.object_detection import (
+    LTDETRObjectDetectionCollateFunction,
 )
 from lightly_train._transforms.task_transform import TaskCollateFunction
 from lightly_train.types import ObjectDetectionDatasetItem
@@ -35,7 +35,7 @@ class ObjectDetectionDataset(TaskDataset):
     dataset_args: COCOObjectDetectionDatasetArgs | YOLOObjectDetectionDatasetArgs  # type: ignore[assignment]
 
     batch_collate_fn_cls: ClassVar[type[TaskCollateFunction]] = (
-        ObjectDetectionCollateFunction
+        LTDETRObjectDetectionCollateFunction
     )
 
     def __getitem__(self, index: int) -> ObjectDetectionDatasetItem:

src/lightly_train/_task_models/dinov2_ltdetr_object_detection/transforms.py

@@ -13,9 +13,11 @@
 from lightning_utilities.core.imports import RequirementCache
 from pydantic import Field
 
-from lightly_train._transforms.object_detection_transform import (
-    ObjectDetectionTransform,
-    ObjectDetectionTransformArgs,
+from lightly_train._transforms.ltdetr_transforms.object_detection import (
+    LTDETRObjectDetectionTransform,
+    LTDETRObjectDetectionTransformArgs,
+)
+from lightly_train._transforms.ltdetr_transforms.utils import (
     resolve_ltdetr_step_schedule_for_augmentation,
 )
 from lightly_train._transforms.transform import (
@@ -199,7 +201,7 @@ class DINOv2LTDETRObjectDetectionResizeArgs(ResizeArgs):
     width: int | Literal["auto"] = "auto"
 
 
-class DINOv2LTDETRObjectDetectionTrainTransformArgs(ObjectDetectionTransformArgs):
+class DINOv2LTDETRObjectDetectionTrainTransformArgs(LTDETRObjectDetectionTransformArgs):
     channel_drop: None = None
     num_channels: int | Literal["auto"] = "auto"
     photometric_distort: (
@@ -301,7 +303,7 @@ def resolve_step_schedule(
         )
 
 
-class DINOv2LTDETRObjectDetectionValTransformArgs(ObjectDetectionTransformArgs):
+class DINOv2LTDETRObjectDetectionValTransformArgs(LTDETRObjectDetectionTransformArgs):
     channel_drop: None = None
     num_channels: int | Literal["auto"] = "auto"
     photometric_distort: None = None
@@ -368,9 +370,9 @@ def resolve_auto(self, model_init_args: dict[str, Any]) -> None:
                     self.num_channels = len(self.normalize.mean)
 
 
-class DINOv2LTDETRObjectDetectionTrainTransform(ObjectDetectionTransform):
+class DINOv2LTDETRObjectDetectionTrainTransform(LTDETRObjectDetectionTransform):
     transform_args_cls = DINOv2LTDETRObjectDetectionTrainTransformArgs
 
 
-class DINOv2LTDETRObjectDetectionValTransform(ObjectDetectionTransform):
+class DINOv2LTDETRObjectDetectionValTransform(LTDETRObjectDetectionTransform):
     transform_args_cls = DINOv2LTDETRObjectDetectionValTransformArgs

src/lightly_train/_task_models/ltdetr_object_detection/transforms.py

@@ -7,19 +7,22 @@
 #
 from __future__ import annotations
 
-import math
 from typing import Any, Literal, Sequence
 
 from albumentations import BboxParams
-from lightning_utilities.core.imports import RequirementCache
 from pydantic import Field
 
 from lightly_train._task_models.object_detection_components.ltdetr_geometry import (
     ltdetr_image_size_divisor,
 )
-from lightly_train._transforms.object_detection_transform import (
-    ObjectDetectionTransform,
-    ObjectDetectionTransformArgs,
+from lightly_train._transforms.ltdetr_transforms.object_detection import (
+    LTDETRObjectDetectionTransform,
+    LTDETRObjectDetectionTransformArgs,
+)
+from lightly_train._transforms.ltdetr_transforms.utils import (
+    ALBUMENTATIONS_VERSION_GREATER_EQUAL_1_4_5,
+    ALBUMENTATIONS_VERSION_GREATER_EQUAL_2_0_1,
+    resolve_image_size_for_patch_size,
     resolve_ltdetr_step_schedule_for_augmentation,
 )
 from lightly_train._transforms.transform import (
@@ -39,39 +42,6 @@
 )
 from lightly_train.types import ImageSizeTuple
 
-ALBUMENTATIONS_VERSION_GREATER_EQUAL_1_4_5 = RequirementCache("albumentations>=1.4.5")
-ALBUMENTATIONS_VERSION_GREATER_EQUAL_2_0_1 = RequirementCache("albumentations>=2.0.1")
-
-
-def _resolve_image_size_for_patch_size(
-    model_init_args: dict[str, Any],
-    *,
-    default_image_size: tuple[int, int],
-    patch_size: int | None,
-) -> tuple[int, int]:
-    provided_image_size = model_init_args.get("image_size")
-    if provided_image_size is not None:
-        image_size = (
-            int(provided_image_size[0]),
-            int(provided_image_size[1]),
-        )
-        if patch_size is not None:
-            divisor = ltdetr_image_size_divisor(patch_size)
-            if any(size % divisor != 0 for size in image_size):
-                raise ValueError(
-                    "When providing an image size in model_init_args, it must be divisible by 2 * the patch size."
-                )
-        return image_size
-
-    if patch_size is None:
-        return default_image_size
-
-    divisor = ltdetr_image_size_divisor(patch_size)
-    return (
-        math.ceil(default_image_size[0] / divisor) * divisor,
-        math.ceil(default_image_size[1] / divisor) * divisor,
-    )
-
 
 class LTDETRObjectDetectionRandomPhotometricDistortArgs(RandomPhotometricDistortArgs):
     prob: float = 0.5
@@ -229,7 +199,7 @@ class LTDETRObjectDetectionResizeArgs(ResizeArgs):
     width: int | Literal["auto"] = "auto"
 
 
-class LTDETRObjectDetectionTrainTransformArgs(ObjectDetectionTransformArgs):
+class LTDETRObjectDetectionTrainTransformArgs(LTDETRObjectDetectionTransformArgs):
     channel_drop: ChannelDropArgs | None = None
     num_channels: int | Literal["auto"] = "auto"
     photometric_distort: LTDETRObjectDetectionRandomPhotometricDistortArgs | None = (
@@ -284,7 +254,7 @@ def resolve_auto(self, model_init_args: dict[str, Any]) -> None:
         patch_size: int | None = model_init_args.get("patch_size")
 
         if self.image_size == "auto":
-            self.image_size = _resolve_image_size_for_patch_size(
+            self.image_size = resolve_image_size_for_patch_size(
                 model_init_args,
                 default_image_size=(640, 640),
                 patch_size=patch_size,
@@ -350,7 +320,7 @@ def resolve_step_schedule(
         )
 
 
-class LTDETRObjectDetectionValTransformArgs(ObjectDetectionTransformArgs):
+class LTDETRObjectDetectionValTransformArgs(LTDETRObjectDetectionTransformArgs):
     channel_drop: None = None
     num_channels: int | Literal["auto"] = "auto"
     photometric_distort: None = None
@@ -387,7 +357,7 @@ def resolve_auto(self, model_init_args: dict[str, Any]) -> None:
         patch_size: int | None = model_init_args.get("patch_size")
 
         if self.image_size == "auto":
-            self.image_size = _resolve_image_size_for_patch_size(
+            self.image_size = resolve_image_size_for_patch_size(
                 model_init_args,
                 default_image_size=(640, 640),
                 patch_size=patch_size,
@@ -421,11 +391,11 @@ def resolve_auto(self, model_init_args: dict[str, Any]) -> None:
                     self.num_channels = len(self.normalize.mean)
 
 
-class LTDETRObjectDetectionTrainTransform(ObjectDetectionTransform):
+class LTDETRObjectDetectionTrainTransform(LTDETRObjectDetectionTransform):
     transform_args_cls = LTDETRObjectDetectionTrainTransformArgs
 
 
-class LTDETRObjectDetectionValTransform(ObjectDetectionTransform):
+class LTDETRObjectDetectionValTransform(LTDETRObjectDetectionTransform):
     transform_args_cls = LTDETRObjectDetectionValTransformArgs
 
 
@@ -592,7 +562,9 @@ class DINOv2LTDETRObjectDetectionResizeArgsV2(ResizeArgs):
     width: int | Literal["auto"] = "auto"
 
 
-class DINOv2LTDETRObjectDetectionTrainTransformArgsV2(ObjectDetectionTransformArgs):
+class DINOv2LTDETRObjectDetectionTrainTransformArgsV2(
+    LTDETRObjectDetectionTransformArgs
+):
     channel_drop: None = None
     num_channels: int | Literal["auto"] = "auto"
     photometric_distort: (
@@ -694,7 +666,7 @@ def resolve_step_schedule(
         )
 
 
-class DINOv2LTDETRObjectDetectionValTransformArgsV2(ObjectDetectionTransformArgs):
+class DINOv2LTDETRObjectDetectionValTransformArgsV2(LTDETRObjectDetectionTransformArgs):
     channel_drop: None = None
     num_channels: int | Literal["auto"] = "auto"
     photometric_distort: None = None
@@ -761,9 +733,9 @@ def resolve_auto(self, model_init_args: dict[str, Any]) -> None:
                     self.num_channels = len(self.normalize.mean)
 
 
-class DINOv2LTDETRObjectDetectionTrainTransformV2(ObjectDetectionTransform):
+class DINOv2LTDETRObjectDetectionTrainTransformV2(LTDETRObjectDetectionTransform):
     transform_args_cls = DINOv2LTDETRObjectDetectionTrainTransformArgsV2
 
 
-class DINOv2LTDETRObjectDetectionValTransformV2(ObjectDetectionTransform):
+class DINOv2LTDETRObjectDetectionValTransformV2(LTDETRObjectDetectionTransform):
     transform_args_cls = DINOv2LTDETRObjectDetectionValTransformArgsV2

src/lightly_train/_task_models/picodet_object_detection/transforms.py

@@ -13,9 +13,9 @@
 from lightning_utilities.core.imports import RequirementCache
 from pydantic import Field
 
-from lightly_train._transforms.object_detection_transform import (
-    ObjectDetectionTransform,
-    ObjectDetectionTransformArgs,
+from lightly_train._transforms.ltdetr_transforms.object_detection import (
+    LTDETRObjectDetectionTransform,
+    LTDETRObjectDetectionTransformArgs,
 )
 from lightly_train._transforms.transform import (
     NormalizeArgs,
@@ -72,7 +72,7 @@ class PicoDetScaleJitterArgs(ScaleJitterArgs):
     divisible_by: int | None = None
 
 
-class PicoDetObjectDetectionTrainTransformArgs(ObjectDetectionTransformArgs):
+class PicoDetObjectDetectionTrainTransformArgs(LTDETRObjectDetectionTransformArgs):
     """PicoDet training transforms aligned with the reference config.
 
     PicoDet defaults mirror LTDETR training augmentations for consistency.
@@ -182,13 +182,13 @@ def resolve_auto(self, model_init_args: dict[str, Any]) -> None:
             self.scale_jitter.divisible_by = 32
 
 
-class PicoDetObjectDetectionTrainTransform(ObjectDetectionTransform):
+class PicoDetObjectDetectionTrainTransform(LTDETRObjectDetectionTransform):
     """Training transforms for PicoDet."""
 
     transform_args_cls = PicoDetObjectDetectionTrainTransformArgs
 
 
-class PicoDetObjectDetectionValTransformArgs(ObjectDetectionTransformArgs):
+class PicoDetObjectDetectionValTransformArgs(LTDETRObjectDetectionTransformArgs):
     """PicoDet validation transforms."""
 
     channel_drop: None = None
@@ -258,7 +258,7 @@ def resolve_auto(self, model_init_args: dict[str, Any]) -> None:
                 self.num_channels = 3
 
 
-class PicoDetObjectDetectionValTransform(ObjectDetectionTransform):
+class PicoDetObjectDetectionValTransform(LTDETRObjectDetectionTransform):
     """Validation transforms for PicoDet."""
 
     transform_args_cls = PicoDetObjectDetectionValTransformArgs

src/lightly_train/_transforms/ltdetr_transforms/__init__.py

@@ -0,0 +1,7 @@
+#
+# Copyright (c) Lightly AG and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+#