update augmentation strategy

Author: Donglai Wei

connectomics/config/auto_config.py

@@ -11,9 +11,8 @@
 Users can manually override any auto-determined parameters.
 """
 
-import torch
 import numpy as np
-from typing import Dict, List, Tuple, Optional, Any
+from typing import Dict, List, Optional, Any
 from dataclasses import dataclass, field
 from omegaconf import OmegaConf, DictConfig
 import warnings

connectomics/config/gpu_utils.py

@@ -7,8 +7,7 @@
 
 import torch
 import psutil
-from typing import Dict, Optional, Tuple
-import warnings
+from typing import Dict, Tuple
 
 
 def get_gpu_info() -> Dict[str, any]:

connectomics/config/hydra_config.py

@@ -614,20 +614,32 @@ class MonitorConfig:
 # Augmentation configurations
 @dataclass
 class FlipConfig:
-    """Random flip augmentation."""
+    """Random flipping augmentation."""
 
     enabled: bool = True
     prob: float = 0.5
     spatial_axis: Optional[List[int]] = None  # None = all axes
 
 
+@dataclass
+class AffineConfig:
+    """Affine transformation augmentation (rotation, scaling, shearing)."""
+
+    enabled: bool = False  # Disabled by default (can be combined with Rotate90d)
+    prob: float = 0.5
+    rotate_range: Tuple[float, float, float] = (0.2, 0.2, 0.2)  # Rotation range in radians (~11°)
+    scale_range: Tuple[float, float, float] = (0.1, 0.1, 0.1)   # Scaling range (±10%)
+    shear_range: Tuple[float, float, float] = (0.1, 0.1, 0.1)   # Shearing range (±10%)
+
+
 @dataclass
 class RotateConfig:
     """Random rotation augmentation."""
 
     enabled: bool = True
     prob: float = 0.5
     max_angle: float = 90.0
+    spatial_axes: Tuple[int, int] = (1, 2)  # Axes to rotate: (1, 2) = Y-X plane (preserves Z)
 
 
 @dataclass
@@ -736,10 +748,12 @@ class CopyPasteConfig:
 class AugmentationConfig:
     """Complete augmentation configuration."""
 
+    preset: str = "some"  # "all", "some", or "none" - controls how enabled flags are interpreted
     enabled: bool = False
 
     # Standard augmentations
     flip: FlipConfig = field(default_factory=FlipConfig)
+    affine: AffineConfig = field(default_factory=AffineConfig) # Added AffineConfig
     rotate: RotateConfig = field(default_factory=RotateConfig)
     elastic: ElasticConfig = field(default_factory=ElasticConfig)
     intensity: IntensityConfig = field(default_factory=IntensityConfig)
@@ -1022,6 +1036,7 @@ def configure_instance_segmentation(cfg: Config, boundary_thickness: int = 5) ->
     # Augmentation configuration
     "AugmentationConfig",
     "FlipConfig",
+    "AffineConfig", # Added AffineConfig
     "RotateConfig",
     "ElasticConfig",
     "IntensityConfig",

connectomics/data/__init__.py

@@ -14,9 +14,6 @@
     from connectomics.data.process import MultiTaskLabelTransformd, create_label_transform_pipeline
 """
 
-from .dataset.dataset_base import *
-from .dataset import *
-from .io import *
 
 # Make submodules available
 from . import augment

connectomics/data/augment/build.py

@@ -6,7 +6,6 @@
 
 from __future__ import annotations
 from typing import Dict
-import numpy as np
 import torch
 from monai.transforms import (
     Compose,
@@ -41,11 +40,10 @@
     RandCutBlurd,
     RandMixupd,
     RandCopyPasted,
-    SqueezeLabeld,
     NormalizeLabelsd,
     SmartNormalizeIntensityd,
 )
