fix zebrafish_neurons.yaml: need multiple loss functions for each channel

Author: Donglai Wei

connectomics/config/hydra_config.py

@@ -911,8 +911,8 @@ class SavePredictionConfig:
     """
 
     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
+    intensity_scale: float = -1.0  # If < 0, keep raw predictions (no normalization/scaling). If > 0, normalize to [0,1] then scale.
+    intensity_dtype: str = "uint8"  # Save as uint8 for visualization (ignored if intensity_scale < 0)
 
 
 @dataclass
@@ -999,6 +999,8 @@ class PostprocessingConfig:
     Note: Intensity scaling and dtype conversion are handled by SavePredictionConfig.
     """
 
+    enabled: bool = False  # Enable postprocessing pipeline
+
     # Binary segmentation refinement (morphological ops, connected components)
     binary: Optional[Dict[str, Any]] = field(
         default_factory=dict
@@ -1059,7 +1061,7 @@ class TestDataConfig:
     test_image: Optional[str] = None
     test_label: Optional[str] = None
     test_mask: Optional[str] = None
-    test_resolution: Optional[List[int]] = None
+    test_resolution: Optional[List[float]] = None
     test_transpose: Optional[List[int]] = None
     output_path: Optional[str] = None
     cache_suffix: str = "_prediction.h5"

connectomics/decoding/segmentation.py

@@ -19,6 +19,7 @@
 from skimage.morphology import dilation, remove_small_objects
 from scipy.ndimage import zoom
 import mahotas
+from connectomics.data.process.target import seg_to_semantic_edt
 
 try:
     from numba import jit
@@ -175,10 +176,7 @@ def decode_instance_binary_contour_distance(
     elif mode == "watershed":
         # Watershed mode requires distance channel
         if distance is None:
-            raise ValueError(
-                "Watershed mode requires distance channel. "
-                "Please specify distance_channels in your decode configuration."
-            )
+            distance = seg_to_semantic_edt(foreground, mode="3d")
         # step 2: compute the instance seeds
         if precomputed_seed is not None:
             seed = precomputed_seed

connectomics/inference/io.py

@@ -22,25 +22,59 @@ def apply_save_prediction_transform(cfg: Config | DictConfig, data: np.ndarray)
 
     This is used when saving intermediate predictions (before decoding).
 
+    Default behavior (no config):
+    - Normalizes predictions to [0, 1] using min-max normalization
+    - Keeps dtype as float32
+
+    Config options:
+    - intensity_scale: If < 0, disables normalization (raw values)
+                       If > 0, normalize to [0, 1] then multiply by scale
+    - intensity_dtype: Target dtype for conversion (uint8, float32, etc.)
+
     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
+    # Default: keep raw predictions if no config
+    intensity_scale = -1.0  # Default: keep raw predictions
+
+    if hasattr(cfg, "inference") and hasattr(cfg.inference, "save_prediction"):
+        save_pred_cfg = cfg.inference.save_prediction
+        intensity_scale = getattr(save_pred_cfg, "intensity_scale", -1.0)
+
+    # Apply intensity scaling (if intensity_scale >= 0, normalize to [0, 1] then scale)
+    if intensity_scale >= 0:
+        # Convert to float32 for normalization
+        data = data.astype(np.float32)
+
+        # Min-max normalization to [0, 1]
+        data_min = data.min()
+        data_max = data.max()
 
-    save_pred_cfg = cfg.inference.save_prediction
+        if data_max > data_min:
+            data = (data - data_min) / (data_max - data_min)
+            print(f"  Normalized predictions to [0, 1] (min={data_min:.4f}, max={data_max:.4f})")
+        else:
+            print(f"  Warning: data_min == data_max ({data_min:.4f}), skipping normalization")
 
-    # 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 scaling factor
+        if intensity_scale != 1.0:
+            data = data * float(intensity_scale)
+            print(f"  Scaled predictions by {intensity_scale} -> range [{data.min():.4f}, {data.max():.4f}]")
+    else:
+        print(f"  Intensity scaling disabled (scale={intensity_scale} < 0), keeping raw predictions")
+        # Skip dtype conversion when intensity_scale < 0 to preserve raw predictions
+        return data
 
     # Apply dtype conversion
-    target_dtype_str = getattr(save_pred_cfg, "intensity_dtype", None)
+    target_dtype_str = None
+    if hasattr(cfg, "inference") and hasattr(cfg.inference, "save_prediction"):
+        save_pred_cfg = cfg.inference.save_prediction
+        target_dtype_str = getattr(save_pred_cfg, "intensity_dtype", None)
+
     if target_dtype_str is not None:
         dtype_map = {
             "uint8": np.uint8,
@@ -68,6 +102,7 @@ def apply_save_prediction_transform(cfg: Config | DictConfig, data: np.ndarray)
         if np.issubdtype(target_dtype, np.integer):
             info = np.iinfo(target_dtype)
             data = np.clip(data, info.min, info.max)
+            print(f"  Converting to {target_dtype_str} (clipped to [{info.min}, {info.max}])")
 
         data = data.astype(target_dtype)
 
@@ -89,6 +124,10 @@ def apply_postprocessing(cfg: Config | DictConfig, data: np.ndarray) -> np.ndarr
 
     postprocessing = cfg.inference.postprocessing
 
+    # Check if postprocessing is enabled
+    if not getattr(postprocessing, "enabled", False):
+        return data
+
     binary_config = getattr(postprocessing, "binary", None)
     if binary_config is not None and getattr(binary_config, "enabled", False):
         from connectomics.decoding.postprocess import apply_binary_postprocessing

scripts/main.py

@@ -144,41 +144,58 @@ def main():
     # Setup run directory (handles DDP coordination and config saving)
     # Determine output base directory from checkpoint for test/tune modes
     if args.mode in ["test", "tune", "tune-test"] and args.checkpoint:
-        # Extract base directory from checkpoint path
-        output_base = get_output_base_from_checkpoint(args.checkpoint)
-        output_base.mkdir(parents=True, exist_ok=True)
+        # Check if test.data.output_path is already set (skip checkpoint-based path extraction)
+        test_output_path_set = (
+            hasattr(cfg, "test")
+            and hasattr(cfg.test, "data")
+            and getattr(cfg.test.data, "output_path", None)
+        )
 
-        # Create mode-specific subdirectories
-        if args.mode in ["tune", "tune-test"]:
-            dirpath = str(output_base / "tuning")
-            results_path = str(output_base / "results")
-            # Override tune output directories in config
-            if cfg.tune is not None:
-                cfg.tune.output.output_dir = dirpath
-                cfg.tune.output.output_pred = results_path
-            # For tune-test, also set test output directory and cache suffix
-            if args.mode == "tune-test":
-                print(f"🔍 Setting test config for tune-test mode")
-                print(f"🔍 cfg.test is None: {cfg.test is None}")
-                if cfg.test is not None:
-                    print(f"🔍 cfg.test.data is None: {cfg.test.data is None}")
-                    if cfg.test.data is not None:
-                        cfg.test.data.output_path = results_path
-                        cfg.test.data.cache_suffix = cfg.tune.output.cache_suffix
-                        print(f"📋 Test output: {cfg.test.data.output_path}")
-                        print(f"📋 Test cache suffix: {cfg.test.data.cache_suffix}")
+        if args.mode == "test" and test_output_path_set:
+            # Use the config value directly, skip checkpoint-based directory creation
+            dirpath = str(cfg.test.data.output_path)
+            output_base = None  # Not needed when using config path
+        else:
+            # Extract base directory from checkpoint path
+            output_base = get_output_base_from_checkpoint(args.checkpoint)
+            output_base.mkdir(parents=True, exist_ok=True)
+
+            # Create mode-specific subdirectories
+            if args.mode in ["tune", "tune-test"]:
+                dirpath = str(output_base / "tuning")
+                results_path = str(output_base / "results")
+                # Override tune output directories in config
+                if cfg.tune is not None:
+                    cfg.tune.output.output_dir = dirpath
+                    cfg.tune.output.output_pred = results_path
+                # For tune-test, also set test output directory and cache suffix
+                if args.mode == "tune-test":
+                    print(f"🔍 Setting test config for tune-test mode")
+                    print(f"🔍 cfg.test is None: {cfg.test is None}")
+                    if cfg.test is not None:
+                        print(f"🔍 cfg.test.data is None: {cfg.test.data is None}")
+                        if cfg.test.data is not None:
+                            cfg.test.data.output_path = results_path
+                            cfg.test.data.cache_suffix = cfg.tune.output.cache_suffix
+                            print(f"📋 Test output: {cfg.test.data.output_path}")
+                            print(f"📋 Test cache suffix: {cfg.test.data.cache_suffix}")
+                        else:
+                            print(f"❌ cfg.test.data is None, cannot set cache_suffix!")
+                    else:
+                        print(f"❌ cfg.test is None, cannot set cache_suffix!")
+            else:  # test mode
+                dirpath = str(output_base / "results")
+                # Override test output directory in config only if not already set
+                if hasattr(cfg, "test") and hasattr(cfg.test, "data"):
+                    if not getattr(cfg.test.data, "output_path", None):
+                        cfg.test.data.output_path = dirpath
                     else:
-                        print(f"❌ cfg.test.data is None, cannot set cache_suffix!")
-                else:
-                    print(f"❌ cfg.test is None, cannot set cache_suffix!")
-        else:  # test mode
-            dirpath = str(output_base / "results")
-            # Override test output directory in config
-            if hasattr(cfg, "test") and hasattr(cfg.test, "data"):
-                cfg.test.data.output_path = dirpath
+                        # Use the config value, but ensure it's a string path
+                        dirpath = str(cfg.test.data.output_path)
 
         run_dir = setup_run_directory(args.mode, cfg, dirpath)
-        print(f"📂 Output base: {output_base}")
+        if output_base is not None:
+            print(f"📂 Output base: {output_base}")
     else:
         # Train mode or no checkpoint - use default config paths
         dirpath = cfg.monitor.checkpoint.dirpath