Updated some batched transforms integration docs

Author: srivarra

tests/transforms/test_adjust_contrast.py

@@ -1,7 +1,6 @@
 import pytest
 import torch
 from monai.transforms import AdjustContrast, Compose
-
 from viscy.transforms import BatchedRandAdjustContrast, BatchedRandAdjustContrastd
 
 

tests/transforms/test_crop.py

@@ -1,7 +1,6 @@
 import pytest
 import torch
 from monai.transforms import Compose
-
 from viscy.transforms._crop import (
     BatchedCenterSpatialCrop,
     BatchedCenterSpatialCropd,

tests/transforms/test_flip.py

@@ -1,6 +1,5 @@
 import pytest
 import torch
-
 from viscy.transforms import BatchedRandFlip, BatchedRandFlipd
 
 

tests/transforms/test_gaussian_smooth.py

@@ -7,7 +7,6 @@
     get_gaussian_kernel3d,
 )
 from monai.transforms.intensity.array import GaussianSmooth
-
 from viscy.transforms import BatchedRandGaussianSmooth, BatchedRandGaussianSmoothd
 from viscy.transforms._gaussian_smooth import filter3d_separable
 

tests/transforms/test_noise.py

@@ -1,7 +1,6 @@
 import pytest
 import torch
 from monai.transforms import Compose
-
 from viscy.transforms import BatchedRandGaussianNoise, BatchedRandGaussianNoised
 
 

tests/transforms/test_scale_intensity.py

@@ -1,7 +1,6 @@
 import pytest
 import torch
 from monai.transforms import RandScaleIntensity
-
 from viscy.transforms import BatchedRandScaleIntensity, BatchedRandScaleIntensityd
 
 

tests/transforms/test_transforms.py

@@ -1,6 +1,5 @@
 import pytest
 import torch
-
 from viscy.transforms._decollate import Decollate
 from viscy.transforms._transforms import (
     BatchedScaleIntensityRangePercentiles,

viscy/data/triplet.py

@@ -207,8 +207,9 @@ def _get_tensorstore(self, position: Position) -> ts.TensorStore:
         return self._tensorstores[fov_name]
 
     def _filter_tracks(self, tracks_tables: list[pd.DataFrame]) -> pd.DataFrame:
-        """Exclude tracks that are too close to the border
-        or do not have the next time point.
+        """
+
+        Exclude tracks that are too close to the border or do not have the next time point.
 
         Parameters
         ----------

viscy/transforms/batched_rand_3d_elasticd.py

@@ -5,7 +5,29 @@
 
 
 class BatchedRand3DElasticd(MapTransform, RandomizableTransform):
-    """Batched 3D elastic deformation for biological structures."""
+    """Apply random 3D elastic deformation to input data.
+
+    Uses Gaussian-smoothed displacement fields to simulate deformation.
+
+    Parameters
+    ----------
+    keys : str or Iterable[str]
+        Keys of the corresponding items to be transformed.
+    sigma_range : tuple[float, float]
+        Range for random sigma values used in Gaussian smoothing.
+    magnitude_range : tuple[float, float]
+        Range for random displacement magnitude values.
+    spatial_size : tuple[int, int, int] or int or None, optional
+        Expected spatial size of input data.
+    prob : float, optional
+        Probability of applying the transform, by default 0.1.
+    mode : str, optional
+        Interpolation mode for grid sampling, by default "bilinear".
+    padding_mode : str, optional
+        Padding mode for grid sampling, by default "reflection".
+    allow_missing_keys : bool, optional
+        Whether to ignore missing keys, by default False.
+    """
 
     def __init__(
         self,
@@ -29,7 +51,6 @@ def __init__(
     def _generate_elastic_field(
         self, shape: torch.Size, device: torch.device
     ) -> Tensor:
-        """Generate batched elastic deformation field."""
         batch_size = shape[0]
         spatial_dims = shape[2:]  # Skip batch and channel
 
@@ -76,6 +97,18 @@ def _generate_elastic_field(
         return torch.stack(displacement_fields)
 
     def __call__(self, sample: dict[str, Tensor]) -> dict[str, Tensor]:
+        """Apply elastic deformation to sample data.
+
+        Parameters
+        ----------
+        sample : dict[str, Tensor]
+            Dictionary containing image tensors to transform.
+
+        Returns
+        -------
+        dict[str, Tensor]
+            Dictionary with transformed tensors.
+        """
         self.randomize(None)
         d = dict(sample)
 

viscy/transforms/batched_rand_histogram_shiftd.py

@@ -5,7 +5,21 @@
 
 
 class BatchedRandHistogramShiftd(MapTransform, RandomizableTransform):
-    """Batched random histogram shifting for intensity distribution changes."""
+    """Apply random histogram shifts to modify intensity distributions.
+
+    Adds random intensity offsets to simulate illumination variations.
+
+    Parameters
+    ----------
+    keys : str or Iterable[str]
+        Keys of the corresponding items to be transformed.
+    shift_range : tuple[float, float], optional
+        Range for random intensity shift values, by default (-0.1, 0.1).
+    prob : float, optional
+        Probability of applying the transform, by default 0.1.
+    allow_missing_keys : bool, optional
+        Whether to ignore missing keys, by default False.
+    """
 
     def __init__(
         self,
@@ -19,6 +33,18 @@ def __init__(
         self.shift_range = shift_range
 
     def __call__(self, sample: dict[str, Tensor]) -> dict[str, Tensor]:
+        """Apply histogram shift to sample data.
+
+        Parameters
+        ----------
+        sample : dict[str, Tensor]
+            Dictionary containing image tensors to transform.
+
+        Returns
+        -------
+        dict[str, Tensor]
+            Dictionary with intensity-shifted tensors.
+        """
         self.randomize(None)
         d = dict(sample)