Unify profile engine to dict-valued profiles and fix legacy tutorial configs

All profile families (loss, decoding, activation) now use dict-valued
profiles with merge semantics instead of bare lists with replacement.
This makes the profile system consistent: every profile is a config
subtree merged into its parent section.

Migrate 12 tutorial configs from legacy flat keys (test_image,
test_label, test_path, test_resolution, reject_sampling) to the
modern nested DataInputConfig structure (test.data.test.*).

Add mask_transform field to DataConfig and wire it into
build_test_transforms and test_pipeline so vesicle_xm.yaml loads.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: Donglai Wei

connectomics/config/pipeline/profile_engine.py

@@ -194,6 +194,7 @@ def __init__(
         target_path: str,
         merge_into_existing: bool = True,
         cleanup_keys: Optional[Tuple[str, ...]] = None,
+        list_key: Optional[str] = None,
     ) -> None:
         if isinstance(selector_paths, str):
             selector_paths = [selector_paths]
@@ -202,6 +203,7 @@ def __init__(
         self.target_path = target_path
         self.merge_into_existing = merge_into_existing
         self.cleanup_keys = cleanup_keys if cleanup_keys is not None else (profiles_key,)
+        self.list_key = list_key
 
     def _apply_list_overrides(self, yaml_conf: DictConfig, selector_path: str) -> None:
         """Apply per-index overrides to a list-valued profile.
@@ -224,7 +226,8 @@ def _apply_list_overrides(self, yaml_conf: DictConfig, selector_path: str) -> No
         if overrides is None or not isinstance(overrides, (DictConfig, dict)):
             return
 
-        target_list = OmegaConf.select(yaml_conf, self.target_path)
+        list_path = f"{self.target_path}.{self.list_key}" if self.list_key else self.target_path
+        target_list = OmegaConf.select(yaml_conf, list_path)
         if not isinstance(target_list, (ListConfig, list)):
             return
 
@@ -343,10 +346,12 @@ def __init__(
         profiles_key: str,
         target_paths: List[str],
         profile_key: str = "profile",
+        list_key: Optional[str] = None,
     ) -> None:
         self.profiles_key = profiles_key
         self.target_paths = target_paths
         self.profile_key = profile_key
+        self.list_key = list_key
         self.cleanup_keys = (profiles_key,)
 
     def apply(self, yaml_conf: DictConfig) -> None:
@@ -376,34 +381,40 @@ def apply(self, yaml_conf: DictConfig) -> None:
                     )
 
                 profile_payload = profiles[profile_name]
+
+                # Extract nested list from dict-valued profiles.
+                profile_list = profile_payload
+                if self.list_key and isinstance(profile_payload, (DictConfig, dict)):
+                    profile_list = profile_payload.get(self.list_key, profile_payload)
+
                 overrides = {k: v for k, v in item.items() if k != self.profile_key}
 
                 if not overrides:
                     # Pure profile reference: expand entire profile list inline
-                    if not isinstance(profile_payload, (ListConfig, list)):
+                    if not isinstance(profile_list, (ListConfig, list)):
                         raise ValueError(
                             f"Profile '{profile_name}' in {self.profiles_key} must resolve to a list "
-                            f"for target '{target_path}', got {type(profile_payload)}"
+                            f"for target '{target_path}', got {type(profile_list)}"
                         )
-                    expanded_items.extend(list(profile_payload))
+                    expanded_items.extend(list(profile_list))
                 else:
                     # Profile + overrides: use first profile item as base, merge overrides on top
-                    if isinstance(profile_payload, (ListConfig, list)):
-                        if not profile_payload:
+                    if isinstance(profile_list, (ListConfig, list)):
+                        if not profile_list:
                             raise ValueError(
                                 f"Profile '{profile_name}' in {self.profiles_key} is empty."
                             )
                         base = OmegaConf.create(
-                            OmegaConf.to_container(profile_payload[0], resolve=False)
+                            OmegaConf.to_container(profile_list[0], resolve=False)
                         )
-                    elif isinstance(profile_payload, (DictConfig, dict)):
+                    elif isinstance(profile_list, (DictConfig, dict)):
                         base = OmegaConf.create(
-                            OmegaConf.to_container(profile_payload, resolve=False)
+                            OmegaConf.to_container(profile_list, resolve=False)
                         )
                     else:
                         raise ValueError(
                             f"Profile '{profile_name}' in {self.profiles_key} must be a list or dict, "
-                            f"got {type(profile_payload)}"
+                            f"got {type(profile_list)}"
                         )
                     merged = OmegaConf.merge(base, OmegaConf.create(overrides))
                     expanded_items.append(merged)
@@ -460,26 +471,28 @@ def _stage_path(stage: str, rel_path: str) -> str:
 
 
 # Each tuple:
-# (profiles_key, stages, selector_rel, target_rel, merge_into_existing)
+# (profiles_key, stages, selector_rel, target_rel, merge_into_existing, list_key)
 # Order matters: pipeline/system first, then arch, then the rest.
-_VALUE_PROFILE_FAMILIES: List[Tuple[str, Tuple[str, ...], str, str, bool]] = [
-    ("pipeline_profiles", (_STAGE_DEFAULT,), "pipeline_profile", "", True),
-    ("system_profiles", (_STAGE_DEFAULT, _STAGE_TRAIN, _STAGE_TEST, _STAGE_TUNE), "system.profile", "system", True),
-    ("arch_profiles", (_STAGE_DEFAULT, _STAGE_TRAIN, _STAGE_TEST, _STAGE_TUNE), "model.arch.profile", "model", True),
-    ("augmentation_profiles", (_STAGE_DEFAULT, _STAGE_TRAIN), "data.augmentation.profile", "data.augmentation", True),
-    ("dataloader_profiles", (_STAGE_DEFAULT, _STAGE_TRAIN, _STAGE_TEST, _STAGE_TUNE), "data.dataloader.profile", "data.dataloader", True),
-    ("optimizer_profiles", (_STAGE_DEFAULT, _STAGE_TRAIN), "optimization.profile", "optimization", True),
-    ("loss_profiles", (_STAGE_DEFAULT, _STAGE_TRAIN), "model.loss.profile", "model.loss.losses", False),
-    ("label_profiles", (_STAGE_DEFAULT, _STAGE_TRAIN), "data.label_transform.profile", "data.label_transform", False),
-    ("decoding_profiles", (_STAGE_DEFAULT, _STAGE_TEST, _STAGE_TUNE), "inference.decoding_profile", "inference.decoding", False),
-    ("activation_profiles", (_STAGE_DEFAULT, _STAGE_TEST, _STAGE_TUNE), "inference.test_time_augmentation.activation_profile", "inference.test_time_augmentation.channel_activations", False),
+# All profiles are dict-valued and use merge semantics. ``list_key`` identifies the
+# nested list within dict-valued profiles (used for positional overrides).
+_VALUE_PROFILE_FAMILIES: List[Tuple[str, Tuple[str, ...], str, str, bool, Optional[str]]] = [
+    ("pipeline_profiles", (_STAGE_DEFAULT,), "pipeline_profile", "", True, None),
+    ("system_profiles", (_STAGE_DEFAULT, _STAGE_TRAIN, _STAGE_TEST, _STAGE_TUNE), "system.profile", "system", True, None),
+    ("arch_profiles", (_STAGE_DEFAULT, _STAGE_TRAIN, _STAGE_TEST, _STAGE_TUNE), "model.arch.profile", "model", True, None),
+    ("augmentation_profiles", (_STAGE_DEFAULT, _STAGE_TRAIN), "data.augmentation.profile", "data.augmentation", True, None),
+    ("dataloader_profiles", (_STAGE_DEFAULT, _STAGE_TRAIN, _STAGE_TEST, _STAGE_TUNE), "data.dataloader.profile", "data.dataloader", True, None),
+    ("optimizer_profiles", (_STAGE_DEFAULT, _STAGE_TRAIN), "optimization.profile", "optimization", True, None),
+    ("loss_profiles", (_STAGE_DEFAULT, _STAGE_TRAIN), "model.loss.profile", "model.loss", True, "losses"),
+    ("label_profiles", (_STAGE_DEFAULT, _STAGE_TRAIN), "data.label_transform.profile", "data.label_transform", True, None),
+    ("decoding_profiles", (_STAGE_DEFAULT, _STAGE_TEST, _STAGE_TUNE), "inference.decoding_profile", "inference", True, "decoding"),
+    ("activation_profiles", (_STAGE_DEFAULT, _STAGE_TEST, _STAGE_TUNE), "inference.test_time_augmentation.activation_profile", "inference.test_time_augmentation", True, "channel_activations"),
 ]
 
 
-def _build_value_profile_specs() -> List[Tuple[List[str], str, str, bool, Tuple[str, ...]]]:
-    specs: List[Tuple[List[str], str, str, bool, Tuple[str, ...]]] = []
+def _build_value_profile_specs() -> List[Tuple[List[str], str, str, bool, Tuple[str, ...], Optional[str]]]:
+    specs: List[Tuple[List[str], str, str, bool, Tuple[str, ...], Optional[str]]] = []
 
-    for profiles_key, stages, selector_rel, target_rel, merge in _VALUE_PROFILE_FAMILIES:
+    for profiles_key, stages, selector_rel, target_rel, merge, list_key in _VALUE_PROFILE_FAMILIES:
         for stage in stages:
             selector_path = _stage_path(stage, selector_rel)
             target_path = _stage_path(stage, target_rel)
@@ -490,6 +503,7 @@ def _build_value_profile_specs() -> List[Tuple[List[str], str, str, bool, Tuple[
                     target_path,
                     merge,
                     (profiles_key, selector_path),
+                    list_key,
                 )
             )
 
@@ -498,10 +512,10 @@ def _build_value_profile_specs() -> List[Tuple[List[str], str, str, bool, Tuple[
 
 # Each tuple: (profiles_key, stages, target_rel)
 _REFERENCE_PROFILE_FAMILIES: List[Tuple[str, Tuple[str, ...], str]] = [
-    ("loss_profiles", (_STAGE_DEFAULT, _STAGE_TRAIN), "model.loss.losses"),
+    ("loss_profiles", (_STAGE_DEFAULT, _STAGE_TRAIN), "model.loss"),
     ("label_profiles", (_STAGE_DEFAULT, _STAGE_TRAIN), "data.label_transform"),
-    ("decoding_profiles", (_STAGE_DEFAULT, _STAGE_TEST, _STAGE_TUNE), "inference.decoding"),
-    ("activation_profiles", (_STAGE_DEFAULT, _STAGE_TEST, _STAGE_TUNE), "inference.test_time_augmentation.channel_activations"),
+    ("decoding_profiles", (_STAGE_DEFAULT, _STAGE_TEST, _STAGE_TUNE), "inference"),
+    ("activation_profiles", (_STAGE_DEFAULT, _STAGE_TEST, _STAGE_TUNE), "inference.test_time_augmentation"),
     ("augmentation_profiles", (_STAGE_DEFAULT, _STAGE_TRAIN), "data.augmentation"),
     ("dataloader_profiles", (_STAGE_DEFAULT, _STAGE_TRAIN, _STAGE_TEST, _STAGE_TUNE), "data.dataloader"),
     ("optimizer_profiles", (_STAGE_DEFAULT, _STAGE_TRAIN), "optimization"),
@@ -517,24 +531,25 @@ def _build_reference_profile_specs() -> List[Tuple[str, List[str]]]:
     return specs
 
 
-# Each tuple: (profiles_key, stages, target_rel)
-_LIST_REFERENCE_FAMILIES: List[Tuple[str, Tuple[str, ...], str]] = [
-    ("decoding_profiles", (_STAGE_DEFAULT, _STAGE_TUNE, _STAGE_TEST), "inference.decoding"),
+# Each tuple: (profiles_key, stages, target_rel, list_key)
+# ``list_key`` identifies the nested list within dict-valued profiles.
+_LIST_REFERENCE_FAMILIES: List[Tuple[str, Tuple[str, ...], str, str]] = [
+    ("decoding_profiles", (_STAGE_DEFAULT, _STAGE_TUNE, _STAGE_TEST), "inference.decoding", "decoding"),
 ]
 
 
-def _build_list_reference_specs() -> List[Tuple[str, List[str]]]:
-    specs: List[Tuple[str, List[str]]] = []
-    for profiles_key, stages, target_rel in _LIST_REFERENCE_FAMILIES:
+def _build_list_reference_specs() -> List[Tuple[str, List[str], str]]:
+    specs: List[Tuple[str, List[str], str]] = []
+    for profiles_key, stages, target_rel, list_key in _LIST_REFERENCE_FAMILIES:
         target_paths = [_stage_path(stage, target_rel) for stage in stages]
-        specs.append((profiles_key, target_paths))
+        specs.append((profiles_key, target_paths, list_key))
     return specs
 
 
 def _build_allowed_selector_paths() -> set[str]:
     paths: set[str] = set()
 
-    for _, stages, selector_rel, _, _ in _VALUE_PROFILE_FAMILIES:
+    for _, stages, selector_rel, _, _, _ in _VALUE_PROFILE_FAMILIES:
         for stage in stages:
             paths.add(_normalize_selector_path(_stage_path(stage, selector_rel)))
 
@@ -574,13 +589,14 @@ def _build_profile_engine() -> YamlProfileEngine:
     appliers: List[YamlProfileApplier] = []
 
     # 1) Value profile appliers (order defined by _VALUE_PROFILE_FAMILIES table)
-    for selector_paths, profiles_key, target_path, merge, cleanup_keys in _VALUE_PROFILE_SPECS:
+    for selector_paths, profiles_key, target_path, merge, cleanup_keys, list_key in _VALUE_PROFILE_SPECS:
         appliers.append(ValueProfileApplier(
             selector_paths=selector_paths,
             profiles_key=profiles_key,
             target_path=target_path,
             merge_into_existing=merge,
             cleanup_keys=cleanup_keys,
+            list_key=list_key,
         ))
 
     # 2) Reference profile appliers
@@ -591,10 +607,11 @@ def _build_profile_engine() -> YamlProfileEngine:
         ))
 
     # 3) List profile reference appliers
-    for profiles_key, target_paths in _LIST_REFERENCE_SPECS:
+    for profiles_key, target_paths, list_key in _LIST_REFERENCE_SPECS:
         appliers.append(ListProfileReferenceApplier(
             profiles_key=profiles_key,
             target_paths=target_paths,
+            list_key=list_key,
         ))
 
     return YamlProfileEngine(appliers)

connectomics/config/schema/data.py

@@ -469,6 +469,9 @@ class DataConfig:
         default_factory=NNUNetPreprocessingConfig
     )
 
+    # Mask-specific transformation (overrides data_transform for masks when present)
+    mask_transform: Optional[DataTransformConfig] = None
+
     # Multi-channel label transformation (for affinity maps, distance transforms, etc.)
     label_transform: LabelTransformConfig = field(default_factory=LabelTransformConfig)
 

connectomics/data/augment/build.py

@@ -319,9 +319,11 @@ def _build_eval_transforms_impl(cfg: Config, mode: str = "val", keys: list[str]
     # Apply resize if configured (before cropping)
     image_resize_factors = getattr(data_cfg.image_transform, "resize", None)
 
+    # Prefer mask_transform over data_transform for mask-specific settings.
+    mask_cfg = getattr(data_cfg, "mask_transform", None) or data_cfg.data_transform
     mask_resize_factors = None
-    if mode in {"test", "tune"} and data_cfg.data_transform.resize is not None:
-        mask_resize_factors = data_cfg.data_transform.resize
+    if mode in {"test", "tune"} and mask_cfg.resize is not None:
+        mask_resize_factors = mask_cfg.resize
 
     if image_resize_factors is not None and image_resize_factors:
         transforms.append(
@@ -366,8 +368,8 @@ def _build_eval_transforms_impl(cfg: Config, mode: str = "val", keys: list[str]
     mask_binarize = False
     mask_threshold = 0.0
     if mode in {"test", "tune"}:
-        mask_binarize = bool(getattr(data_cfg.data_transform, "binarize", False))
-        mask_threshold = float(getattr(data_cfg.data_transform, "threshold", 0.0))
+        mask_binarize = bool(getattr(mask_cfg, "binarize", False))
+        mask_threshold = float(getattr(mask_cfg, "threshold", 0.0))
 
     if "mask" in keys and mask_binarize:
         transforms.append(

connectomics/training/lightning/test_pipeline.py

@@ -831,7 +831,9 @@ def run_test_step(module, batch: Dict[str, torch.Tensor], batch_idx: int) -> STE
     logger.info("Starting sliding-window inference...")
 
     mask_align_to_image = False
-    mask_transform_cfg = getattr(module.cfg.data, "data_transform", None)
+    mask_transform_cfg = getattr(module.cfg.data, "mask_transform", None) or getattr(
+        module.cfg.data, "data_transform", None
+    )
     if mask_transform_cfg is not None:
         mask_align_to_image = bool(getattr(mask_transform_cfg, "align_to_image", False))
 

tutorials/bases/activation_profiles.yaml

@@ -1,11 +1,14 @@
 activation_profiles:
   act_binary:
-    - {channels: ":", activation: sigmoid}
+    channel_activations:
+      - {channels: ":", activation: sigmoid}
 
   act_bcd:
-    - {channels: "0:2", activation: sigmoid}
-    - {channels: "2:3", activation: tanh}
+    channel_activations:
+      - {channels: "0:2", activation: sigmoid}
+      - {channels: "2:3", activation: tanh}
 
   act_bd:
-    - {channels: "0:-1", activation: sigmoid}
-    - {channels: "-1:", activation: tanh}
+    channel_activations:
+      - {channels: "0:-1", activation: sigmoid}
+      - {channels: "-1:", activation: tanh}

tutorials/bases/decoding_profiles.yaml

@@ -1,21 +1,23 @@
 decoding_profiles:
   decoding_bcd:
-    - name: decode_instance_binary_contour_distance
-      kwargs:
-        binary_threshold: [0.9, 0.85]
-        contour_threshold: [0.8, 1.1]
-        distance_threshold: [0.5, -0.5]
-        min_instance_size: 30
-        min_seed_size: 5
+    decoding:
+      - name: decode_instance_binary_contour_distance
+        kwargs:
+          binary_threshold: [0.9, 0.85]
+          contour_threshold: [0.8, 1.1]
+          distance_threshold: [0.5, -0.5]
+          min_instance_size: 30
+          min_seed_size: 5
   decoding_abiss:
-    - name: decode_abiss
-      kwargs:
-        command: "{python_exe} scripts/run_abiss_single.py --input {input_h5} --output {output_h5}"
-        # Affinity thresholds accept absolute values (e.g. 0.80) or
-        # percentiles (e.g. "80%") resolved from the data.
-        cli_args:
-          ws_high_threshold: 80%
-          ws_low_threshold: 1%
-          ws_size_threshold: 800
-          ws_dust_threshold: 600
-          ws_merge_threshold: 20%
+    decoding:
+      - name: decode_abiss
+        kwargs:
+          command: "{python_exe} scripts/run_abiss_single.py --input {input_h5} --output {output_h5}"
+          # Affinity thresholds accept absolute values (e.g. 0.80) or
+          # percentiles (e.g. "80%") resolved from the data.
+          cli_args:
+            ws_high_threshold: 80%
+            ws_low_threshold: 1%
+            ws_size_threshold: 800
+            ws_dust_threshold: 600
+            ws_merge_threshold: 20%