fix multi-task loss and label mapping

Author: Donglai Wei

connectomics/config/hydra_config.py

@@ -313,6 +313,8 @@ class ImageTransformConfig:
     clip_percentile_high: float = (
         1.0  # Upper percentile for clipping (1.0 = no clip, 0.95 = 95th percentile)
     )
+    pad_size: Optional[List[int]] = None  # Reflection padding for context [D, H, W] or [H, W]
+    pad_mode: str = "reflect"  # Padding mode: 'reflect', 'replicate', 'constant'
 
 
 @dataclass

connectomics/data/augment/build.py

@@ -25,6 +25,8 @@
     CenterSpatialCropd,
     SpatialPadd,
     Resized,
+    LoadImaged,  # For filename-based datasets (PNG, JPG, etc.)
+    EnsureChannelFirstd,  # Ensure channel-first format for 2D/3D images
 )
 
 # Import custom loader for HDF5/TIFF volumes
@@ -71,11 +73,20 @@ def build_train_transforms(
 
     # Load images first (unless using pre-cached dataset)
     if not skip_loading:
-        # Get transpose axes for training data
-        train_transpose = cfg.data.train_transpose if cfg.data.train_transpose else []
-        transforms.append(
-            LoadVolumed(keys=keys, transpose_axes=train_transpose if train_transpose else None)
-        )
+        # Use appropriate loader based on dataset type
+        dataset_type = getattr(cfg.data, "dataset_type", "volume")  # Default to volume for backward compatibility
+
+        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
+            train_transpose = cfg.data.train_transpose if cfg.data.train_transpose else []
+            transforms.append(
+                LoadVolumed(keys=keys, transpose_axes=train_transpose if train_transpose else None)
+            )
 
     # Apply volumetric split if enabled
     if cfg.data.split_enabled:
@@ -212,12 +223,20 @@ def build_val_transforms(cfg: Config, keys: list[str] = None) -> Compose:
 
     transforms = []
 
-    # Load images first
-    # Get transpose axes for validation data
-    val_transpose = cfg.data.val_transpose if cfg.data.val_transpose else []
-    transforms.append(
-        LoadVolumed(keys=keys, transpose_axes=val_transpose if val_transpose else None)
-    )
+    # Load images first - use appropriate loader based on dataset type
+    dataset_type = getattr(cfg.data, "dataset_type", "volume")  # Default to volume for backward compatibility
+
+    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
+        val_transpose = cfg.data.val_transpose if cfg.data.val_transpose else []
+        transforms.append(
+            LoadVolumed(keys=keys, transpose_axes=val_transpose if val_transpose else None)
+        )
 
     # Apply volumetric split if enabled
     if cfg.data.split_enabled:
@@ -342,20 +361,29 @@ def build_test_transforms(cfg: Config, keys: list[str] = None) -> Compose:
 
     transforms = []
 
-    # Load images first
-    # Get transpose axes for test data (check both data.test_transpose and inference.data.test_transpose)
-    test_transpose = []
-    if cfg.data.test_transpose:
-        test_transpose = cfg.data.test_transpose
-    if (
-        hasattr(cfg, "inference")
-        and hasattr(cfg.inference, "data")
-        and hasattr(cfg.inference.data, "test_transpose")
-        and cfg.inference.data.test_transpose
-    ):
-        test_transpose = cfg.inference.data.test_transpose  # inference takes precedence
-    transforms.append(
-        LoadVolumed(keys=keys, transpose_axes=test_transpose if test_transpose else None)
+    # Load images first - use appropriate loader based on dataset type
+    dataset_type = getattr(cfg.data, "dataset_type", "volume")  # Default to volume for backward compatibility
+
+    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 for test data (check both data.test_transpose and inference.data.test_transpose)
+        test_transpose = []
+        if cfg.data.test_transpose:
+            test_transpose = cfg.data.test_transpose
+        if (
+            hasattr(cfg, "inference")
+            and hasattr(cfg.inference, "data")
+            and hasattr(cfg.inference.data, "test_transpose")
+            and cfg.inference.data.test_transpose
+        ):
+            test_transpose = cfg.inference.data.test_transpose  # inference takes precedence
+        transforms.append(
+            LoadVolumed(keys=keys, transpose_axes=test_transpose if test_transpose else None)
     )
 
     # Apply volumetric split if enabled (though typically not used for test)

connectomics/data/dataset/dataset_volume_cached.py

@@ -79,6 +79,8 @@ def __init__(
         iter_num: int = 500,
         transforms: Optional[Compose] = None,
         mode: str = "train",
+        pad_size: Optional[Tuple[int, ...]] = None,
+        pad_mode: str = "reflect",
     ):
         self.image_paths = image_paths
         self.label_paths = label_paths if label_paths else [None] * len(image_paths)
@@ -97,6 +99,8 @@ def __init__(
         self.iter_num = iter_num if iter_num > 0 else len(image_paths)
         self.transforms = transforms
         self.mode = mode
+        self.pad_size = pad_size
+        self.pad_mode = pad_mode
 
         # Load all volumes into memory
         print(f"  Loading {len(image_paths)} volumes into memory...")
@@ -116,6 +120,11 @@ def __init__(
                 img = img[None, ...]  # Add channel for 2D
             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)
+
             self.cached_images.append(img)
 
             # Load label if available
@@ -126,6 +135,11 @@ def __init__(
                     lbl = lbl[None, ...]  # Add channel for 2D
                 elif lbl.ndim == 3:
                     lbl = lbl[None, ...]  # Add channel for 3D
+
+                # Apply padding if specified (same padding as image)
+                if self.pad_size is not None:
+                    lbl = self._apply_padding(lbl, mode='constant', constant_values=0)  # Use constant 0 for labels
+
                 self.cached_labels.append(lbl)
             else:
                 self.cached_labels.append(None)
@@ -135,6 +149,11 @@ def __init__(
                 mask = read_volume(mask_path)
                 if mask.ndim == 3:
                     mask = mask[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)
+
                 self.cached_masks.append(mask)
             else:
                 self.cached_masks.append(None)
@@ -148,6 +167,40 @@ def __init__(
         ndim = len(self.patch_size)
         self.volume_sizes = [img.shape[-ndim:] for img in self.cached_images]  # (Z, Y, X) or (Y, X)
 
+    def _apply_padding(
+        self, volume: np.ndarray, mode: Optional[str] = None, constant_values: float = 0
+    ) -> np.ndarray:
+        """
+        Apply padding to a volume using np.pad.
+
+        Args:
+            volume: Input volume with channel dimension (C, D, H, W) or (C, H, W)
+            mode: Padding mode ('reflect', 'constant', etc.). If None, uses self.pad_mode
+            constant_values: Value for constant padding
+
+        Returns:
+            Padded volume
+        """
+        if self.pad_size is None:
+            return volume
+
+        mode = mode if mode is not None else self.pad_mode
+
+        # Build padding tuple for np.pad: ((before, after), ...)
+        # For channel dimension: no padding (0, 0)
+        # For spatial dimensions: pad according to pad_size
+        pad_width = [(0, 0)]  # No padding on channel dimension
+        for p in self.pad_size:
+            pad_width.append((p, p))
+
+        # Apply padding using np.pad
+        if mode == 'constant':
+            padded = np.pad(volume, pad_width, mode=mode, constant_values=constant_values)
+        else:
+            padded = np.pad(volume, pad_width, mode=mode)
+
+        return padded
+
     def __len__(self) -> int:
         return self.iter_num
 

connectomics/data/process/build.py

@@ -123,9 +123,16 @@ def create_label_transform_pipeline(cfg: Any = None, **kwargs: Any) -> Compose:
     cfg = _coerce_config(cfg, kwargs)
 
     # Keys configuration
-    keys_attr = getattr(cfg, 'keys', None)
-    if keys_attr is None:
-        keys_option = [getattr(cfg, 'input_key', 'label')]
+    # Note: Must check if 'keys' exists in config to avoid getting dict.keys() method
+    if hasattr(cfg, '__dict__') and 'keys' in cfg.__dict__:
+        keys_attr = cfg.keys
+    elif hasattr(cfg, '__contains__') and 'keys' in cfg:
+        keys_attr = cfg['keys'] if isinstance(cfg, dict) else getattr(cfg, 'keys')
+    else:
+        keys_attr = None
+
+    if keys_attr is None or callable(keys_attr):  # Protect against dict.keys() method
+        keys_option = [getattr(cfg, 'input_key', None) or cfg.get('input_key', 'label') if isinstance(cfg, dict) else 'label']
     elif isinstance(keys_attr, str):
         keys_option = [keys_attr]
     else:

connectomics/data/process/distance.py

@@ -256,20 +256,13 @@ def skeleton_aware_distance_transform(
         label = np.pad(label, pad_size, mode="constant", constant_values=0)
 
     label_shape = label.shape
-    all_bg_sample = False
 
     skeleton = np.zeros(label_shape, dtype=np.uint8)
     distance = np.zeros(label_shape, dtype=np.float32)
 
     indices = np.unique(label)
-    if indices[0] == 0:
-        if len(indices) > 1:  # exclude background
-            indices = indices[1:]
-        else:  # all-background sample
-            all_bg_sample = True
-
-    if not all_bg_sample:
-        for idx in indices:
+    if len(indices) > 1:
+        for idx in indices[indices > 0]:
             temp2 = remove_small_holes(label == idx, 16, connectivity=1)
             binary = temp2.copy()
 

connectomics/data/process/monai_transforms.py

@@ -718,23 +718,56 @@ def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
             label = d[key]
             label_np, had_batch_dim = self._prepare_label(label)
 
-            # Remove channel dimension if it's 1 (target functions expect [D, H, W] not [1, D, H, W])
-            if label_np.ndim == 4 and label_np.shape[0] == 1:
-                label_np = label_np[0]
+            # Determine if input is 2D or 3D based on original dimensions
+            # After EnsureChannelFirstd: 2D images are [1, H, W], 3D volumes are [1, D, H, W]
+            is_2d_input = label_np.ndim == 3 and label_np.shape[0] == 1
+            is_3d_input = label_np.ndim == 4 and label_np.shape[0] == 1
+
+            # Remove channel dimension (target functions don't expect it)
+            if is_3d_input:
+                label_np = label_np[0]  # [1, D, H, W] -> [D, H, W]
+            elif is_2d_input:
+                # For 2D, keep as [1, H, W] since some functions (boundary, edt) expect 3D input
+                # even in 2D mode (they treat first dim as Z=1)
+                pass  # Keep [1, H, W]
 
             outputs: List[np.ndarray] = []
             for spec in self.task_specs:
-                result = spec["fn"](label_np, **spec["kwargs"])
+                try:
+                    result = spec["fn"](label_np, **spec["kwargs"])
+                except Exception as e:
+                    raise RuntimeError(
+                        f"Task '{spec['name']}' failed with error: {e}\n"
+                        f"Label shape: {label_np.shape}, dtype: {label_np.dtype}\n"
+                        f"Task kwargs: {spec['kwargs']}"
+                    ) from e
                 if result is None:
-                    raise RuntimeError(f"Task '{spec['name']}' returned None.")
+                    raise RuntimeError(
+                        f"Task '{spec['name']}' returned None.\n"
+                        f"Label shape: {label_np.shape}, dtype: {label_np.dtype}\n"
+                        f"Task kwargs: {spec['kwargs']}"
+                    )
                 result_arr = np.asarray(
                     result, dtype=np.float32
                 )  # Convert to float32 (handles bool->float)
 
-                # Ensure each output has a channel dimension [C, D, H, W]
-                # If output is [D, H, W], expand to [1, D, H, W]
-                if result_arr.ndim == 3:
-                    result_arr = result_arr[np.newaxis, ...]  # Add channel dimension
+                # Normalize output dimensions:
+                # For 2D images (input [1, H, W]): functions return [H, W] or [1, H, W]
+                # For 3D volumes (input [D, H, W]): functions return [D, H, W]
+                # Goal: Add channel dimension to get [1, H, W] for 2D or [1, D, H, W] for 3D
+
+                if is_2d_input:
+                    # 2D case: some functions return [H, W], others return [1, H, W]
+                    if result_arr.ndim == 3 and result_arr.shape[0] == 1:
+                        # Function returned [1, H, W], squeeze Z dimension
+                        result_arr = result_arr[0]  # [1, H, W] -> [H, W]
+                    # Now result_arr is [H, W], add channel dimension
+                    if result_arr.ndim == 2:
+                        result_arr = result_arr[np.newaxis, ...]  # [H, W] -> [1, H, W]
+                elif is_3d_input:
+                    # 3D case: functions return [D, H, W], add channel dimension
+                    if result_arr.ndim == 3:
+                        result_arr = result_arr[np.newaxis, ...]  # [D, H, W] -> [1, D, H, W]
 
                 outputs.append(result_arr)
 

connectomics/lightning/callbacks.py

@@ -129,7 +129,12 @@ def on_train_epoch_end(self, trainer, pl_module):
 
             print(f"✓ Saved visualization for epoch {trainer.current_epoch}")
         except Exception as e:
+            import traceback
             print(f"Epoch-end visualization failed: {e}")
+            print(f"Error type: {type(e).__name__}")
+            if hasattr(e, '__traceback__'):
+                print("Traceback:")
+                traceback.print_exception(type(e), e, e.__traceback__)
 
     def on_validation_epoch_end(self, trainer, pl_module):
         """Visualize at end of validation epoch based on log_every_n_epochs."""
@@ -172,7 +177,12 @@ def on_validation_epoch_end(self, trainer, pl_module):
                     prefix='val'  # Single tab name (no epoch prefix)
                 )
         except Exception as e:
+            import traceback
             print(f"Validation epoch-end visualization failed: {e}")
+            print(f"Error type: {type(e).__name__}")
+            if hasattr(e, '__traceback__'):
+                print("Traceback:")
+                traceback.print_exception(type(e), e, e.__traceback__)
 
 
 class NaNDetectionCallback(Callback):