Normalize package layout and add MedNeXt multi-head support

Author: Donglai Wei

.claude/feature/mednext_multi_head.md

@@ -39,7 +39,7 @@
     SlidingWindowConfig,
     TestTimeAugmentationConfig,
 )
-from .model import LossConfig, ModelArchConfig, ModelConfig
+from .model import LossConfig, ModelArchConfig, ModelConfig, ModelHeadConfig
 from .model_mednext import MedNeXtConfig
 from .model_monai import MonaiConfig, TransformerConfig
 from .model_nnunet import NNUNetConfig
@@ -80,6 +80,7 @@
     # Model configuration
     "ModelConfig",
     "ModelArchConfig",
+    "ModelHeadConfig",
     "MonaiConfig",
     "TransformerConfig",
     "MedNeXtConfig",

connectomics/config/pipeline/config_io.py

@@ -442,6 +442,9 @@ class DataConfig:
     - 2D data support with do_2d parameter
     """
 
+    # Root path prepended to train/val/test split paths (empty = no prefix)
+    root_path: str = ""
+
     # Train/Val Split (inspired by DeepEM)
     split_enabled: bool = False  # Enable automatic train/val split (default: False)
     split_train_range: List[float] = field(default_factory=lambda: [0.0, 0.8])  # Train: 0-80%

connectomics/config/schema/__init__.py

@@ -79,6 +79,7 @@ class SavePredictionConfig:
     output_formats: List[str] = field(default_factory=lambda: ["h5"])  # Any of: h5, tiff, png
     output_path: Optional[str] = None
     cache_suffix: str = "_x1_prediction.h5"
+    save_all_heads: bool = False
 
     # Data scaling and output typing
     # -1 keeps native float probabilities/logits; >0 scales and casts to integer dtype if chosen.
@@ -189,6 +190,9 @@ class InferenceConfig:
     Note: stage-specific overrides are merged before runtime; consumers should read `cfg.inference`.
     """
 
+    # Named output head selection for multi-head models. When unset, falls back
+    # to model.primary_head or the sole configured head.
+    head: Optional[str] = None
     sliding_window: SlidingWindowConfig = field(default_factory=SlidingWindowConfig)
     test_time_augmentation: TestTimeAugmentationConfig = field(
         default_factory=TestTimeAugmentationConfig

connectomics/config/schema/data.py

@@ -64,6 +64,16 @@ class ModelArchConfig:
     params: Dict[str, Any] = field(default_factory=dict)
 
 
+@dataclass
+class ModelHeadConfig:
+    """Named task head configuration for MedNeXt multi-head models."""
+
+    out_channels: int = 1
+    num_blocks: int = 0
+    hidden_channels: Optional[int] = None
+    target_slice: Optional[Any] = None
+
+
 @dataclass
 class ModelConfig:
     """Model architecture configuration.
@@ -87,7 +97,12 @@ class ModelConfig:
     input_size: List[int] = field(default_factory=lambda: [128, 128, 128])
     output_size: List[int] = field(default_factory=lambda: [128, 128, 128])
     in_channels: int = 1
+    # Legacy global output width used by single-tensor models and some fallback
+    # code paths. Multi-head MedNeXt models should define per-head widths in
+    # model.heads and should not rely on this to mirror the sum of all heads.
     out_channels: int = 1
+    primary_head: Optional[str] = None
+    heads: Dict[str, ModelHeadConfig] = field(default_factory=dict)
 
     # Architecture-specific nested blocks
     monai: MonaiConfig = field(default_factory=MonaiConfig)

connectomics/config/schema/inference.py

@@ -58,6 +58,7 @@ class ImageLoggingConfig:
     channels: Optional[Tuple[int, ...]] = None
     channel_mode: str = "all"  # "argmax", "all", or "selected"
     selected_channels: Optional[List[int]] = None
+    head: Optional[str] = None
 
 
 @dataclass

connectomics/config/schema/model.py

@@ -1,26 +1,23 @@
-"""
-Data package for PyTorch Connectomics.
+"""Data package for PyTorch Connectomics.
 
 This package provides:
-- Dataset classes (dataset/)
-- Data augmentation (augment/)
-- Data processing transforms (process/)
+- Dataset classes (datasets/)
+- Data augmentation (augmentation/)
+- Data processing transforms (processing/)
 - I/O utilities (io/)
 - DataModules for PyTorch Lightning (see training/lightning/data.py)
 
 Recommended imports:
-    from connectomics.data.dataset import CachedVolumeDataset
-    from connectomics.data.augment import RandMisAlignmentd, build_train_transforms
-    from connectomics.data.process import MultiTaskLabelTransformd, create_label_transform_pipeline
+    from connectomics.data.datasets import CachedVolumeDataset
+    from connectomics.data.augmentation import RandMisAlignmentd, build_train_transforms
+    from connectomics.data.processing import MultiTaskLabelTransformd, create_label_transform_pipeline
 """
 
-# Make submodules available
-from . import augment, dataset, io, process
+from . import augmentation, datasets, io, processing
 
 __all__ = [
-    # Submodules
-    "augment",
-    "dataset",
+    "augmentation",
+    "datasets",
     "io",
-    "process",
+    "processing",
 ]