disable decoding and no saving

Author: Donglai Wei

connectomics/config/gpu_utils.py

@@ -7,6 +7,7 @@
 
 import torch
 import psutil
+import warnings
 from typing import Dict, Tuple
 
 
@@ -33,26 +34,39 @@ def get_gpu_info() -> Dict[str, any]:
     if not torch.cuda.is_available():
         return info
 
-    info["num_gpus"] = torch.cuda.device_count()
+    try:
+        info["num_gpus"] = torch.cuda.device_count()
+    except Exception as e:
+        warnings.warn(f"Failed to query CUDA device count: {e}")
+        return info
 
     for i in range(info["num_gpus"]):
-        # Get GPU name
-        info["gpu_names"].append(torch.cuda.get_device_name(i))
-
-        # Get memory info
-        props = torch.cuda.get_device_properties(i)
-        total_memory = props.total_memory / (1024**3)  # Convert to GB
-        info["total_memory_gb"].append(total_memory)
-
-        # Try to get available memory (may require GPU to be initialized)
         try:
-            torch.cuda.set_device(i)
-            torch.cuda.empty_cache()
-            available_memory = (props.total_memory - torch.cuda.memory_allocated(i)) / (1024**3)
-            info["available_memory_gb"].append(available_memory)
-        except Exception:
-            # Fallback: assume 90% is available
-            info["available_memory_gb"].append(total_memory * 0.9)
+            # Get GPU name
+            info["gpu_names"].append(torch.cuda.get_device_name(i))
+
+            # Get memory info
+            props = torch.cuda.get_device_properties(i)
+            total_memory = props.total_memory / (1024**3)  # Convert to GB
+            info["total_memory_gb"].append(total_memory)
+
+            # Try to get available memory (may require GPU to be initialized)
+            try:
+                torch.cuda.set_device(i)
+                torch.cuda.empty_cache()
+                available_memory = (props.total_memory - torch.cuda.memory_allocated(i)) / (1024**3)
+                info["available_memory_gb"].append(available_memory)
+            except Exception:
+                # Fallback: assume 90% is available
+                info["available_memory_gb"].append(total_memory * 0.9)
+        except Exception as e:
+            # Keep the process alive when a specific GPU index is invalid/broken.
+            # This is especially useful on clusters with dead devices or mismatched
+            # CUDA_VISIBLE_DEVICES / SLURM GPU cgroup mappings.
+            warnings.warn(f"Skipping GPU index {i} during CUDA probe: {e}")
+            info["gpu_names"].append(f"<unavailable:{i}>")
+            info["total_memory_gb"].append(0.0)
+            info["available_memory_gb"].append(0.0)
 
     return info
 

connectomics/config/hydra_config.py

@@ -204,6 +204,8 @@ class ModelConfig:
     )
     deep_supervision_clamp_min: float = -20.0  # Clamp logits to prevent numerical instability
     deep_supervision_clamp_max: float = 20.0  # Especially important at coarser scales
+    enable_nan_detection: bool = True  # Detect NaN/Inf losses and print diagnostics
+    debug_on_nan: bool = True  # Enter pdb on NaN/Inf detection (disable for non-interactive jobs)
 
     # Loss configuration
     loss_functions: List[str] = field(default_factory=lambda: ["DiceLoss", "BCEWithLogitsLoss"])
@@ -362,6 +364,8 @@ class NNUNetPreprocessingConfig:
     source_spacing: Optional[List[float]] = None  # If None, falls back to *_resolution fields
     normalization: str = "zscore"  # "zscore", "none", "0-1", or "divide-K"
     normalization_use_nonzero_mask: bool = True
+    clip_percentile_low: float = 0.0  # Optional clipping before normalization (fraction)
+    clip_percentile_high: float = 1.0  # Optional clipping before normalization (fraction)
     force_separate_z: Optional[bool] = None  # None = auto
     anisotropy_threshold: float = 3.0
     image_order: int = 3