-from ...config.hydra_config import Config, AugmentationConfig, LabelTransformConfig
+from ...config.hydra_config import Config, AugmentationConfig
 
 
 def build_train_transforms(
@@ -522,23 +520,59 @@ def _build_augmentations(aug_cfg: AugmentationConfig, keys: list[str]) -> list:
         List of MONAI transforms
     """
     transforms = []
-
+    
+    # Get preset mode (default to "some" for backward compatibility)
+    preset = getattr(aug_cfg, "preset", "some")
+    
+    # Helper function to check if augmentation should be applied
+    def should_augment(aug_name: str, aug_enabled: bool) -> bool:
+        """
+        Determine if augmentation should be applied based on preset mode.
+        
+        - "all": enabled = True by default, set to False to disable
+        - "some": enabled = False by default, set to True to enable
+        - "none": always False (no augmentations)
+        """
+        if preset == "none":
+            return False
+        elif preset == "all":
+            # All enabled by default, respect False overrides
+            return aug_enabled
+        else:  # preset == "some"
+            # None enabled by default, only use True values
+            return aug_enabled
+    
     # Standard geometric augmentations
-    if aug_cfg.flip.enabled:
+    if should_augment("flip", aug_cfg.flip.enabled):
         transforms.append(
             RandFlipd(keys=keys, prob=aug_cfg.flip.prob, spatial_axis=aug_cfg.flip.spatial_axis)
         )
 
-    if aug_cfg.rotate.enabled:
+    if should_augment("rotate", aug_cfg.rotate.enabled):
+        # Use spatial_axes from config if available
+        spatial_axes = getattr(aug_cfg.rotate, "spatial_axes", (1, 2))
         transforms.append(
             RandRotate90d(
                 keys=keys,
                 prob=aug_cfg.rotate.prob,
-                spatial_axes=(1, 2),  # Rotate in Y-X plane to preserve anisotropic Z
+                spatial_axes=spatial_axes,  # Rotate in specified plane
+            )
+        )
+
+    if should_augment("affine", aug_cfg.affine.enabled):
+        transforms.append(
+            RandAffined(
+                keys=keys,
+                prob=aug_cfg.affine.prob,
+                rotate_range=aug_cfg.affine.rotate_range,
+                scale_range=aug_cfg.affine.scale_range,
+                shear_range=aug_cfg.affine.shear_range,
+                mode="bilinear",
+                padding_mode="reflection",
             )
         )
 
-    if aug_cfg.elastic.enabled:
+    if should_augment("elastic", aug_cfg.elastic.enabled):
         transforms.append(
             Rand3DElasticd(
                 keys=keys,
@@ -549,7 +583,7 @@ def _build_augmentations(aug_cfg: AugmentationConfig, keys: list[str]) -> list:
         )
 
     # Intensity augmentations (only for images)
-    if aug_cfg.intensity.enabled:
+    if should_augment("intensity", aug_cfg.intensity.enabled):
         if aug_cfg.intensity.gaussian_noise_prob > 0:
             transforms.append(
                 RandGaussianNoised(
@@ -578,7 +612,7 @@ def _build_augmentations(aug_cfg: AugmentationConfig, keys: list[str]) -> list:
             )
 
     # EM-specific augmentations
-    if aug_cfg.misalignment.enabled:
+    if should_augment("misalignment", aug_cfg.misalignment.enabled):
         transforms.append(
             RandMisAlignmentd(
                 keys=keys,
@@ -588,7 +622,7 @@ def _build_augmentations(aug_cfg: AugmentationConfig, keys: list[str]) -> list:
             )
         )
 
-    if aug_cfg.missing_section.enabled:
+    if should_augment("missing_section", aug_cfg.missing_section.enabled):
         transforms.append(
             RandMissingSectiond(
                 keys=keys,
@@ -597,7 +631,7 @@ def _build_augmentations(aug_cfg: AugmentationConfig, keys: list[str]) -> list:
             )
         )
 
-    if aug_cfg.motion_blur.enabled:
+    if should_augment("motion_blur", aug_cfg.motion_blur.enabled):
         transforms.append(
             RandMotionBlurd(
                 keys=["image"],
@@ -607,7 +641,7 @@ def _build_augmentations(aug_cfg: AugmentationConfig, keys: list[str]) -> list:
             )
         )
 
-    if aug_cfg.cut_noise.enabled:
+    if should_augment("cut_noise", aug_cfg.cut_noise.enabled):
         transforms.append(
             RandCutNoised(
                 keys=["image"],
@@ -617,7 +651,7 @@ def _build_augmentations(aug_cfg: AugmentationConfig, keys: list[str]) -> list:
             )
         )
 
-    if aug_cfg.cut_blur.enabled:
+    if should_augment("cut_blur", aug_cfg.cut_blur.enabled):
         transforms.append(
             RandCutBlurd(
                 keys=["image"],
@@ -628,7 +662,7 @@ def _build_augmentations(aug_cfg: AugmentationConfig, keys: list[str]) -> list:
             )
         )
 
-    if aug_cfg.missing_parts.enabled:
+    if should_augment("missing_parts", aug_cfg.missing_parts.enabled):
         transforms.append(
             RandMissingPartsd(
                 keys=keys,
@@ -638,14 +672,14 @@ def _build_augmentations(aug_cfg: AugmentationConfig, keys: list[str]) -> list:
         )
 
     # Advanced augmentations
-    if aug_cfg.mixup.enabled:
+    if should_augment("mixup", aug_cfg.mixup.enabled):
         transforms.append(
             RandMixupd(
                 keys=["image"], prob=aug_cfg.mixup.prob, alpha_range=aug_cfg.mixup.alpha_range
             )
         )
 
-    if aug_cfg.copy_paste.enabled:
+    if should_augment("copy_paste", aug_cfg.copy_paste.enabled):
         transforms.append(
             RandCopyPasted(
                 keys=["image"],

connectomics/data/augment/monai_transforms.py

@@ -13,7 +13,6 @@
 import cv2
 from monai.config import KeysCollection
 from monai.transforms import MapTransform, RandomizableTransform
-from monai.utils import ensure_tuple_rep
 
 
 class RandMisAlignmentd(RandomizableTransform, MapTransform):
@@ -807,7 +806,6 @@ def _find_best_paste(
         label_flipped: torch.Tensor,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         """Find best rotation and position with minimal overlap."""
-        import torchvision.transforms.functional as tf
         from scipy.ndimage import binary_dilation
 
         labels = torch.stack([label_orig, label_flipped])

connectomics/data/dataset/dataset_base.py

@@ -7,22 +7,17 @@
 """
 
 from __future__ import annotations
-from typing import Dict, List, Any, Optional, Union, Callable, Sequence, Tuple
-from abc import ABC, abstractmethod
-import os
+from typing import Dict, Any, Optional, Sequence, Tuple
 import numpy as np
 
 import torch
 import torch.utils.data
-import pytorch_lightning as pl
-from torch.utils.data import DataLoader
 
 # MONAI imports
 from monai.data import Dataset, CacheDataset, PersistentDataset
 from monai.transforms import Compose
 from monai.utils import ensure_tuple_rep
 
-from ..io import read_volume
 
 
 class MonaiConnectomicsDataset(Dataset):

connectomics/data/dataset/dataset_filename.py

@@ -7,19 +7,16 @@
 """
 
 from __future__ import annotations
-from typing import Dict, List, Any, Optional, Union, Sequence, Tuple
+from typing import Dict, Any, Optional, Tuple
 import json
 import random
 from pathlib import Path
 import warnings
 
-import numpy as np
 import torch
 from monai.data import Dataset
-from monai.transforms import Compose, LoadImaged, EnsureChannelFirstd
-from monai.utils import ensure_tuple_rep
+from monai.transforms import Compose
 
-from .dataset_base import MonaiConnectomicsDataset
 
 
 class MonaiFilenameDataset(Dataset):

connectomics/data/dataset/dataset_tile.py

@@ -6,18 +6,15 @@
 """
 
 from __future__ import annotations
-from typing import Dict, List, Any, Optional, Union, Sequence, Tuple
-import numpy as np
+from typing import Dict, List, Any, Optional, Tuple
 import json
-import random
 
-import torch
-from monai.data import Dataset, CacheDataset
-from monai.transforms import Compose, LoadImaged, EnsureChannelFirstd
+from monai.data import CacheDataset
+from monai.transforms import Compose, EnsureChannelFirstd
 from monai.utils import ensure_tuple_rep
 
 from .dataset_base import MonaiConnectomicsDataset
-from ..io import create_tile_metadata, reconstruct_volume_from_tiles, TileLoaderd
+from ..io import TileLoaderd
 
 
 class MonaiTileDataset(MonaiConnectomicsDataset):

connectomics/data/dataset/dataset_volume.py

@@ -6,19 +6,13 @@
 """
 
 from __future__ import annotations
-from typing import Dict, List, Any, Optional, Union, Callable, Sequence, Tuple
-import numpy as np
-import random
-import warnings
-
-import torch
-from monai.data import Dataset, CacheDataset
-from monai.transforms import Compose, RandSpatialCropd, CenterSpatialCropd, LoadImaged, EnsureChannelFirstd
+from typing import List, Optional, Tuple
+
+from monai.data import CacheDataset
+from monai.transforms import Compose, RandSpatialCropd, CenterSpatialCropd
 from monai.utils import ensure_tuple_rep
 
 from .dataset_base import MonaiConnectomicsDataset
-from monai.config import KeysCollection
-from monai.transforms import MapTransform
 from ..io.monai_transforms import LoadVolumed
 
 
@@ -210,8 +204,8 @@ def __init__(
         sample_size = kwargs.get('sample_size', (32, 256, 256))
         mode = kwargs.get('mode', 'train')
         do_2d = kwargs.get('do_2d', False)
-        data_mean = kwargs.get('data_mean', 0.5)
-        data_std = kwargs.get('data_std', 0.5)
+        kwargs.get('data_mean', 0.5)
+        kwargs.get('data_std', 0.5)
         transpose_axes = kwargs.get('transpose_axes', None)
 
         # Create data dictionaries