re-define cropping transforms

mehta-lab · Jan 27, 2025 · ca8095d · ca8095d
1 parent e05d802
commit ca8095d
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 3 deletions.
diff --git a/viscy/transforms/__init__.py b/viscy/transforms/__init__.py
@@ -1,9 +1,11 @@
 from viscy.transforms._redef import (
+    CenterSpatialCropd,
     RandAdjustContrastd,
     RandAffined,
     RandGaussianNoised,
     RandGaussianSmoothd,
     RandScaleIntensityd,
+    RandSpatialCropd,
     RandWeightedCropd,
     ScaleIntensityRangePercentilesd,
 )
@@ -17,13 +19,15 @@
 
 __all__ = [
     "BatchedZoom",
+    "CenterSpatialCropd",
     "NormalizeSampled",
     "RandAdjustContrastd",
     "RandAffined",
     "RandGaussianNoised",
     "RandGaussianSmoothd",
     "RandInvertIntensityd",
     "RandScaleIntensityd",
+    "RandSpatialCropd",
     "RandWeightedCropd",
     "ScaleIntensityRangePercentilesd",
     "StackChannelsd",

diff --git a/viscy/transforms/_redef.py b/viscy/transforms/_redef.py
@@ -3,11 +3,13 @@
 from typing import Sequence
 
 from monai.transforms import (
+    CenterSpatialCropd,
     RandAdjustContrastd,
     RandAffined,
     RandGaussianNoised,
     RandGaussianSmoothd,
     RandScaleIntensityd,
+    RandSpatialCropd,
     RandWeightedCropd,
     ScaleIntensityRangePercentilesd,
 )
@@ -37,9 +39,9 @@ def __init__(
         self,
         keys: Sequence[str] | str,
         prob: float,
-        rotate_range: Sequence[float] | float,
-        shear_range: Sequence[float] | float,
-        scale_range: Sequence[float] | float,
+        rotate_range: Sequence[float | Sequence[float]] | float,
+        shear_range: Sequence[float | Sequence[float]] | float,
+        scale_range: Sequence[float | Sequence[float]] | float,
         **kwargs,
     ):
         super().__init__(
@@ -132,3 +134,29 @@ def __init__(
             dtype=dtype,
             allow_missing_keys=allow_missing_keys,
         )
+
+
+class RandSpatialCropd(RandSpatialCropd):
+    def __init__(
+        self,
+        keys: Sequence[str] | str,
+        roi_size: Sequence[int] | int,
+        random_center: bool = True,
+        **kwargs,
+    ):
+        super().__init__(
+            keys=keys,
+            roi_size=roi_size,
+            random_center=random_center,
+            **kwargs,
+        )
+
+
+class CenterSpatialCropd(CenterSpatialCropd):
+    def __init__(
+        self,
+        keys: Sequence[str] | str,
+        roi_size: Sequence[int] | int,
+        **kwargs,
+    ):
+        super().__init__(keys=keys, roi_size=roi_size, **kwargs)