add optuna parameter search

Author: Donglai Wei

OPTUNA_INTEGRATION.md

@@ -99,9 +99,9 @@ The Lucchi++ dataset contains mitochondria segmentation data from EM images.
 ```bash
 # Download from HuggingFace (recommended)
 mkdir -p datasets/
-wget https://huggingface.co/datasets/pytc/tutorial/resolve/main/Lucchi%2B%2B.zip
-unzip Lucchi++.zip -d datasets/
-rm Lucchi++.zip
+wget https://huggingface.co/datasets/pytc/tutorial/resolve/main/lucchi%2B%2B.zip
+unzip lucchi++.zip -d datasets/
+rm lucchi++.zip
 ```
 
 **Size:** ~100 MB

QUICKSTART.md

@@ -321,8 +321,8 @@ data:
 **Solution - Manual download:**
 ```bash
 # Download from HuggingFace
-wget https://huggingface.co/datasets/pytc/tutorial/resolve/main/Lucchi%2B%2B.zip
-unzip Lucchi++.zip -d datasets/
+wget https://huggingface.co/datasets/pytc/tutorial/resolve/main/lucchi%2B%2B.zip
+unzip lucchi++.zip -d datasets/
 
 # Or use git-lfs
 git lfs install

TROUBLESHOOTING.md

@@ -140,20 +140,20 @@ def decode_binary_thresholding(
 
     # Apply morphological opening (erosion + dilation) - removes small objects
     if opening_iterations > 0:
-        binary_mask = ndimage.binary_opening(
-            binary_mask, iterations=opening_iterations
-        ).astype(np.uint8)
+        binary_mask = ndimage.binary_opening(binary_mask, iterations=opening_iterations).astype(
+            np.uint8
+        )
 
     # Apply morphological closing (dilation + erosion) - fills small holes
     if closing_iterations > 0:
-        binary_mask = ndimage.binary_closing(
-            binary_mask, iterations=closing_iterations
-        ).astype(np.uint8)
+        binary_mask = ndimage.binary_closing(binary_mask, iterations=closing_iterations).astype(
+            np.uint8
+        )
 
     # Apply connected components filtering
-    if connected_components and connected_components.get('enabled', False):
-        remove_small = connected_components.get('remove_small', 0)
-        connectivity = connected_components.get('connectivity', 26)
+    if connected_components and connected_components.get("enabled", False):
+        remove_small = connected_components.get("remove_small", 0)
+        connectivity = connected_components.get("connectivity", 26)
 
         if remove_small > 0:
             # Label connected components
@@ -359,7 +359,7 @@ def decode_binary_contour_distance_watershed(
     min_seed_size: int = 32,
     return_seed: bool = False,
     precomputed_seed: Optional[np.ndarray] = None,
-    prediction_scale: int = 255,
+    prediction_scale: int = 1,
     binary_channels: Optional[List[int]] = None,
     contour_channels: Optional[List[int]] = None,
     distance_channels: Optional[List[int]] = None,
@@ -456,15 +456,19 @@ def decode_binary_contour_distance_watershed(
     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]}"
+            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]}"
+            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)
@@ -476,10 +480,19 @@ def decode_binary_contour_distance_watershed(
     # Convert thresholds based on prediction scale
     if prediction_scale == 255:
         distance = (distance / prediction_scale) * 2.0 - 1.0
-        binary_threshold = (binary_threshold[0] * prediction_scale, binary_threshold[1] * 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)
+            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
@@ -492,10 +505,7 @@ def decode_binary_contour_distance_watershed(
             )
         else:
             # No contour constraint - only binary and distance
-            seed_map = (
-                (binary > binary_threshold[0])
-                * (distance > distance_threshold[0])
-            )
+            seed_map = (binary > binary_threshold[0]) * (distance > distance_threshold[0])
         seed = cc3d.connected_components(seed_map)
         seed = remove_small_objects(seed, min_seed_size)
 
@@ -507,10 +517,7 @@ def decode_binary_contour_distance_watershed(
         )
     else:
         # No contour constraint - only binary and distance
-        foreground = (
-            (binary > binary_threshold[1])
-            * (distance > distance_threshold[1])
-        )
+        foreground = (binary > binary_threshold[1]) * (distance > distance_threshold[1])
 
     segmentation = mahotas.cwatershed(-distance.astype(np.float64), seed)
     segmentation[~foreground] = (

connectomics/decoding/segmentation.py

@@ -64,51 +64,32 @@ def create_loss(
     loss_registry = {
         # MONAI Dice variants
         'DiceLoss': DiceLoss,
-        'Dice': DiceLoss,  # Alias
         'DiceCELoss': DiceCELoss,
-        'DiceCE': DiceCELoss,  # Alias
         'DiceFocalLoss': DiceFocalLoss,
-        'DiceFocal': DiceFocalLoss,  # Alias
         'GeneralizedDiceLoss': GeneralizedDiceLoss,
-        'GDiceLoss': GeneralizedDiceLoss,  # Alias
 
         # MONAI other losses
         'FocalLoss': FocalLoss,
-        'Focal': FocalLoss,  # Alias
         'TverskyLoss': TverskyLoss,
-        'Tversky': TverskyLoss,  # Alias
 
         # PyTorch standard losses (for convenience)
         'BCEWithLogitsLoss': nn.BCEWithLogitsLoss,
-        'BCE': nn.BCEWithLogitsLoss,  # Alias
         'CrossEntropyLoss': CrossEntropyLossWrapper,  # Use wrapper for shape handling
-        'CE': CrossEntropyLossWrapper,  # Alias
         'MSELoss': nn.MSELoss,
-        'MSE': nn.MSELoss,  # Alias
         'L1Loss': nn.L1Loss,
-        'L1': nn.L1Loss,  # Alias
 
         # Custom connectomics losses
         'WeightedBCEWithLogitsLoss': WeightedBCEWithLogitsLoss,
-        'WeightedBCE': WeightedBCEWithLogitsLoss,  # Alias
         'WeightedMSELoss': WeightedMSELoss,
-        'WeightedMSE': WeightedMSELoss,  # Alias
         'WeightedMAELoss': WeightedMAELoss,
-        'WeightedMAE': WeightedMAELoss,  # Alias
         'GANLoss': GANLoss,
-        'GAN': GANLoss,  # Alias
 
         # Regularization losses
         'BinaryRegularization': BinaryRegularization,
-        'BinaryReg': BinaryRegularization,  # Alias
         'ForegroundDistanceConsistency': ForegroundDistanceConsistency,
-        'FgDTConsistency': ForegroundDistanceConsistency,  # Alias
         'ContourDistanceConsistency': ContourDistanceConsistency,
-        'ContourDTConsistency': ContourDistanceConsistency,  # Alias
         'ForegroundContourConsistency': ForegroundContourConsistency,
-        'FgContourConsistency': ForegroundContourConsistency,  # Alias
         'NonOverlapRegularization': NonOverlapRegularization,
-        'NonoverlapReg': NonOverlapRegularization,  # Alias
     }
 
     if loss_name not in loss_registry:

connectomics/models/loss/build.py

@@ -15,26 +15,26 @@
 DATASETS = {
     "lucchi": {
         "name": "Lucchi++ Mitochondria Segmentation",
-        "url": "https://huggingface.co/datasets/pytc/tutorial/resolve/main/Lucchi%2B%2B.zip",
+        "url": "https://huggingface.co/datasets/pytc/tutorial/resolve/main/lucchi%2B%2B.zip",
         "size_mb": 100,
         "description": "EM images with mitochondria annotations from Lucchi et al.",
         "files": [
-            "datasets/Lucchi++/train_image.h5",
-            "datasets/Lucchi++/train_label.h5",
-            "datasets/Lucchi++/test_image.h5",
-            "datasets/Lucchi++/test_label.h5",
+            "datasets/lucchi++/train_image.h5",
+            "datasets/lucchi++/train_label.h5",
+            "datasets/lucchi++/test_image.h5",
+            "datasets/lucchi++/test_label.h5",
         ],
     },
     "lucchi++": {  # Alias
         "name": "Lucchi++ Mitochondria Segmentation",
-        "url": "https://huggingface.co/datasets/pytc/tutorial/resolve/main/Lucchi%2B%2B.zip",
+        "url": "https://huggingface.co/datasets/pytc/tutorial/resolve/main/lucchi%2B%2B.zip",
         "size_mb": 100,
         "description": "EM images with mitochondria annotations from Lucchi et al.",
         "files": [
-            "datasets/Lucchi++/train_image.h5",
-            "datasets/Lucchi++/train_label.h5",
-            "datasets/Lucchi++/test_image.h5",
-            "datasets/Lucchi++/test_label.h5",
+            "datasets/lucchi++/train_image.h5",
+            "datasets/lucchi++/train_label.h5",
+            "datasets/lucchi++/test_image.h5",
+            "datasets/lucchi++/test_label.h5",
         ],
     },
 }

connectomics/utils/download.py

@@ -56,6 +56,26 @@ test dataset arch_or_ckpt ckpt_or_args='' *ARGS='':
         python scripts/main.py --config tutorials/{{dataset}}.yaml model.architecture={{arch_or_ckpt}} --mode test --checkpoint {{ckpt_or_args}} {{ARGS}}
     fi
 
+# Tune decoding parameters on validation set (e.g., just tune hydra-lv ckpt.pt)
+tune dataset ckpt *ARGS='':
+    python scripts/main.py --config tutorials/{{dataset}}.yaml --mode tune --checkpoint {{ckpt}} {{ARGS}}
+
+# Tune parameters then test (recommended for optimal results)
+tune-test dataset ckpt *ARGS='':
+    python scripts/main.py --config tutorials/{{dataset}}.yaml --mode tune+test --checkpoint {{ckpt}} {{ARGS}}
+
+# Quick tuning with 20 trials (for testing)
+tune-quick dataset ckpt *ARGS='':
+    python scripts/main.py --config tutorials/{{dataset}}.yaml --mode tune --checkpoint {{ckpt}} --tune-trials 20 {{ARGS}}
+
+# Test with specific parameter file
+test-with-params dataset ckpt params *ARGS='':
+    python scripts/main.py --config tutorials/{{dataset}}.yaml --mode test --checkpoint {{ckpt}} --params {{params}} {{ARGS}}
+
+# Inference (alias for test, clearer naming)
+infer dataset ckpt *ARGS='':
+    python scripts/main.py --config tutorials/{{dataset}}.yaml --mode infer --checkpoint {{ckpt}} {{ARGS}}
+
 # ============================================================================
 # Monitoring Commands
 # ============================================================================