Merge pull request #182 from QiongWang1/master

Refactor decoding with minimal changes and reuse of existing pipeline

Author: donglaiw

connectomics/config/hydra_config.py

@@ -440,6 +440,7 @@ class DataConfig:
 
     # Sampling (for volumetric datasets)
     iter_num_per_epoch: Optional[int] = None  # Alias for iter_num (if set, overrides iter_num)
+    val_iter_num: Optional[int] = None  # Validation iterations per epoch (auto-calculated if None)
     use_preloaded_cache: bool = (
         True  # Preload volumes into memory for fast random cropping (default: True)
     )
@@ -480,6 +481,10 @@ class SchedulerConfig:
     warmup_start_lr: float = 0.0001
     min_lr: float = 0.00001
 
+    # Scheduler interval control
+    interval: str = "epoch"  # "epoch" or "step" - controls when scheduler steps
+    frequency: int = 1  # How often to step the scheduler
+
     # CosineAnnealing-specific
     t_max: Optional[int] = None
 
@@ -941,6 +946,9 @@ class SavePredictionConfig:
         enabled: Enable saving intermediate predictions (default: True)
         intensity_scale: Scale factor for predictions (e.g., 255 for uint8 visualization)
         intensity_dtype: Data type for saved predictions (e.g., 'uint8', 'float32')
+        output_formats: List of output formats to save predictions in (e.g., ['h5', 'tiff', 'nii.gz'])
+                       Supported formats: 'h5', 'tiff', 'nii', 'nii.gz', 'png'
+                       Default: ['h5', 'nii.gz'] for backward compatibility
     """
 
     enabled: bool = True  # Enable saving intermediate predictions
@@ -951,6 +959,9 @@ class SavePredictionConfig:
     intensity_dtype: str = (
         "uint8"  # Save as uint8 for visualization (ignored if intensity_scale < 0)
     )
+    output_formats: List[str] = field(
+        default_factory=lambda: ["h5", "nii.gz"]  # Default: HDF5 + NIfTI for backward compatibility
+    )
 
 
 @dataclass

connectomics/config/hydra_utils.py

@@ -183,8 +183,14 @@ def validate_config(cfg: Config) -> None:
         raise ValueError("optimization.optimizer.weight_decay must be non-negative")
 
     # Training validation
-    if cfg.optimization.max_epochs <= 0:
-        raise ValueError("optimization.max_epochs must be positive")
+    # [FIX 2] Allow max_epochs to be 0 or negative when using step-based training
+    max_steps_cfg = getattr(cfg.optimization, "max_steps", None)
+    if max_steps_cfg is None or max_steps_cfg <= 0:
+        # Epoch-based training: max_epochs must be positive
+        if cfg.optimization.max_epochs <= 0:
+            raise ValueError("optimization.max_epochs must be positive when max_steps is not set")
+    # If max_steps is set, max_epochs can be anything (will be overridden to -1 in trainer)
+    
     if cfg.optimization.gradient_clip_val < 0:
         raise ValueError("optimization.gradient_clip_val must be non-negative")
     if cfg.optimization.accumulate_grad_batches <= 0:

connectomics/data/augment/build.py

@@ -196,7 +196,7 @@ def build_train_transforms(
     return Compose(transforms)
 
 
-def _build_eval_transforms_impl(cfg: Config, mode: str = "val", keys: list[str] = None) -> Compose:
+def _build_eval_transforms_impl(cfg: Config, mode: str = "val", keys: list[str] = None, skip_loading: bool = False) -> Compose:
     """
     Internal implementation for building evaluation transforms (validation or test).
 
