feat: [WIP/Proposal] Support copick global labels as the persisted label space

octopi/datasets/augment.py

@@ -1,9 +1,10 @@
 from octopi.datasets.loader import LoadCopickd
+import numpy as np
 from monai.transforms import (
-    Compose, 
-    RandFlipd, 
-    Orientationd, 
-    RandRotate90d, 
+    Compose,
+    RandFlipd,
+    Orientationd,
+    RandRotate90d,
     NormalizeIntensityd,
     EnsureChannelFirstd,
     RandCropByPosNegLabeld,
@@ -13,21 +14,37 @@
     RandShiftIntensityd,
     RandAdjustContrastd,
     RandGaussianNoised,
-    ScaleIntensityRanged,  
+    ScaleIntensityRanged,
     RandomOrder,
     RandAffined,
+    MapLabelValued,
 )
 
-def get_transforms():
+def get_transforms(orig_labels=None, target_labels=None, label_dtype=np.uint8):
     """
     Returns non-random transforms.
+
+    When ``orig_labels``/``target_labels`` are provided, a MapLabelValued is inserted
+    right after loading to compact the persisted copick *global* labels into the model's
+    dense ``[0..K]`` space (the loss/metric/one-hot ops require dense, contiguous labels).
+    The map is identity-collapsing for legacy sequential data, so it is a no-op there.
     """
-    return Compose([
-        LoadCopickd(),
+    transforms = [LoadCopickd()]
+    if orig_labels is not None:
+        transforms.append(
+            MapLabelValued(
+                keys=["label"],
+                orig_labels=orig_labels,
+                target_labels=target_labels,
+                dtype=label_dtype,
+            )
+        )
+    transforms += [
         EnsureChannelFirstd(keys=["image", "label"], channel_dim="no_channel"),
         NormalizeIntensityd(keys="image"),
-        Orientationd(keys=["image", "label"], axcodes="RAS")
-    ])
+        Orientationd(keys=["image", "label"], axcodes="RAS"),
+    ]
+    return Compose(transforms)
 
 def get_random_transforms( input_dim, num_samples, Nclasses, bg_ratio: float = 0.0):
     """

octopi/datasets/generators.py

@@ -92,7 +92,7 @@ def get_data_splits(
         # Get Class Info from the Training Dataset (classes are identical across
         # resolutions, so read from the first requested voxel size).
         target_info = (self.target_name, self.target_session_id, self.target_user_id)
-        self.Nclasses, self.class_names = utils.get_class_info(
+        self.Nclasses, self.class_names, self.label_space, self.model_labels = utils.get_class_info(
             self.config, self.myRunIDs['train'].keys(), target_info, self.resolutions[0][1] )
 
         return self.myRunIDs
@@ -119,6 +119,11 @@ def create(self,
         # Define the Input Dimensions
         self.input_dim = crop_size, crop_size, crop_size
 
+        # Build the global->dense compaction map once (shared by train/val transforms).
+        # Labels are painted with copick global values on disk; the model needs dense
+        # [0..K]. Identity-collapsing for legacy sequential data.
+        orig_labels, target_labels = utils.build_compaction_map(self.root, self.model_labels, self.label_space)
+
         # Create the list of training files. One entry per (run, resolution):
         # each (alg, voxel_size) pair is loaded at its native spacing with its
         # own matching target, so the model trains across all resolutions.
@@ -133,7 +138,7 @@ def create(self,
         # Default Train Transforms for Particle Picking
         if train_transforms is None:
             train_transforms = Compose([
-                augment.get_transforms(),
+                augment.get_transforms(orig_labels=orig_labels, target_labels=target_labels),
                 augment.get_random_transforms(self.input_dim, num_samples, self.Nclasses, self.bgr)
             ])
 
@@ -194,7 +199,7 @@ def create(self,
         # systematically underestimates F1 by splitting particles at patch
         # boundaries. GPU memory stays bounded by the trainer's sw_batch_size
         # (only a few ROI windows resident at once), not by the volume size.
-        val_ds = CacheDataset(data=val_files, transform=augment.get_transforms(), cache_rate=1.0, num_workers=4)
+        val_ds = CacheDataset(data=val_files, transform=augment.get_transforms(orig_labels=orig_labels, target_labels=target_labels), cache_rate=1.0, num_workers=4)
 
         # batch_size=1: validation is per full volume (volumes differ in shape
         # and cannot be stacked). Throughput is governed by the trainer's
@@ -372,9 +377,26 @@ def get_data_splits(
 
         # Get Class Info from the Training Dataset (classes identical across
         # resolutions, so read from the first requested voxel size).
-        self.Nclasses, self.class_names = utils.get_class_info(
+        self.first_session = first_session
+        self.Nclasses, self.class_names, self.label_space, self.model_labels = utils.get_class_info(
             config_path, self.myRunIDs['train'][first_session].keys(), target_info, self.resolutions[0][1] )
 
+        # A single shared compaction map is derived from the first session only, so every
+        # config must assign the same copick global label to each selected object. Fail loudly
+        # otherwise (per-config maps are out of scope). Only meaningful in copick label space.
+        if self.label_space == 'copick':
+            ref = {name: self.roots[first_session].get_object(name).label for name in self.class_names}
+            for session_key, root in self.roots.items():
+                other = {
+                    name: (root.get_object(name).label if root.get_object(name) is not None else None)
+                    for name in self.class_names
+                }
+                if other != ref:
+                    raise ValueError(
+                        f"Copick global labels disagree across configs: session '{first_session}' has "
+                        f"{ref} but session '{session_key}' has {other}. MultiCopickDataModule derives one "
+                        f"shared copick->model map from the first session; per-config maps are unsupported." )
+
         return self.myRunIDs
 
     def create(self, 
@@ -390,6 +412,12 @@ def create(self,
         # Define the Input Dimensions
         self.input_dim = crop_size, crop_size, crop_size
 
+        # Build the copick->model compaction map once from the first session's root
+        # (all configs agree on the schema, asserted in get_data_splits). Shared by
+        # train/val transforms; (None, None) for legacy targets → no-op.
+        orig_labels, target_labels = utils.build_compaction_map(
+            self.roots[self.first_session], self.model_labels, self.label_space)
+
         # Create the list of training files (one entry per session × run × resolution).
         train_files = [
             { 'run': run_id,
@@ -404,7 +432,7 @@ def create(self,
         # Default Train Transforms for Particle Picking
         if train_transforms is None:
             train_transforms = Compose([
-                augment.get_transforms(),
+                augment.get_transforms(orig_labels=orig_labels, target_labels=target_labels),
                 augment.get_random_transforms(self.input_dim, num_samples, self.Nclasses, self.bgr)
             ])
 
@@ -455,7 +483,7 @@ def create(self,
         # systematically underestimates F1 by splitting particles at patch
         # boundaries. GPU memory stays bounded by the trainer's sw_batch_size
         # (only a few ROI windows resident at once), not by the volume size.
-        val_ds = CacheDataset(data=val_files, transform=augment.get_transforms(), cache_rate=1.0, num_workers=4)
+        val_ds = CacheDataset(data=val_files, transform=augment.get_transforms(orig_labels=orig_labels, target_labels=target_labels), cache_rate=1.0, num_workers=4)
 
         # batch_size=1: validation is per full volume (volumes differ in shape
         # and cannot be stacked). Throughput is governed by the trainer's

octopi/datasets/helpers.py

@@ -245,7 +245,7 @@ def get_data_splits(
 
     return myRunIDs
 
-def get_class_info(config, runIDs, target_info, voxel_size) -> tuple[int, list[str]]:
+def get_class_info(config, runIDs, target_info, voxel_size) -> tuple[int, list[str], str, dict]:
     """
     Get the number of classes and class names from a segmentation.
 
