update mito documentation

Author: Donglai Wei

connectomics/config/hydra_config.py

@@ -181,6 +181,12 @@ class ModelConfig:
     rsunet_act_negative_slope: float = 0.01  # For LeakyReLU
     rsunet_act_init: float = 0.25  # For PReLU
 
+    # nnUNet-specific parameters (for loading pretrained models)
+    nnunet_checkpoint: Optional[str] = None  # Path to pretrained checkpoint (.pth file)
+    nnunet_plans: Optional[str] = None  # Path to plans.json file
+    nnunet_dataset: Optional[str] = None  # Path to dataset.json file
+    nnunet_device: str = "cuda"  # Device for model loading ('cuda' or 'cpu')
+
     # Deep supervision (supported by MedNeXt, RSUNet, and some MONAI models)
     deep_supervision: bool = False
     deep_supervision_weights: Optional[List[float]] = (
@@ -872,7 +878,10 @@ class SlidingWindowConfig:
 
 @dataclass
 class TestTimeAugmentationConfig:
-    """Test-time augmentation configuration."""
+    """Test-time augmentation configuration.
+
+    Note: Saving predictions is now handled by SavePredictionConfig.
+    """
 
     enabled: bool = False
     flip_axes: Any = (
@@ -886,12 +895,24 @@ class TestTimeAugmentationConfig:
     )
     ensemble_mode: str = "mean"  # Ensemble mode for TTA: 'mean', 'min', 'max'
     apply_mask: bool = False  # Multiply each channel by corresponding test_mask after ensemble
-    save_predictions: bool = (
-        False  # Save intermediate TTA predictions (before decoding) to disk (default: False)
-    )
-    save_dtype: Optional[str] = (
-        None  # Data type for saving predictions: "float16", "float32", "uint8", "uint16", or None (keep original)
-    )
+
+
+@dataclass
+class SavePredictionConfig:
+    """Configuration for saving intermediate predictions during inference.
+
+    Controls how raw model predictions are saved before any decoding or postprocessing.
+    Useful for debugging, visualization, or running multiple decoding strategies.
+
+    Attributes:
+        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')
+    """
+
+    enabled: bool = True  # Enable saving intermediate predictions
+    intensity_scale: float = 255.0  # Scale predictions to [0, 255] for saving
+    intensity_dtype: str = "uint8"  # Save as uint8 for visualization
 
 
 @dataclass
@@ -971,28 +992,18 @@ class ConnectedComponentsConfig:
 class PostprocessingConfig:
     """Postprocessing configuration for inference output.
 
-    Controls how predictions are transformed before saving:
+    Controls how predictions are transformed after saving:
     - Binary refinement: Morphological operations and connected components filtering
-    - Scaling: Multiply intensity values (e.g., 255 for uint8)
-    - Dtype conversion: Convert to target data type with proper clamping
     - Transpose: Reorder axes (e.g., [2,1,0] for zyx->xyz)
+
+    Note: Intensity scaling and dtype conversion are handled by SavePredictionConfig.
     """
 
     # Binary segmentation refinement (morphological ops, connected components)
     binary: Optional[Dict[str, Any]] = field(
         default_factory=dict
     )  # Binary postprocessing config (e.g., {'opening_iterations': 2})
 
-    # Intensity scaling
-    intensity_scale: Optional[float] = (
-        None  # Scale predictions for saving (e.g., 255.0 for uint8). None = no scaling
-    )
-
-    # Data type conversion
-    intensity_dtype: Optional[str] = (
-        None  # Output data type: 'uint8', 'uint16', 'float32'. None = no conversion (keep as-is)
-    )
-
     # Axis permutation
     output_transpose: List[int] = field(
         default_factory=list
@@ -1017,6 +1028,7 @@ class InferenceConfig:
     Key Features:
     - Sliding window inference for large volumes
     - Test-time augmentation (TTA) support
+    - Saving intermediate predictions
     - Multiple decoding strategies
     - Postprocessing and evaluation
     - System resource overrides for inference
@@ -1028,6 +1040,7 @@ class InferenceConfig:
     test_time_augmentation: TestTimeAugmentationConfig = field(
         default_factory=TestTimeAugmentationConfig
     )
+    save_prediction: SavePredictionConfig = field(default_factory=SavePredictionConfig)
     decoding: Optional[List[DecodeModeConfig]] = None  # List of decode modes to apply sequentially
     postprocessing: PostprocessingConfig = field(default_factory=PostprocessingConfig)
     evaluation: EvaluationConfig = field(default_factory=EvaluationConfig)
@@ -1255,6 +1268,7 @@ def configure_instance_segmentation(cfg: Config, boundary_thickness: int = 5) ->
     "InferenceDataConfig",
     "SlidingWindowConfig",
     "TestTimeAugmentationConfig",
+    "SavePredictionConfig",
     "PostprocessingConfig",
     "BinaryPostprocessingConfig",
     "ConnectedComponentsConfig",

connectomics/inference/__init__.py

@@ -1,12 +1,19 @@
 """Inference utilities package."""
 
 from .manager import InferenceManager
-from .io import apply_postprocessing, apply_decode_mode, resolve_output_filenames, write_outputs
+from .io import (
+    apply_save_prediction_transform,
+    apply_postprocessing,
+    apply_decode_mode,
+    resolve_output_filenames,
+    write_outputs,
+)
 from .sliding import build_sliding_inferer, resolve_inferer_roi_size, resolve_inferer_overlap
 from .tta import TTAPredictor
 
 __all__ = [
     "InferenceManager",
+    "apply_save_prediction_transform",
     "apply_postprocessing",
     "apply_decode_mode",
     "resolve_output_filenames",

connectomics/inference/io.py

@@ -16,14 +16,73 @@
 from ..config import Config
 
 
+def apply_save_prediction_transform(cfg: Config | DictConfig, data: np.ndarray) -> np.ndarray:
+    """
+    Apply intensity scaling and dtype conversion from save_prediction config.
+
+    This is used when saving intermediate predictions (before decoding).
+
+    Args:
+        cfg: Configuration object
+        data: Predictions array to transform
+
+    Returns:
+        Transformed predictions with applied scaling and dtype conversion
+    """
+    if not hasattr(cfg, "inference") or not hasattr(cfg.inference, "save_prediction"):
+        return data
+
+    save_pred_cfg = cfg.inference.save_prediction
+
+    # Apply intensity scaling
+    intensity_scale = getattr(save_pred_cfg, "intensity_scale", None)
+    if intensity_scale is not None and intensity_scale != 1.0:
+        data = data * float(intensity_scale)
+
+    # Apply dtype conversion
+    target_dtype_str = getattr(save_pred_cfg, "intensity_dtype", None)
+    if target_dtype_str is not None:
+        dtype_map = {
+            "uint8": np.uint8,
+            "int8": np.int8,
+            "uint16": np.uint16,
+            "int16": np.int16,
+            "uint32": np.uint32,
+            "int32": np.int32,
+            "float16": np.float16,
+            "float32": np.float32,
+            "float64": np.float64,
+        }
+
+        if target_dtype_str not in dtype_map:
+            warnings.warn(
+                f"Unknown dtype '{target_dtype_str}' in save_prediction config. "
+                f"Supported: {list(dtype_map.keys())}. Keeping current dtype.",
+                UserWarning,
+            )
+            return data
+
+        target_dtype = dtype_map[target_dtype_str]
+
+        # Get dtype info for proper clamping
+        if np.issubdtype(target_dtype, np.integer):
+            info = np.iinfo(target_dtype)
+            data = np.clip(data, info.min, info.max)
+
+        data = data.astype(target_dtype)
+
+    return data
+
+
 def apply_postprocessing(cfg: Config | DictConfig, data: np.ndarray) -> np.ndarray:
     """
     Apply postprocessing transformations to predictions.
 
-    This method applies (in order):
+    This method applies:
     1. Binary postprocessing (morphological operations, connected components filtering)
-    2. Scaling (intensity_scale or output_scale): Multiply predictions by scale factor
-    3. Dtype conversion (intensity_dtype or output_dtype): Convert to target dtype with clamping
+    2. Axis transposition (output_transpose)
+
+    Note: Intensity scaling and dtype conversion are handled by apply_save_prediction_transform()
     """
     if not hasattr(cfg, "inference") or not hasattr(cfg.inference, "postprocessing"):
         return data
@@ -82,53 +141,16 @@ def apply_postprocessing(cfg: Config | DictConfig, data: np.ndarray) -> np.ndarr
 
         data = np.stack(results, axis=0)
 
-    intensity_scale = getattr(postprocessing, "intensity_scale", None)
-    output_scale = getattr(postprocessing, "output_scale", None)
-    scale = intensity_scale if intensity_scale is not None else output_scale
-    if scale is not None:
-        data = data * float(scale)
-
-    target_dtype_str = getattr(postprocessing, "intensity_dtype", None)
-    if target_dtype_str is None:
-        target_dtype_str = getattr(postprocessing, "output_dtype", None)
-
-    if target_dtype_str is not None:
-        dtype_map = {
-            "uint8": np.uint8,
-            "int8": np.int8,
-            "uint16": np.uint16,
-            "int16": np.int16,
-            "uint32": np.uint32,
-            "int32": np.int32,
-            "float16": np.float16,
-            "float32": np.float32,
-            "float64": np.float64,
-        }
-
-        if target_dtype_str not in dtype_map:
+    # Apply axis transposition if configured
+    output_transpose = getattr(postprocessing, "output_transpose", [])
+    if output_transpose and len(output_transpose) > 0:
+        try:
+            data = np.transpose(data, axes=output_transpose)
+        except Exception as e:
             warnings.warn(
-                f"Unknown dtype '{target_dtype_str}'. Supported: {list(dtype_map.keys())}. "
-                f"Keeping float32.",
+                f"Transpose failed with axes {output_transpose}: {e}. Keeping original shape.",
                 UserWarning,
             )
-            return data
-
-        target_dtype = dtype_map[target_dtype_str]
-
-        if target_dtype_str == "uint8":
-            data = np.clip(data, 0, 255)
-        elif target_dtype_str == "int8":
-            data = np.clip(data, -128, 127)
-        elif target_dtype_str == "uint16":
-            data = np.clip(data, 0, 65535)
-        elif target_dtype_str == "int16":
-            data = np.clip(data, -32768, 32767)
-        elif target_dtype_str == "uint32":
-            data = np.clip(data, 0, 4294967295)
-        elif target_dtype_str == "int32":
-            data = np.clip(data, -2147483648, 2147483647)
-
-        data = data.astype(target_dtype)
 
     return data
 
@@ -342,8 +364,8 @@ def write_outputs(
         sample = np.squeeze(sample)
         write_hdf5(
             output_path,
-            "main",
             sample.astype(np.float32) if not np.issubdtype(sample.dtype, np.integer) else sample,
+            dataset="main",
         )
 
         print(f"  ✓ Saved prediction: {output_path}")

connectomics/models/arch/__init__.py

@@ -57,6 +57,13 @@ def build_my_model(cfg):
 except ImportError:
     _RSUNET_AVAILABLE = False
 
+# Import nnUNet models to trigger registration
+try:
+    from . import nnunet_models
+    _NNUNET_AVAILABLE = True
+except ImportError:
+    _NNUNET_AVAILABLE = False
+
 # Check what's available
 def get_available_architectures() -> dict:
     """
@@ -75,6 +82,7 @@ def get_available_architectures() -> dict:
         'monai': [a for a in all_archs if a.startswith('monai_')] if _MONAI_AVAILABLE else [],
         'mednext': [a for a in all_archs if a.startswith('mednext')] if _MEDNEXT_AVAILABLE else [],
         'rsunet': [a for a in all_archs if a.startswith('rsunet')] if _RSUNET_AVAILABLE else [],
+        'nnunet': [a for a in all_archs if a.startswith('nnunet')] if _NNUNET_AVAILABLE else [],
     }
 
     return info
@@ -107,6 +115,13 @@ def print_available_architectures():
         for arch in info['rsunet']:
             print(f"  - {arch}")
 
+    if info['nnunet']:
+        print(f"\nnnUNet Models ({len(info['nnunet'])}):")
+        for arch in info['nnunet']:
+            print(f"  - {arch}")
+    else:
+        print("\nnnUNet Models: Not available (install with: pip install nnunetv2)")
+
     print(f"\nTotal: {len(info['all'])} architectures")
     print("="*60 + "\n")