load external mednext model

Author: Donglai Wei

connectomics/config/hydra_config.py

@@ -202,6 +202,11 @@ class ModelConfig:
         None  # None = single task (apply all losses to all channels)
     )
 
+    # External model weights loading
+    # For loading pretrained weights from external checkpoints (e.g., BANIS, nnUNet)
+    external_weights_path: Optional[str] = None  # Path to external checkpoint file
+    external_weights_key_prefix: str = "model."  # Prefix to strip from state_dict keys
+
 
 # Label transformation configurations
 @dataclass
@@ -858,6 +863,7 @@ class SlidingWindowConfig:
     )
     padding_mode: str = "constant"  # Padding mode at volume boundaries
     pad_size: Optional[List[int]] = None  # Padding size for context (e.g., [16, 32, 32])
+    save_channels: Optional[List[int]] = None  # Channel indices to save (None = all channels)
 
 
 @dataclass

connectomics/data/io/io.py

@@ -16,6 +16,7 @@
 import h5py
 import numpy as np
 import imageio
+import nibabel as nib
 
 # Avoid PIL "IOError: image file truncated"
 from PIL import ImageFile
@@ -263,7 +264,7 @@ def write_pickle_file(filename: str, data: object) -> None:
 def read_volume(
     filename: str, dataset: Optional[str] = None, drop_channel: bool = False
 ) -> np.ndarray:
-    """Load volumetric data in HDF5, TIFF or PNG formats.
+    """Load volumetric data in HDF5, TIFF, PNG, or NIfTI formats.
 
     Args:
         filename: Path to the volume file