@@ -256,8 +256,10 @@ def get_class_info(config, runIDs, target_info, voxel_size) -> tuple[int, list[s
         voxel_size: The voxel size of the segmentation.
 
     Returns:
-        Nclasses: The number of classes.
-        class_names: The list of class names.
+        Nclasses: The number of classes (foreground objects + background).
+        class_names: The object names sorted by ascending recorded label (model channel).
+        label_space: 'copick' when the persisted seg holds copick global labels, else None.
+        model_labels: The recorded {name: model_label} map from the targets YAML.
     """
 
     # Load the Copick Config
@@ -285,14 +287,46 @@ def get_class_info(config, runIDs, target_info, voxel_size) -> tuple[int, list[s
             config, target_name, 'targets', 
             target_user_id, target_session_id
         )
-        class_names = target_config['input']['labels']
-        Nclasses = len(class_names) + 1
-        class_names = [name for name, idx in sorted(class_names.items(), key=lambda x: x[1])]
+        # `labels` maps name -> recorded label (model channel in copick space, else the
+        # legacy sequential value). `label_space` marks whether the seg holds copick globals.
+        model_labels = target_config['input']['labels']
+        label_space = target_config['input'].get('label_space')
+        Nclasses = len(model_labels) + 1
+        # Order object names by their model channel so class_names[i] aligns to channel i+1.
+        class_names = [name for name, _ in sorted(model_labels.items(), key=lambda kv: kv[1])]
 
         # We Only need to read One Segmentation to Get Class Info
-        break      
+        break
 
-    return Nclasses, class_names
+    return Nclasses, class_names, label_space, model_labels
+
+def build_compaction_map(root, model_labels, label_space):
+    """
+    Build the (orig_labels, target_labels) pair for MONAI's MapLabelValued that compacts the
+    persisted copick *global* labels into the model's dense channels, for ``label_space ==
+    'copick'`` targets. The mapping is read explicitly from the recorded ``model_labels``
+    ({name: model_label}) joined with the project config's global labels by name — no sorting.
+
+    Legacy targets (``label_space`` not 'copick') are already in model/dense space, so no remap
+    is needed and ``(None, None)`` is returned (``get_transforms`` then skips MapLabelValued).
+
+    In copick mode every project object's global label maps to its model channel; objects that
+    are painted but not selected as training targets fold to background 0 (safe — copick global
+    labels are unique). Background maps to background.
+
+    Args:
+        root: A copick root for the (training) project the labels are painted in.
+        model_labels: The recorded {name: model_label} map (from get_class_info).
+        label_space: 'copick' to build the map, otherwise legacy no-op.
+
+    Returns:
+        (orig_labels, target_labels) for ``MapLabelValued``, or (None, None) for legacy.
+    """
+    if label_space != 'copick':
+        return None, None
+    orig_labels = [0] + [o.label for o in root.pickable_objects]
+    target_labels = [0] + [model_labels.get(o.name, 0) for o in root.pickable_objects]
+    return orig_labels, target_labels
 
 def build_target_uri(name: str, sessionid: str | None, userid: str | None, voxel_size: float) -> str:
     """