@@ -894,7 +898,6 @@ class InferenceDataConfig:
     - Inference: outputs/experiment_name/YYYYMMDD_HHMMSS/inference/last.ckpt/{output_name}
     """
 
-    test_path: str = ""  # Base path for test data (e.g., "/path/to/dataset/test/")
     test_image: Any = None  # str, List[str], or None - Can be single file or list of files
     test_label: Any = None  # str, List[str], or None - Can be single file or list of files
     test_mask: Any = None  # str, List[str], or None - Optional mask for inference
@@ -1150,6 +1153,7 @@ class TestDataConfig:
 
     # These can be strings (single file), lists (multiple files), or None
     # Using Any to support both str and List[str] (OmegaConf doesn't support Union of containers)
+    test_path: str = ""  # Base path for test data (prepended to test_* paths if set)
     test_image: Any = None  # str, List[str], or None
     test_label: Any = None  # str, List[str], or None
     test_mask: Any = None  # str, List[str], or None

connectomics/data/augment/build.py

@@ -104,6 +104,8 @@ def build_train_transforms(
                 normalization_use_nonzero_mask=getattr(
                     nnunet_pre_cfg, "normalization_use_nonzero_mask", True
                 ),
+                clip_percentile_low=getattr(nnunet_pre_cfg, "clip_percentile_low", 0.0),
+                clip_percentile_high=getattr(nnunet_pre_cfg, "clip_percentile_high", 1.0),
                 force_separate_z=getattr(nnunet_pre_cfg, "force_separate_z", None),
                 anisotropy_threshold=getattr(nnunet_pre_cfg, "anisotropy_threshold", 3.0),
                 image_order=getattr(nnunet_pre_cfg, "image_order", 3),
@@ -350,6 +352,8 @@ def _build_eval_transforms_impl(cfg: Config, mode: str = "val", keys: list[str]
                 normalization_use_nonzero_mask=getattr(
                     nnunet_pre_cfg, "normalization_use_nonzero_mask", True
                 ),
+                clip_percentile_low=getattr(nnunet_pre_cfg, "clip_percentile_low", 0.0),
+                clip_percentile_high=getattr(nnunet_pre_cfg, "clip_percentile_high", 1.0),
                 force_separate_z=getattr(nnunet_pre_cfg, "force_separate_z", None),
                 anisotropy_threshold=getattr(nnunet_pre_cfg, "anisotropy_threshold", 3.0),
                 image_order=getattr(nnunet_pre_cfg, "image_order", 3),

connectomics/data/dataset/dataset_volume_cached.py

@@ -159,6 +159,7 @@ def __init__(
         patch_size: Tuple[int, int, int] = (112, 112, 112),
         iter_num: int = 500,
         transforms: Optional[Compose] = None,
+        pre_cache_transforms: Optional[Any] = None,
         mode: str = "train",
         pad_size: Optional[Tuple[int, ...]] = None,
         pad_mode: str = "reflect",
@@ -181,6 +182,7 @@ def __init__(
 
         self.iter_num = iter_num if iter_num > 0 else len(image_paths)
         self.transforms = transforms
+        self.pre_cache_transforms = pre_cache_transforms
         self.mode = mode
         self.pad_size = pad_size
         self.pad_mode = pad_mode
@@ -218,17 +220,8 @@ def __init__(
             elif img.ndim == 3:
                 img = img[None, ...]  # Add channel for 3D
 
-            # Apply padding if specified
-            if self.pad_size is not None:
-                img = self._apply_padding(img)
-
-            # Ensure volume is at least as large as patch_size in all dimensions
-            # This prevents crops from being smaller than patch_size
-            img = self._ensure_minimum_size(img)
-
-            self.cached_images.append(img)
-
             # Load label if available
+            lbl = None
             if lbl_path:
                 lbl = read_volume(lbl_path)
                 # Add channel dimension for both 2D and 3D
@@ -237,6 +230,37 @@ def __init__(
                 elif lbl.ndim == 3:
                     lbl = lbl[None, ...]  # Add channel for 3D
 
+            # Load mask if available
+            mask = None
+            if mask_path:
+                mask = read_volume(mask_path)
+                if mask.ndim == 2:
+                    mask = mask[None, ...]
+                elif mask.ndim == 3:
+                    mask = mask[None, ...]
+
+            # Apply one-time preprocessing before caching (e.g., nnU-Net crop/resample/normalize).
+            if self.pre_cache_transforms is not None:
+                sample = {"image": img}
+                if lbl is not None:
+                    sample["label"] = lbl
+                if mask is not None:
+                    sample["mask"] = mask
+                sample = self.pre_cache_transforms(sample)
+                img = sample["image"]
+                lbl = sample.get("label")
+                mask = sample.get("mask")
+
+            # Apply padding if specified
+            if self.pad_size is not None:
+                img = self._apply_padding(img)
+
+            # Ensure volume is at least as large as patch_size in all dimensions
+            # This prevents crops from being smaller than patch_size
+            img = self._ensure_minimum_size(img)
+            self.cached_images.append(img)
+
+            if lbl is not None:
                 # Apply padding if specified (same padding as image)
                 if self.pad_size is not None:
                     lbl = self._apply_padding(
@@ -245,24 +269,17 @@ def __init__(
 
                 # Ensure label is at least as large as patch_size
                 lbl = self._ensure_minimum_size(lbl, mode="constant", constant_values=0)
-
                 self.cached_labels.append(lbl)
             else:
                 self.cached_labels.append(None)
 
-            # Load mask if available
-            if mask_path:
-                mask = read_volume(mask_path)
-                if mask.ndim == 3:
-                    mask = mask[None, ...]
-
+            if mask is not None:
                 # Apply padding if specified (same padding as label)
                 if self.pad_size is not None:
                     mask = self._apply_padding(mask, mode="constant", constant_values=0)
 
                 # Ensure mask is at least as large as patch_size
                 mask = self._ensure_minimum_size(mask, mode="constant", constant_values=0)
-
                 self.cached_masks.append(mask)
             else:
                 self.cached_masks.append(None)

connectomics/data/io/monai_transforms.py

@@ -101,6 +101,7 @@ class NNUNetPreprocessd(MapTransform):
     This transform optionally applies:
     - foreground crop
     - spacing-aware resampling
+    - optional percentile clipping
     - intensity normalization
 
     It stores enough metadata in ``image_meta_dict['nnunet_preprocess']`` to
@@ -117,6 +118,8 @@ def __init__(
         target_spacing: Optional[Sequence[float]] = None,
         normalization: str = "zscore",
         normalization_use_nonzero_mask: bool = True,
+        clip_percentile_low: float = 0.0,
+        clip_percentile_high: float = 1.0,
         force_separate_z: Optional[bool] = None,
         anisotropy_threshold: float = 3.0,
         image_order: int = 3,
@@ -132,11 +135,27 @@ def __init__(
         self.target_spacing = list(target_spacing) if target_spacing is not None else None
         self.normalization = normalization
         self.normalization_use_nonzero_mask = normalization_use_nonzero_mask
+        self.clip_percentile_low = float(clip_percentile_low)
+        self.clip_percentile_high = float(clip_percentile_high)
         self.force_separate_z = force_separate_z
         self.anisotropy_threshold = anisotropy_threshold
         self.image_order = image_order
         self.label_order = label_order
         self.order_z = order_z
+        self._debug_stats_printed = False  # Print one-line pre/post stats once per transform instance
+        if not (0.0 <= self.clip_percentile_low <= 1.0):
+            raise ValueError(
+                f"clip_percentile_low must be in [0, 1], got {self.clip_percentile_low}"
+            )
+        if not (0.0 <= self.clip_percentile_high <= 1.0):
+            raise ValueError(
+                f"clip_percentile_high must be in [0, 1], got {self.clip_percentile_high}"
+            )
+        if self.clip_percentile_low > self.clip_percentile_high:
+            raise ValueError(
+                "clip_percentile_low must be <= clip_percentile_high, got "
+                f"{self.clip_percentile_low} > {self.clip_percentile_high}"
+            )
 
     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
         d = dict(data)
@@ -164,6 +183,8 @@ def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
             "target_spacing": target_spacing,
             "normalization": self.normalization,
             "normalization_use_nonzero_mask": self.normalization_use_nonzero_mask,
+            "clip_percentile_low": self.clip_percentile_low,
+            "clip_percentile_high": self.clip_percentile_high,
             "crop_bbox": None,
             "cropped_spatial_shape": list(image_spatial_shape),
             "resampled_spatial_shape": list(image_spatial_shape),
@@ -230,13 +251,37 @@ def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
                     self._get_spatial_shape(image_np, spatial_dims)
                 )
 
-        # Normalize image after spatial transforms.
+        pre_stats = None
+        if not self._debug_stats_printed:
+            pre_stats = self._format_debug_stats(image_np)
+
+        # Optional clipping before normalization (mirrors SmartNormalizeIntensityd order).
+        image_np = self._clip_image_percentiles(
+            image_np,
+            spatial_dims=spatial_dims,
+            low=self.clip_percentile_low,
+            high=self.clip_percentile_high,
+            use_nonzero_mask=self.normalization_use_nonzero_mask,
+        )
+
+        # Normalize image after spatial transforms and clipping.
         image_np = self._normalize_image(
             image_np,
             spatial_dims=spatial_dims,
             mode=self.normalization,
             use_nonzero_mask=self.normalization_use_nonzero_mask,
         )
+        if not self._debug_stats_printed:
+            post_stats = self._format_debug_stats(image_np)
+            print(
+                "[NNUNetPreprocessd] "
+                f"key={self.image_key} mode={self.normalization} "
+                f"clip=({self.clip_percentile_low:.3f},{self.clip_percentile_high:.3f}) "
+                f"nonzero_mask={self.normalization_use_nonzero_mask} "
+                f"pre={pre_stats} post={post_stats}",
+                flush=True,
+            )
+            self._debug_stats_printed = True
         d[self.image_key] = self._from_numpy(image_np, image_state)
 
         meta_dict["nnunet_preprocess"] = preprocess_meta
@@ -299,6 +344,20 @@ def _normalize_spacing(
             return values[-spatial_dims:]
         return None
 
+    @staticmethod
+    def _format_debug_stats(array: np.ndarray) -> str:
+        arr = np.asarray(array, dtype=np.float32)
+        if arr.size == 0:
+            return "empty"
+        finite = arr[np.isfinite(arr)]
+        if finite.size == 0:
+            return "all_nonfinite"
+        return (
+            f"shape={tuple(int(v) for v in arr.shape)} "
+            f"min={float(finite.min()):.4f} max={float(finite.max()):.4f} "
+            f"mean={float(finite.mean()):.4f} std={float(finite.std()):.4f}"
+        )
+
     @staticmethod
     def _create_nonzero_mask(image: np.ndarray, spatial_dims: int) -> Optional[np.ndarray]:
         if image.ndim == spatial_dims + 1:
@@ -449,6 +508,48 @@ def _resample_spatial(
             )
         return stacked
 
+    @staticmethod
+    def _clip_image_percentiles(
+        image: np.ndarray,
+        spatial_dims: int,
+        low: float,
+        high: float,
+        use_nonzero_mask: bool,
+    ) -> np.ndarray:
+        if low <= 0.0 and high >= 1.0:
+            return image
+
+        image = image.astype(np.float32, copy=False)
+        low_q = float(low) * 100.0
+        high_q = float(high) * 100.0
+
+        def _apply(volume: np.ndarray, mask: Optional[np.ndarray]) -> np.ndarray:
+            if mask is not None and np.any(mask):
+                values = volume[mask]
+            else:
+                values = volume.reshape(-1)
+            if values.size == 0:
+                return volume
+            low_val = float(np.percentile(values, low_q))
+            high_val = float(np.percentile(values, high_q))
+            if high_val < low_val:
+                low_val, high_val = high_val, low_val
+            return np.clip(volume, low_val, high_val)
+
+        if image.ndim == spatial_dims + 1:
+            nonzero_mask = np.any(image != 0, axis=0) if use_nonzero_mask else None
+            for c in range(image.shape[0]):
+                image[c] = _apply(image[c], nonzero_mask)
+                if nonzero_mask is not None:
+                    image[c][~nonzero_mask] = 0
+            return image
+
+        mask = image != 0 if use_nonzero_mask else None
+        image = _apply(image, mask)
+        if mask is not None:
+            image[~mask] = 0
+        return image
+
     @staticmethod
     def _normalize_image(
         image: np.ndarray,

connectomics/data/process/segment.py

@@ -2,12 +2,12 @@
 Segmentation processing functions for PyTorch Connectomics.
 """
 
-from __future__ import print_function, division
-from typing import Optional, Union, List
+from typing import List, Optional, Union
+
+import cc3d
 import numpy as np
 import torch
 from skimage.morphology import binary_dilation, dilation, erosion
-import cc3d
 
 RATES_TYPE = Optional[Union[List[int], int]]