@@ -276,7 +277,11 @@ def read_volume(
     Raises:
         ValueError: If file format is not recognized
     """
-    image_suffix = filename[filename.rfind(".") + 1 :].lower()
+    # Handle .nii.gz files specially
+    if filename.endswith('.nii.gz'):
+        image_suffix = 'nii.gz'
+    else:
+        image_suffix = filename[filename.rfind(".") + 1 :].lower()
 
     if image_suffix in ["h5", "hdf5"]:
         data = read_hdf5(filename, dataset)
@@ -314,9 +319,23 @@ def read_volume(
         if data.ndim == 4:
             # Convert (D, H, W, C) to (C, D, H, W) order
             data = data.transpose(3, 0, 1, 2)
+    elif image_suffix in ["nii", "nii.gz"]:
+        # NIfTI format (.nii or .nii.gz)
+        nii_img = nib.load(filename)
+        data = np.asarray(nii_img.dataobj)
+        # NIfTI is typically (X, Y, Z) or (X, Y, Z, C)
+        # Convert to our (D, H, W) or (C, D, H, W) format
+        # X=W (width), Y=H (height), Z=D (depth)
+        if data.ndim == 3:
+            # (X, Y, Z) -> (Z, Y, X) = (D, H, W)
+            data = data.transpose(2, 1, 0)
+        elif data.ndim == 4:
+            # (X, Y, Z, C) -> (C, Z, Y, X) = (C, D, H, W)
+            data = data.transpose(3, 2, 1, 0)
     else:
         raise ValueError(
-            f"Unrecognizable file format for {filename}. " f"Expected: h5, hdf5, tif, tiff, or png"
+            f"Unrecognizable file format for {filename}. "
+            f"Expected: h5, hdf5, tif, tiff, png, nii, or nii.gz"
         )
 
     # if data.ndim not in [3, 4]:
@@ -342,7 +361,7 @@ def save_volume(
         filename: Output filename or directory path
         volume: Volume data to save
         dataset: Dataset name for HDF5 format
-        file_format: Output format ('h5' or 'png')
+        file_format: Output format ('h5', 'png', 'nii', or 'nii.gz')
 
     Raises:
         ValueError: If file format is not supported
@@ -351,8 +370,21 @@ def save_volume(
         write_hdf5(filename, volume, dataset=dataset)
     elif file_format == "png":
         save_images(filename, volume)
+    elif file_format in ["nii", "nii.gz"]:
+        # NIfTI format
+        # Convert from our (D, H, W) or (C, D, H, W) to NIfTI (X, Y, Z) or (X, Y, Z, C)
+        if volume.ndim == 3:
+            # (D, H, W) -> (W, H, D) = (X, Y, Z)
+            nii_data = volume.transpose(2, 1, 0)
+        elif volume.ndim == 4:
+            # (C, D, H, W) -> (W, H, D, C) = (X, Y, Z, C)
+            nii_data = volume.transpose(3, 2, 1, 0)
+        else:
+            nii_data = volume
+        nii_img = nib.Nifti1Image(nii_data, affine=np.eye(4))
+        nib.save(nii_img, filename)
     else:
-        raise ValueError(f"Unsupported format: {file_format}. " f"Supported formats: h5, png")
+        raise ValueError(f"Unsupported format: {file_format}. " f"Supported formats: h5, png, nii, nii.gz")
 
 
 def get_vol_shape(filename: str, dataset: Optional[str] = None) -> tuple:
@@ -382,7 +414,11 @@ def get_vol_shape(filename: str, dataset: Optional[str] = None) -> tuple:
     if not os.path.exists(filename):
         raise FileNotFoundError(f"File not found: {filename}")
 
-    image_suffix = filename[filename.rfind(".") + 1 :].lower()
+    # Handle .nii.gz files specially
+    if filename.endswith('.nii.gz'):
+        image_suffix = 'nii.gz'
+    else:
+        image_suffix = filename[filename.rfind(".") + 1 :].lower()
 
     if image_suffix in ["h5", "hdf5"]:
         # HDF5: Read shape from metadata (no data loading)
@@ -424,9 +460,24 @@ def get_vol_shape(filename: str, dataset: Optional[str] = None) -> tuple:
         else:
             raise ValueError(f"Unsupported PNG dimensions: {first_image.ndim}D")
 
+    elif image_suffix in ["nii", "nii.gz"]:
+        # NIfTI: Read shape from header (no data loading)
+        nii_img = nib.load(filename)
+        nii_shape = nii_img.header.get_data_shape()
+        # Convert from NIfTI (X, Y, Z) or (X, Y, Z, C) to our (D, H, W) or (C, D, H, W)
+        if len(nii_shape) == 3:
+            # (X, Y, Z) -> (Z, Y, X) = (D, H, W)
+            return (nii_shape[2], nii_shape[1], nii_shape[0])
+        elif len(nii_shape) == 4:
+            # (X, Y, Z, C) -> (C, D, H, W)
+            return (nii_shape[3], nii_shape[2], nii_shape[1], nii_shape[0])
+        else:
+            return nii_shape
+
     else:
         raise ValueError(
-            f"Unrecognizable file format for {filename}. " f"Expected: h5, hdf5, tif, tiff, or png"
+            f"Unrecognizable file format for {filename}. "
+            f"Expected: h5, hdf5, tif, tiff, png, nii, or nii.gz"
         )
 
 

connectomics/decoding/segmentation.py

@@ -15,7 +15,7 @@
 """
 
 from __future__ import print_function, division
-from typing import Optional, Tuple
+from typing import Optional, Tuple, List
 import numpy as np
 import cc3d
 import fastremap
@@ -352,14 +352,17 @@ def decode_binary_contour_watershed(
 def decode_binary_contour_distance_watershed(
     predictions: np.ndarray,
     binary_threshold: Tuple[float, float] = (0.9, 0.85),
-    contour_threshold: Tuple[float, float] = (0.8, 1.1),
+    contour_threshold: Optional[Tuple[float, float]] = (0.8, 1.1),
     distance_threshold: Tuple[float, float] = (0.5, 0),
     min_instance_size: int = 128,
     remove_small_mode: str = "background",
     min_seed_size: int = 32,
     return_seed: bool = False,
     precomputed_seed: Optional[np.ndarray] = None,
     prediction_scale: int = 255,
+    binary_channels: Optional[List[int]] = None,
+    contour_channels: Optional[List[int]] = None,
+    distance_channels: Optional[List[int]] = None,
 ):
     r"""Convert binary foreground probability maps, instance contours and signed distance
     transform to instance masks via watershed segmentation algorithm.
@@ -369,48 +372,145 @@ def decode_binary_contour_distance_watershed(
         function that converts the input image into ``np.float64`` data type for processing. Therefore please make sure enough memory is allocated when handling large arrays.
 
     Args:
-        predictions (numpy.ndarray): foreground, contour, and distance probability of shape :math:`(3, Z, Y, X)`.
+        predictions (numpy.ndarray): foreground, contour, and distance probability of shape :math:`(3, Z, Y, X)`
+            or :math:`(2, Z, Y, X)` if contour is disabled.
         binary_threshold (tuple): tuple of two floats (seed_threshold, foreground_threshold) for binary mask.
             The first value is used for seed generation, the second for foreground mask. Default: (0.9, 0.85)
-        contour_threshold (tuple): tuple of two floats (seed_threshold, foreground_threshold) for instance contours.
-            The first value is used for seed generation, the second for foreground mask. Default: (0.8, 1.1)
+        contour_threshold (tuple or None): tuple of two floats (seed_threshold, foreground_threshold) for instance contours.
+            The first value is used for seed generation, the second for foreground mask.
+            Set to None to disable contour constraints (for BANIS-style binary+distance only). Default: (0.8, 1.1)
         distance_threshold (tuple): tuple of two floats (seed_threshold, foreground_threshold) for signed distance.
             The first value is used for seed generation, the second for foreground mask. Default: (0.5, -0.5)
         min_instance_size (int): minimum size threshold for instances to keep. Default: 128
         remove_small_mode (str): ``'background'``, ``'neighbor'`` or ``'none'``. Default: ``'background'``
         min_seed_size (int): minimum size of seed objects. Default: 32
         return_seed (bool): whether to return the seed map. Default: False
         precomputed_seed (numpy.ndarray, optional): precomputed seed map. Default: None
-        prediction_scale (int): scale of input predictions (255 for uint8 range). Default: 255
+        prediction_scale (int): scale of input predictions (255 for uint8 range, 1 for 0-1 range). Default: 255
+        binary_channels (list of int, optional): channel indices for binary mask. If multiple, they are averaged.
+            Default: None (uses position-based assignment)
+        contour_channels (list of int, optional): channel indices for contour. If multiple, they are averaged.
+            Default: None (uses position-based assignment)
+        distance_channels (list of int, optional): channel indices for distance. If multiple, they are averaged.
+            Default: None (uses position-based assignment)
 
     Returns:
         numpy.ndarray or tuple: Instance segmentation mask, or (mask, seed) if return_seed=True.
-    """
-    assert predictions.shape[0] == 3
-    binary, contour, distance = predictions[0], predictions[1], predictions[2]
 
+    Example:
+        >>> # Standard 3-channel (binary, contour, distance)
+        >>> seg = decode_binary_contour_distance_watershed(predictions)
+
+        >>> # BANIS-style 2-channel (binary, distance) - no contour
+        >>> seg = decode_binary_contour_distance_watershed(
+        ...     predictions,  # shape (2, Z, Y, X)
+        ...     binary_threshold=(0.5, 0.5),
+        ...     contour_threshold=None,  # Disable contour
+        ...     distance_threshold=(0.0, -1.0),
+        ...     prediction_scale=1,
+        ... )
+
+        >>> # Explicit channel selection with averaging
+        >>> seg = decode_binary_contour_distance_watershed(
+        ...     predictions,  # shape (3, Z, Y, X) with channels [aff_x, aff_y, SDT]
+        ...     binary_channels=[0, 1],  # Average channels 0 and 1 for binary
+        ...     contour_channels=None,   # No contour
+        ...     distance_channels=[2],   # Channel 2 for distance
+        ...     contour_threshold=None,
+        ...     prediction_scale=1,
+        ... )
+    """
+    # Check if contour is disabled
+    use_contour = contour_threshold is not None
+
+    # Extract channels using explicit selection or position-based fallback
+    if binary_channels is not None or distance_channels is not None:
+        # Explicit channel selection mode
+        if binary_channels is not None:
+            if len(binary_channels) > 1:
+                binary = predictions[binary_channels].mean(axis=0)
+            else:
+                binary = predictions[binary_channels[0]]
+        else:
+            binary = predictions[0]
+
+        if distance_channels is not None:
+            if len(distance_channels) > 1:
+                distance = predictions[distance_channels].mean(axis=0)
+            else:
+                distance = predictions[distance_channels[0]]
+        else:
+            distance = predictions[-1]
+
+        if use_contour:
+            if contour_channels is not None:
+                if len(contour_channels) > 1:
+                    contour = predictions[contour_channels].mean(axis=0)
+                else:
+                    contour = predictions[contour_channels[0]]
+            else:
+                # Default: assume contour is second-to-last if using contour
+                contour = predictions[-2]
+        else:
+            contour = None
+    else:
+        # Position-based fallback (legacy behavior)
+        if use_contour:
+            assert predictions.shape[0] >= 3, f"Expected at least 3 channels (binary, contour, distance), got {predictions.shape[0]}"
+            # If more than 3 channels, first N-2 channels are binary (average them)
+            if predictions.shape[0] > 3:
+                binary = predictions[:-2].mean(axis=0)
+                contour, distance = predictions[-2], predictions[-1]
+            else:
+                binary, contour, distance = predictions[0], predictions[1], predictions[2]
+        else:
+            assert predictions.shape[0] >= 2, f"Expected at least 2 channels (binary, distance) when contour disabled, got {predictions.shape[0]}"
+            # If more than 2 channels, first N-1 channels are binary (average them)
+            if predictions.shape[0] > 2:
+                binary = predictions[:-1].mean(axis=0)
+                distance = predictions[-1]
+            else:
+                binary, distance = predictions[0], predictions[1]
+            contour = None
+
+    # Convert thresholds based on prediction scale
     if prediction_scale == 255:
         distance = (distance / prediction_scale) * 2.0 - 1.0
-        binary_threshold = binary_threshold * prediction_scale
-        contour_threshold = contour_threshold * prediction_scale
-        distance_threshold = distance_threshold * prediction_scale
+        binary_threshold = (binary_threshold[0] * prediction_scale, binary_threshold[1] * prediction_scale)
+        if use_contour:
+            contour_threshold = (contour_threshold[0] * prediction_scale, contour_threshold[1] * prediction_scale)
+        distance_threshold = (distance_threshold[0] * prediction_scale, distance_threshold[1] * prediction_scale)
 
     if precomputed_seed is not None:
         seed = precomputed_seed
     else:  # compute the instance seeds
-        seed_map = (
-            (binary > binary_threshold[0])
-            * (contour < contour_threshold[0])
-            * (distance > distance_threshold[0])
-        )
+        if use_contour:
+            seed_map = (
+                (binary > binary_threshold[0])
+                * (contour < contour_threshold[0])
+                * (distance > distance_threshold[0])
+            )
+        else:
+            # No contour constraint - only binary and distance
+            seed_map = (
+                (binary > binary_threshold[0])
+                * (distance > distance_threshold[0])
+            )
         seed = cc3d.connected_components(seed_map)
         seed = remove_small_objects(seed, min_seed_size)
 
-    foreground = (
-        (binary > binary_threshold[1])
-        * (contour < contour_threshold[1])
-        * (distance > distance_threshold[1])
-    )
+    if use_contour:
+        foreground = (
+            (binary > binary_threshold[1])
+            * (contour < contour_threshold[1])
+            * (distance > distance_threshold[1])
+        )
+    else:
+        # No contour constraint - only binary and distance
+        foreground = (
+            (binary > binary_threshold[1])
+            * (distance > distance_threshold[1])
+        )
 
     segmentation = mahotas.cwatershed(-distance.astype(np.float64), seed)
     segmentation[~foreground] = (

connectomics/lightning/inference.py

@@ -802,6 +802,11 @@ def write_outputs(
     if hasattr(cfg.inference, "postprocessing"):
         output_transpose = getattr(cfg.inference.postprocessing, "output_transpose", [])
 
+    # Get save_channels from sliding_window config (to reduce memory/storage)
+    save_channels = None
+    if hasattr(cfg.inference, "sliding_window"):
+        save_channels = getattr(cfg.inference.sliding_window, "save_channels", None)
+
     # Determine actual batch size from predictions
     # Handle both batched (B, ...) and unbatched (...) predictions
     if predictions.ndim >= 4:
@@ -839,6 +844,20 @@ def write_outputs(
         filename = filenames[idx]
         output_path = output_dir / f"{filename}_{suffix}.h5"
 
+        # Select specific channels if save_channels is specified
+        # save_channels can be a list of indices like [0, 6] to save only channels 0 and 6
+        # Skip if channels were already filtered during inference (check if num_channels matches)
+        if save_channels is not None and sample.ndim >= 4:
+            channel_indices = list(save_channels)
+            num_channels = sample.shape[0]
+            # Only filter if not already filtered (channels > len(save_channels))
+            if num_channels > len(channel_indices):
+                try:
+                    sample = sample[channel_indices]
+                    print(f"  Selected channels {channel_indices} from {predictions[idx].shape[0]} channels")
+                except Exception as e:
+                    print(f"  WARNING: write_outputs - channel selection failed: {e}, keeping all channels")
+
         # Transpose if needed (output_transpose: list of axis permutation)
         if output_transpose and len(output_transpose) > 0:
             try:

connectomics/lightning/lit_model.py

@@ -286,7 +286,17 @@ def _extract_main_output(self, outputs: Union[torch.Tensor, Dict[str, torch.Tens
     def _sliding_window_predict(self, inputs: torch.Tensor) -> torch.Tensor:
         """Wrapper used by MONAI inferer to obtain primary model predictions."""
         outputs = self.forward(inputs)
-        return self._extract_main_output(outputs)
+        result = self._extract_main_output(outputs)
+
+        # Filter channels during inference to reduce memory during accumulation
+        # This is applied before the sliding window aggregates results
+        if hasattr(self.cfg, 'inference') and hasattr(self.cfg.inference, 'sliding_window'):
+            save_channels = getattr(self.cfg.inference.sliding_window, 'save_channels', None)
+            if save_channels is not None:
+                # Select only specified channels (e.g., [0, 6] for foreground + SDT)
+                result = result[:, save_channels, ...]
+
+        return result
 
     def _apply_tta_preprocessing(self, tensor: torch.Tensor) -> torch.Tensor:
         """