@@ -207,6 +207,7 @@ def _build_eval_transforms_impl(cfg: Config, mode: str = "val", keys: list[str]
         cfg: Hydra Config object
         mode: 'val' or 'test' mode
         keys: Keys to transform (default: auto-detected based on mode)
+        skip_loading: Skip LoadVolumed (for pre-cached datasets)
 
     Returns:
         Composed MONAI transforms (no augmentation)
@@ -259,32 +260,34 @@ def _build_eval_transforms_impl(cfg: Config, mode: str = "val", keys: list[str]
     transforms = []
 
     # Load images first - use appropriate loader based on dataset type
-    dataset_type = getattr(cfg.data, "dataset_type", "volume")
-
-    if dataset_type == "filename":
-        # For filename-based datasets (PNG, JPG, etc.), use MONAI's LoadImaged
-        transforms.append(LoadImaged(keys=keys, image_only=False))
-        # Ensure channel-first format [C, H, W] or [C, D, H, W]
-        transforms.append(EnsureChannelFirstd(keys=keys))
-    else:
-        # For volume-based datasets (HDF5, TIFF volumes), use custom LoadVolumed
-        # Get transpose axes based on mode
-        if mode == "val":
-            transpose_axes = cfg.data.val_transpose if cfg.data.val_transpose else []
-        else:  # mode == "test"
-            # Use test.data.test_transpose
-            transpose_axes = []
-            if (
-                hasattr(cfg, "test")
-                and hasattr(cfg.test, "data")
-                and hasattr(cfg.test.data, "test_transpose")
-                and cfg.test.data.test_transpose
-            ):
-                transpose_axes = cfg.test.data.test_transpose
+    # Skip loading if using pre-cached datasets
+    if not skip_loading:
+        dataset_type = getattr(cfg.data, "dataset_type", "volume")
 
-        transforms.append(
-            LoadVolumed(keys=keys, transpose_axes=transpose_axes if transpose_axes else None)
-        )
+        if dataset_type == "filename":
+            # For filename-based datasets (PNG, JPG, etc.), use MONAI's LoadImaged
+            transforms.append(LoadImaged(keys=keys, image_only=False))
+            # Ensure channel-first format [C, H, W] or [C, D, H, W]
+            transforms.append(EnsureChannelFirstd(keys=keys))
+        else:
+            # For volume-based datasets (HDF5, TIFF volumes), use custom LoadVolumed
+            # Get transpose axes based on mode
+            if mode == "val":
+                transpose_axes = cfg.data.val_transpose if cfg.data.val_transpose else []
+            else:  # mode == "test"
+                # Use test.data.test_transpose
+                transpose_axes = []
+                if (
+                    hasattr(cfg, "test")
+                    and hasattr(cfg.test, "data")
+                    and hasattr(cfg.test.data, "test_transpose")
+                    and cfg.test.data.test_transpose
+                ):
+                    transpose_axes = cfg.test.data.test_transpose
+
+            transforms.append(
+                LoadVolumed(keys=keys, transpose_axes=transpose_axes if transpose_axes else None)
+            )
 
     # Apply volumetric split if enabled
     if cfg.data.split_enabled:
@@ -441,18 +444,19 @@ def _build_eval_transforms_impl(cfg: Config, mode: str = "val", keys: list[str]
     return Compose(transforms)
 
 
-def build_val_transforms(cfg: Config, keys: list[str] = None) -> Compose:
+def build_val_transforms(cfg: Config, keys: list[str] = None, skip_loading: bool = False) -> Compose:
     """
     Build validation transforms from Hydra config.
 
     Args:
         cfg: Hydra Config object
         keys: Keys to transform (default: auto-detected as ['image', 'label'])
+        skip_loading: Skip LoadVolumed (for pre-cached datasets)
 
     Returns:
         Composed MONAI transforms (no augmentation, center cropping)
     """
-    return _build_eval_transforms_impl(cfg, mode="val", keys=keys)
+    return _build_eval_transforms_impl(cfg, mode="val", keys=keys, skip_loading=skip_loading)
 
 
 def build_test_transforms(cfg: Config, keys: list[str] = None) -> Compose:

connectomics/data/augment/monai_transforms.py

@@ -24,6 +24,7 @@ class RandMisAlignmentd(RandomizableTransform, MapTransform):
     Simulates section misalignment artifacts common in EM volumes.
     """
 
+
     def __init__(
         self,
         keys: KeysCollection,
@@ -1141,10 +1142,24 @@ def _normalize(
         self, volume: Union[np.ndarray, torch.Tensor]
     ) -> Union[np.ndarray, torch.Tensor]:
         """Apply normalization to volume."""
+        from ...utils.debug_utils import print_tensor_stats
+        
         is_numpy = isinstance(volume, np.ndarray)
         if not is_numpy:
             volume = volume.numpy()
 
+        # DEBUG: Print raw input before normalization
+        print_tensor_stats(
+            volume,
+            stage_name="STAGE 1: RAW IMAGE (before normalization)",
+            tensor_name="image",
+            print_once=True,
+            extra_info={
+                "normalization_mode": self.mode,
+                "clip_percentiles": f"[{self.clip_percentile_low}, {self.clip_percentile_high}]"
+            }
+        )
+
         # Step 1: Percentile clipping (if enabled by non-default values)
         if self.clip_percentile_low > 0.0 or self.clip_percentile_high < 1.0:
             low_val = np.percentile(volume, self.clip_percentile_low * 100)
@@ -1171,6 +1186,18 @@ def _normalize(
             # Simple divide by K (e.g., divide-255 for uint8 images)
             volume = volume / self.divide_value
 
+        # DEBUG: Print after normalization
+        print_tensor_stats(
+            volume,
+            stage_name="STAGE 2: AFTER IMAGE NORMALIZATION",
+            tensor_name="image",
+            print_once=True,
+            extra_info={
+                "normalization_applied": self.mode,
+                "expected_range": "[0, 1]" if self.mode == "0-1" else "varies"
+            }
+        )
+
         return volume if is_numpy else torch.from_numpy(volume)
 
 

connectomics/data/dataset/dataset_base.py

@@ -212,7 +212,24 @@ def __init__(
             self.dataset_length = self.iter_num
         else:
             self.dataset_length = len(data_dicts)
+        
+        # [FIX] Add validation reseeding support
+        self.base_seed = 0
+        self.current_epoch = 0
 
+    def __getitem__(self, index: int) -> Dict[str, Any]:
+        """
+        Get a data sample with caching.
+        
+        When iter_num > len(data), we need to map the requested index
+        to an actual data index by using modulo operation.
+        """
+        # Map the requested index to actual data index
+        actual_index = index % len(self.data)
+        
+        # Call parent's __getitem__ with the mapped index
+        return super().__getitem__(actual_index)
+    
     def __len__(self) -> int:
         """
         Return dataset length.
@@ -226,13 +243,77 @@ def __len__(self) -> int:
             # Partial caching: return cached length for validation
             # For training with iter_num, we still want to iterate iter_num times
             if self.mode == 'train' and self.iter_num > 0:
-                return self.dataset_length
+                result = self.dataset_length
             else:
                 # For validation/test, only iterate over cached items
-                return len(self._cache)
+                result = len(self._cache)
+        else:
+            # Full caching or no caching: use dataset_length
+            result = self.dataset_length
 
-        # Full caching or no caching: use dataset_length
-        return self.dataset_length
+        return result
+    
+    def set_epoch(self, epoch: int, base_seed: int = 0):
+        """
+        Set current epoch for epoch-based validation reseeding.
+        
+        This method enables validation to sample different patches each epoch
+        while maintaining determinism. For training, this has no effect since
+        training already uses random sampling.
+        
+        Args:
+            epoch: Current training epoch
+            base_seed: Base random seed (typically from cfg.system.seed)
+        
+        Usage:
+            Called by ValidationReseedingCallback at the start of each validation epoch.
+        """
+        if self.mode == "val":
+            import random
+            self.base_seed = base_seed
+            self.current_epoch = epoch
+            effective_seed = self.base_seed + epoch
+            random.seed(effective_seed)
+            
+            # IMPORTANT: Print to verify reseeding is happening
+            print(f"[Validation] Set epoch={epoch}, base_seed={base_seed}, effective_seed={effective_seed}")
+            print(f"[Validation] Dataset: {type(self).__name__}@{id(self)}, mode={self.mode}, iter_num={self.iter_num}")
+    
+    def get_sampling_fingerprint(self, num_samples: int = 5) -> str:
+        """
+        Generate a deterministic fingerprint of validation sampling.
+        
+        This allows verification that validation patches change across epochs.
+        For MonaiCachedConnectomicsDataset, we sample indices that would be used.
+        
+        Args:
+            num_samples: Number of random samples to include in fingerprint
+        
+        Returns:
+            String representing the sampling fingerprint
+        """
+        if self.mode != "val":
+            return "N/A (training mode)"
+        
+        import random
+        # Save current RNG state
+        state = random.getstate()
+        
+        try:
+            # Generate deterministic samples
+            samples = []
+            for _ in range(num_samples):
+                # Sample index (same logic as __getitem__)
+                idx = random.randint(0, len(self.data) - 1)
+                samples.append(idx)
+            
+            # Create fingerprint string
+            fingerprint = ", ".join([f"idx{i}" for i in samples])
+            return fingerprint
+        
+        finally:
+            # Restore RNG state
+            random.setstate(state)
 
 
 class MonaiPersistentConnectomicsDataset(PersistentDataset):