change ther refactoring of transform and collate function

Author: yutong-xiang-97

src/lightly_train/_task_models/ltdetr_object_detection/transforms.py

@@ -15,14 +15,6 @@
 from lightly_train._task_models.object_detection_components.ltdetr_geometry import (
     ltdetr_image_size_divisor,
 )
-from lightly_train._transforms.ltdetr_transforms.base import (
-    LTDETRMosaicArgs,
-    LTDETRRandomFlipArgs,
-    LTDETRRandomIoUCropArgs,
-    LTDETRRandomPhotometricDistortArgs,
-    LTDETRRandomZoomOutArgs,
-    LTDETRResizeArgs,
-)
 from lightly_train._transforms.ltdetr_transforms.object_detection import (
     LTDETRObjectDetectionTransform,
     LTDETRObjectDetectionTransformArgs,
@@ -51,6 +43,63 @@
 from lightly_train.types import ImageSizeTuple
 
 
+class LTDETRObjectDetectionRandomPhotometricDistortArgs(RandomPhotometricDistortArgs):
+    prob: float = 0.5
+
+    brightness: tuple[float, float] = (0.875, 1.125)
+    contrast: tuple[float, float] = (0.5, 1.5)
+    saturation: tuple[float, float] = (0.5, 1.5)
+    hue: tuple[float, float] = (-0.05, 0.05)
+
+    # "auto" resolves to epoch 4, or to floor(total_epochs / 3) for runs
+    # with <= 12 epochs.
+    step_start: int | Literal["auto"] = "auto"
+    # "auto" resolves to epoch total_epochs - no_aug_epoch. For shorter runs,
+    # no_aug_epoch is scaled following a certain rule. See :func:`resolve_ltdetr_step_schedule` for the full algorithm.
+    # None means photometric distort is always on.
+    step_stop: int | Literal["auto"] | None = "auto"
+
+
+class LTDETRObjectDetectionRandomZoomOutArgs(RandomZoomOutArgs):
+    prob: float = 0.5
+
+    fill: float = 0.0
+    side_range: tuple[float, float] = (1.0, 4.0)
+
+    # "auto" resolves to epoch 4, or to floor(total_epochs / 3) for runs
+    # with <= 12 epochs.
+    step_start: int | Literal["auto"] = "auto"
+    # "auto" resolves to epoch total_epochs - no_aug_epoch. For shorter runs,
+    # no_aug_epoch is scaled following a certain rule. See :func:`resolve_ltdetr_step_schedule` for the full algorithm.
+    # None means random zoom out is always on.
+    step_stop: int | Literal["auto"] | None = "auto"
+
+
+class LTDETRObjectDetectionRandomIoUCropArgs(RandomIoUCropArgs):
+    prob: float = 0.8
+
+    min_scale: float = 0.3
+    max_scale: float = 1.0
+    min_aspect_ratio: float = 0.5
+    max_aspect_ratio: float = 2.0
+    sampler_options: Sequence[float] | None = None
+    crop_trials: int = 40
+    iou_trials: int = 1000
+
+    # "auto" resolves to epoch 4, or to floor(total_epochs / 3) for runs
+    # with <= 12 epochs.
+    step_start: int | Literal["auto"] = "auto"
+    # "auto" resolves to epoch total_epochs - no_aug_epoch. For shorter runs,
+    # no_aug_epoch is scaled following a certain rule. See :func:`resolve_ltdetr_step_schedule` for the full algorithm.
+    # None means random IoU crop is always on.
+    step_stop: int | Literal["auto"] | None = "auto"
+
+
+class LTDETRObjectDetectionRandomFlipArgs(RandomFlipArgs):
+    horizontal_prob: float = 0.5
+    vertical_prob: float = 0.0
+
+
 class LTDETRObjectDetectionScaleJitterArgs(ScaleJitterArgs):
     sizes: Sequence[tuple[int, int]] | None = [
         (480, 480),
@@ -124,29 +173,57 @@ class LTDETRObjectDetectionCopyBlendArgs(CopyBlendArgs):
     step_stop: int | Literal["auto"] | None = "auto"
 
 
+class LTDETRObjectDetectionMosaicArgs(MosaicArgs):
+    prob: float = 0.5
+
+    output_size: int = 320
+    max_size: int | None = None
+    rotation_range: float = 10.0
+    translation_range: tuple[float, float] = (0.1, 0.1)
+    scaling_range: tuple[float, float] = (0.5, 1.5)
+    fill_value: int | float = 0
+    max_cached_images: int = 50
+    random_pop: bool = True
+
+    # "auto" resolves to epoch 4, or to floor(total_epochs / 3) for runs
+    # with <= 12 epochs.
+    step_start: int | Literal["auto"] = "auto"
+    # "auto" uses a compressed short-run schedule for <= 12 epochs and
+    # transitions to the midpoint rule on longer runs.
+    # None means mosaic is always on.
+    step_stop: int | Literal["auto"] | None = "auto"
+
+
+class LTDETRObjectDetectionResizeArgs(ResizeArgs):
+    height: int | Literal["auto"] = "auto"
+    width: int | Literal["auto"] = "auto"
+
+
 class LTDETRObjectDetectionTrainTransformArgs(LTDETRObjectDetectionTransformArgs):
     channel_drop: ChannelDropArgs | None = None
     num_channels: int | Literal["auto"] = "auto"
-    photometric_distort: LTDETRRandomPhotometricDistortArgs | None = Field(
-        default_factory=LTDETRRandomPhotometricDistortArgs
+    photometric_distort: LTDETRObjectDetectionRandomPhotometricDistortArgs | None = (
+        Field(default_factory=LTDETRObjectDetectionRandomPhotometricDistortArgs)
     )
-    random_zoom_out: LTDETRRandomZoomOutArgs | None = Field(
-        default_factory=LTDETRRandomZoomOutArgs
+    random_zoom_out: LTDETRObjectDetectionRandomZoomOutArgs | None = Field(
+        default_factory=LTDETRObjectDetectionRandomZoomOutArgs
     )
-    random_iou_crop: LTDETRRandomIoUCropArgs | None = Field(
-        default_factory=LTDETRRandomIoUCropArgs
+    random_iou_crop: LTDETRObjectDetectionRandomIoUCropArgs | None = Field(
+        default_factory=LTDETRObjectDetectionRandomIoUCropArgs
     )
-    random_flip: LTDETRRandomFlipArgs | None = Field(
-        default_factory=LTDETRRandomFlipArgs
+    random_flip: LTDETRObjectDetectionRandomFlipArgs | None = Field(
+        default_factory=LTDETRObjectDetectionRandomFlipArgs
     )
     random_rotate_90: RandomRotate90Args | None = None
     random_rotate: RandomRotationArgs | None = None
     image_size: ImageSizeTuple | Literal["auto"] = "auto"
-    resize: ResizeArgs | None = Field(default_factory=LTDETRResizeArgs)
+    resize: ResizeArgs | None = Field(default_factory=LTDETRObjectDetectionResizeArgs)
     scale_jitter: LTDETRObjectDetectionScaleJitterArgs | None = Field(
         default_factory=LTDETRObjectDetectionScaleJitterArgs
     )
-    mosaic: LTDETRMosaicArgs | None = Field(default_factory=LTDETRMosaicArgs)
+    mosaic: LTDETRObjectDetectionMosaicArgs | None = Field(
+        default_factory=LTDETRObjectDetectionMosaicArgs
+    )
     mixup: LTDETRObjectDetectionMixUpArgs | None = Field(
         default_factory=LTDETRObjectDetectionMixUpArgs
     )
@@ -253,9 +330,9 @@ class LTDETRObjectDetectionValTransformArgs(LTDETRObjectDetectionTransformArgs):
     random_rotate_90: RandomRotate90Args | None = None
     random_rotate: RandomRotationArgs | None = None
     image_size: ImageSizeTuple | Literal["auto"] = "auto"
-    resize: ResizeArgs | None = Field(default_factory=LTDETRResizeArgs)
+    resize: ResizeArgs | None = Field(default_factory=LTDETRObjectDetectionResizeArgs)
     scale_jitter: ScaleJitterArgs | None = None
-    mosaic: LTDETRMosaicArgs | None = None
+    mosaic: LTDETRObjectDetectionMosaicArgs | None = None
     mixup: LTDETRObjectDetectionMixUpArgs | None = None
     copyblend: LTDETRObjectDetectionCopyBlendArgs | None = None
     bbox_params: BboxParams = Field(
@@ -324,7 +401,9 @@ class LTDETRObjectDetectionValTransform(LTDETRObjectDetectionTransform):
 
 # TODO (Lionel, 06/26): Remove all the `v2` naming once the DINOv2 LT-DETR models are
 # completely migrated to the generic LTDETR pipeline.
-class DINOv2LTDETRRandomPhotometricDistortArgsV2(RandomPhotometricDistortArgs):
+class DINOv2LTDETRObjectDetectionRandomPhotometricDistortArgsV2(
+    RandomPhotometricDistortArgs
+):
     prob: float = 0.5
 
     brightness: tuple[float, float] = (0.875, 1.125)
@@ -341,7 +420,7 @@ class DINOv2LTDETRRandomPhotometricDistortArgsV2(RandomPhotometricDistortArgs):
     step_stop: int | Literal["auto"] | None = "auto"
 
 
-class DINOv2LTDETRRandomZoomOutArgsV2(RandomZoomOutArgs):
+class DINOv2LTDETRObjectDetectionRandomZoomOutArgsV2(RandomZoomOutArgs):
     prob: float = 0.5
 
     fill: float = 0.0
@@ -356,7 +435,7 @@ class DINOv2LTDETRRandomZoomOutArgsV2(RandomZoomOutArgs):
     step_stop: int | Literal["auto"] | None = "auto"
 
 
-class DINOv2LTDETRRandomIoUCropArgsV2(RandomIoUCropArgs):
+class DINOv2LTDETRObjectDetectionRandomIoUCropArgsV2(RandomIoUCropArgs):
     prob: float = 0.8
 
     min_scale: float = 0.3
@@ -376,7 +455,7 @@ class DINOv2LTDETRRandomIoUCropArgsV2(RandomIoUCropArgs):
     step_stop: int | Literal["auto"] | None = "auto"
 
 
-class DINOv2LTDETRRandomFlipArgsV2(RandomFlipArgs):
+class DINOv2LTDETRObjectDetectionRandomFlipArgsV2(RandomFlipArgs):
     horizontal_prob: float = 0.5
     vertical_prob: float = 0.0
 
@@ -457,7 +536,7 @@ class DINOv2LTDETRObjectDetectionCopyBlendArgsV2(CopyBlendArgs):
     step_stop: int | Literal["auto"] | None = "auto"
 
 
-class DINOv2LTDETRMosaicArgsV2(MosaicArgs):
+class DINOv2LTDETRObjectDetectionMosaicArgsV2(MosaicArgs):
     prob: float = 0.5
 
     output_size: int = 320
@@ -478,7 +557,7 @@ class DINOv2LTDETRMosaicArgsV2(MosaicArgs):
     step_stop: int | Literal["auto"] | None = "auto"
 
 
-class DINOv2LTDETRResizeArgsV2(ResizeArgs):
+class DINOv2LTDETRObjectDetectionResizeArgsV2(ResizeArgs):
     height: int | Literal["auto"] = "auto"
     width: int | Literal["auto"] = "auto"
 
@@ -488,27 +567,29 @@ class DINOv2LTDETRObjectDetectionTrainTransformArgsV2(
 ):
     channel_drop: None = None
     num_channels: int | Literal["auto"] = "auto"
-    photometric_distort: DINOv2LTDETRRandomPhotometricDistortArgsV2 | None = Field(
-        default_factory=DINOv2LTDETRRandomPhotometricDistortArgsV2
-    )
-    random_zoom_out: DINOv2LTDETRRandomZoomOutArgsV2 | None = Field(
-        default_factory=DINOv2LTDETRRandomZoomOutArgsV2
+    photometric_distort: (
+        DINOv2LTDETRObjectDetectionRandomPhotometricDistortArgsV2 | None
+    ) = Field(default_factory=DINOv2LTDETRObjectDetectionRandomPhotometricDistortArgsV2)
+    random_zoom_out: DINOv2LTDETRObjectDetectionRandomZoomOutArgsV2 | None = Field(
+        default_factory=DINOv2LTDETRObjectDetectionRandomZoomOutArgsV2
     )
-    random_iou_crop: DINOv2LTDETRRandomIoUCropArgsV2 | None = Field(
-        default_factory=DINOv2LTDETRRandomIoUCropArgsV2
+    random_iou_crop: DINOv2LTDETRObjectDetectionRandomIoUCropArgsV2 | None = Field(
+        default_factory=DINOv2LTDETRObjectDetectionRandomIoUCropArgsV2
     )
-    random_flip: DINOv2LTDETRRandomFlipArgsV2 | None = Field(
-        default_factory=DINOv2LTDETRRandomFlipArgsV2
+    random_flip: DINOv2LTDETRObjectDetectionRandomFlipArgsV2 | None = Field(
+        default_factory=DINOv2LTDETRObjectDetectionRandomFlipArgsV2
     )
     random_rotate_90: RandomRotate90Args | None = None
     random_rotate: RandomRotationArgs | None = None
     image_size: ImageSizeTuple | Literal["auto"] = "auto"
-    resize: ResizeArgs | None = Field(default_factory=DINOv2LTDETRResizeArgsV2)
+    resize: ResizeArgs | None = Field(
+        default_factory=DINOv2LTDETRObjectDetectionResizeArgsV2
+    )
     scale_jitter: DINOv2LTDETRObjectDetectionScaleJitterArgsV2 | None = Field(
         default_factory=DINOv2LTDETRObjectDetectionScaleJitterArgsV2
     )
-    mosaic: DINOv2LTDETRMosaicArgsV2 | None = Field(
-        default_factory=DINOv2LTDETRMosaicArgsV2
+    mosaic: DINOv2LTDETRObjectDetectionMosaicArgsV2 | None = Field(
+        default_factory=DINOv2LTDETRObjectDetectionMosaicArgsV2
     )
     mixup: DINOv2LTDETRObjectDetectionMixUpArgsV2 | None = Field(
         default_factory=DINOv2LTDETRObjectDetectionMixUpArgsV2
@@ -595,9 +676,11 @@ class DINOv2LTDETRObjectDetectionValTransformArgsV2(LTDETRObjectDetectionTransfo
     random_rotate_90: RandomRotate90Args | None = None
     random_rotate: RandomRotationArgs | None = None
     image_size: ImageSizeTuple | Literal["auto"] = "auto"
-    resize: ResizeArgs | None = Field(default_factory=DINOv2LTDETRResizeArgsV2)
+    resize: ResizeArgs | None = Field(
+        default_factory=DINOv2LTDETRObjectDetectionResizeArgsV2
+    )
     scale_jitter: ScaleJitterArgs | None = None
-    mosaic: DINOv2LTDETRMosaicArgsV2 | None = None
+    mosaic: DINOv2LTDETRObjectDetectionMosaicArgsV2 | None = None
     mixup: DINOv2LTDETRObjectDetectionMixUpArgsV2 | None = None
     copyblend: DINOv2LTDETRObjectDetectionCopyBlendArgsV2 | None = None
     bbox_params: BboxParams = Field(