From 33123517f5089b5560627390ff38a1b05b3002a8 Mon Sep 17 00:00:00 2001 From: uermel Date: Wed, 1 Jul 2026 13:19:54 -0700 Subject: [PATCH] proposal for copick-consistent input/output --- octopi/datasets/augment.py | 37 ++++-- octopi/datasets/generators.py | 40 +++++- octopi/datasets/helpers.py | 50 ++++++-- octopi/entry_points/run_create_targets.py | 116 ++++++++++++++---- .../processing/create_targets_from_picks.py | 24 ++-- octopi/pytorch/segmentation.py | 69 ++++++++++- octopi/utils/io.py | 8 +- octopi/workflows.py | 13 +- 8 files changed, 291 insertions(+), 66 deletions(-) diff --git a/octopi/datasets/augment.py b/octopi/datasets/augment.py index 25d8464..2074598 100755 --- a/octopi/datasets/augment.py +++ b/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): """ diff --git a/octopi/datasets/generators.py b/octopi/datasets/generators.py index 4ae7049..8b5ac3d 100755 --- a/octopi/datasets/generators.py +++ b/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 diff --git a/octopi/datasets/helpers.py b/octopi/datasets/helpers.py index 7e95a75..2fd5755 100644 --- a/octopi/datasets/helpers.py +++ b/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: """ diff --git a/octopi/entry_points/run_create_targets.py b/octopi/entry_points/run_create_targets.py index 1c13c63..0b181b7 100755 --- a/octopi/entry_points/run_create_targets.py +++ b/octopi/entry_points/run_create_targets.py @@ -8,27 +8,31 @@ def create_sub_train_targets( pick_targets: List[Tuple[str, Union[str, None], Union[str, None]]], seg_targets: List[Tuple[str, Union[str, None], Union[str, None]]], voxel_size: float, - radius_scale: float, + radius_scale: float, tomogram_algorithm: str, target_segmentation_name: str, target_user_id: str, target_session_id: str, - run_ids: List[str], + run_ids: List[str], + label_space: str = 'model', + from_model: str = None, ): import octopi.processing.create_targets_from_picks as create_targets import copick - # Load Copick Project + # Load Copick Project root = copick.from_file(config) # Create empty dictionary for all targets train_targets = defaultdict(dict) - # Create dictionary for particle targets - value = 1 + # Create dictionary for particle targets. Labels are assigned sequentially in + # argument order (legacy model/dense space); `finalize_label_space` may re-express + # them in copick global space below when `--label-space copick` is requested. + value = 1 for t in pick_targets: # Parse the target - obj_name, user_id, session_id = t + obj_name, user_id, session_id = t obj = root.get_object(obj_name) # Check if the object is valid @@ -39,9 +43,9 @@ def create_sub_train_targets( if obj_name in train_targets: print(f'Warning - Skipping Particle Target: "{obj_name}, {user_id}, {session_id}", as it has already been added to the target list.') continue - - # Get the label and radius of the object - label = value # Assign labels sequentially + + # Assign labels sequentially + label = value info = { "label": label, "user_id": user_id, @@ -52,13 +56,16 @@ def create_sub_train_targets( train_targets[obj_name] = info value += 1 - # Create dictionary for segmentation targets - train_targets = add_segmentation_targets(root, seg_targets, train_targets, value) + # Create dictionary for segmentation targets (sequential, continuing the count) + train_targets = add_segmentation_targets(root, seg_targets, train_targets, value) + + # Optionally re-express targets in copick GLOBAL label space (explicit label_space). + label_space = finalize_label_space(root, train_targets, label_space, from_model) create_targets.generate_targets( config, train_targets, voxel_size, tomogram_algorithm, radius_scale, - target_segmentation_name, target_user_id, - target_session_id, run_ids + target_segmentation_name, target_user_id, + target_session_id, run_ids, label_space=label_space, ) @@ -73,12 +80,14 @@ def create_all_train_targets( target_segmentation_name: str, target_user_id: str, target_session_id: str, - run_ids: List[str], - ): + run_ids: List[str], + label_space: str = 'model', + from_model: str = None, + ): import octopi.processing.create_targets_from_picks as create_targets import copick - # Load Copick Project + # Load Copick Project root = copick.from_file(config) # Create empty dictionary for all targets @@ -98,17 +107,20 @@ def create_all_train_targets( # Create dictionary for segmentation targets target_objects = add_segmentation_targets(root, seg_targets, target_objects) + # Optionally re-express targets in copick GLOBAL label space (explicit label_space). + label_space = finalize_label_space(root, target_objects, label_space, from_model) + create_targets.generate_targets( - config, target_objects, voxel_size, tomogram_algorithm, - radius_scale, target_segmentation_name, target_user_id, - target_session_id, run_ids + config, target_objects, voxel_size, tomogram_algorithm, + radius_scale, target_segmentation_name, target_user_id, + target_session_id, run_ids, label_space=label_space, ) def add_segmentation_targets( root, seg_targets, train_targets: dict, - start_value: int = -1 + start_value: int = -1, ): # Create dictionary for segmentation targets @@ -117,7 +129,8 @@ def add_segmentation_targets( # Parse Segmentation Target obj_name, user_id, session_id = s - # Add Segmentation Target + # Assign a sequential label (continuing the particle count) or fall back to the + # object's global label; `finalize_label_space` re-expresses these when needed. if start_value > 0: value = start_value start_value += 1 @@ -129,8 +142,8 @@ def add_segmentation_targets( "label": value, "user_id": user_id, "session_id": session_id, - "is_particle_target": False, - "radius": None, + "is_particle_target": False, + "radius": None, } train_targets[obj_name] = info @@ -138,7 +151,48 @@ def add_segmentation_targets( except: print(f'Warning - Skipping Segmentation Name: "{obj_name}", as it is not a valid object in the Copick project.') - return train_targets + return train_targets + + +def finalize_label_space(root, train_targets, label_space, from_model): + """ + Optionally re-express the target labels in copick GLOBAL label space. + + Default (``label_space='model'``): no-op — the recorded ``label`` stays the model/dense + value and is what gets painted into the segmentation (legacy behavior, unchanged). + + ``label_space='copick'``: the segmentation is painted with each object's copick GLOBAL + label (``root.get_object(name).label``), while the recorded ``label`` becomes the + model's dense channel (``model_label``). Channels are minted by ascending global label, + or, with ``from_model``, inherited per-name from a pretrained model config so fine-tuning + keeps the pretrained channel<->object binding. Returns the marker to persist in the + targets YAML (``'copick'`` or ``None``). + """ + if label_space != 'copick': + return None + + from octopi.utils import io + + names = list(train_targets.keys()) + copick_labels = {n: root.get_object(n).label for n in names} + + if from_model: + pretrained = io.load_yaml(from_model).get('labels') or {} + missing = [n for n in names if n not in pretrained] + if missing: + raise ValueError( + f"--from-model {from_model} does not define model labels for {missing}; cannot " + f"align channels for fine-tuning (fine-tuning onto a new object set is out of scope)." ) + model_labels = {n: int(pretrained[n]) for n in names} + else: + # Mint contiguous dense channels 1..K by ascending copick global label. + order = sorted(names, key=lambda n: copick_labels[n]) + model_labels = {n: i + 1 for i, n in enumerate(order)} + + for n in names: + train_targets[n]['paint_label'] = copick_labels[n] # canonical copick label on disk + train_targets[n]['label'] = model_labels[n] # recorded model (dense) channel + return 'copick' @click.command('create-targets', no_args_is_help=True) @@ -172,9 +226,17 @@ def add_segmentation_targets( help='Target specifications: "name" or "name,user_id,session_id"') @click.option('-c', '--config', type=click.Path(exists=True), required=True, help="Path to the CoPick configuration file") +@click.option('--label-space', type=click.Choice(['model', 'copick']), default='model', + help="Label space of the persisted targets. 'model' (default): legacy dense/sequential " + "labels. 'copick': paint canonical copick GLOBAL labels and record model channels " + "(name -> model_label) plus a `label_space: copick` marker in the targets YAML.") +@click.option('--from-model', type=click.Path(exists=True), default=None, + help="With --label-space copick: inherit each object's model channel from this pretrained " + "model_config.yaml (by name) so fine-tuning keeps the pretrained channel binding.") def cli(config, target, picks_session_id, picks_user_id, seg_target, run_ids, tomo_alg, radius_scale, voxel_size, - target_segmentation_name, target_user_id, target_session_id): + target_segmentation_name, target_user_id, target_session_id, + label_space, from_model): """ Generate segmentation targets from CoPick configurations. @@ -215,6 +277,8 @@ def cli(config, target, picks_session_id, picks_user_id, seg_target, run_ids, target_user_id=target_user_id, target_session_id=target_session_id, run_ids=run_ids, + label_space=label_space, + from_model=from_model, ) else: # If no --target is provided, call create_all_train_targets @@ -230,6 +294,8 @@ def cli(config, target, picks_session_id, picks_user_id, seg_target, run_ids, target_user_id=target_user_id, target_session_id=target_session_id, run_ids=run_ids, + label_space=label_space, + from_model=from_model, ) diff --git a/octopi/processing/create_targets_from_picks.py b/octopi/processing/create_targets_from_picks.py index c234b46..64c8e2e 100755 --- a/octopi/processing/create_targets_from_picks.py +++ b/octopi/processing/create_targets_from_picks.py @@ -60,6 +60,7 @@ def generate_targets( target_user_name: str = 'octopi', target_session_id: str = '1', run_ids: List[str] = None, + label_space: str = None, ): """ Generate segmentation targets from picks in CoPick configuration. @@ -130,9 +131,10 @@ def generate_targets( voxel_size=voxel_size ) - # Add Segmentations to Target + # Add Segmentations to Target. Paint the value the seg should hold on disk: + # 'paint_label' (copick global) when in copick label space, else the recorded 'label'. for seg in query_seg: - classLabel = train_targets[seg.name]['label'] + classLabel = train_targets[seg.name].get('paint_label', train_targets[seg.name]['label']) segvol = seg.numpy() if segvol.shape != target.shape: try: @@ -163,10 +165,11 @@ def generate_targets( # Add Picks to Target for pick in query: numPicks += len(pick.points) - target = from_picks(pick, - target, - train_targets[pick.pickable_object_name]['radius'] * radius_scale, - train_targets[pick.pickable_object_name]['label'], + pinfo = train_targets[pick.pickable_object_name] + target = from_picks(pick, + target, + pinfo['radius'] * radius_scale, + pinfo.get('paint_label', pinfo['label']), voxel_size ) @@ -196,6 +199,7 @@ def generate_targets( "target_session_id": target_session_id, "voxel_size": voxel_size, "labels": labels, + "label_space": label_space, } target_query = f'{target_user_name}_{target_session_id}_{target_segmentation_name}' print(f'💾 Saving parameters to {basepath}/targets-{target_query}.yaml') @@ -214,13 +218,17 @@ def save_parameters(args, basepath: str, target_query: str): basepath: Path to save the YAML file. target_query: Query string for target identification. """ - # Prepare input group + # Prepare input group. `labels` records name -> recorded label (model channel in copick + # mode, sequential in legacy). The optional `label_space: copick` marker tells consumers the + # persisted seg holds copick GLOBAL labels and to remap copick<->model by name. keys = ['user_id', 'session_id'] input_group = { "config": args['config'], - "labels": {name: info['label'] for name, info in args['train_targets'].items()}, # <-- Added comma here + "labels": {name: info['label'] for name, info in args['train_targets'].items()}, "targets": {name: {k: info[k] for k in keys} for name, info in args['train_targets'].items()} } + if args.get('label_space'): + input_group['label_space'] = args['label_space'] # Organize parameters into subgroups new_entry = { diff --git a/octopi/pytorch/segmentation.py b/octopi/pytorch/segmentation.py index 002f7e7..7fd5a8a 100755 --- a/octopi/pytorch/segmentation.py +++ b/octopi/pytorch/segmentation.py @@ -131,10 +131,12 @@ def _load_models(self, model_config: List[str], model_weights: List[str]): """Load a single model or multiple models for soup.""" self.models = [] + configs_meta = [] for i, (config_path, weights_path) in enumerate(zip(model_config, model_weights)): # Load the Model Config and Model Builder current_modelconfig = io.load_yaml(config_path) + configs_meta.append((current_modelconfig.get('labels'), current_modelconfig.get('label_space'))) model_builder = common.get_model(current_modelconfig['model']['architecture']) # Check if the weights file exists @@ -164,13 +166,68 @@ def _load_models(self, model_config: List[str], model_weights: List[str]): self.model = self.models[0] # Set the Number of Classes and Input Dimensions - Assume All Models are the Same - self.Nclass = current_modelconfig['model']['num_classes'] + self.Nclass = current_modelconfig['model']['num_classes'] self.dim_in = current_modelconfig['model']['dim_in'] self.input_dim = None + # Build the dense->global expansion map (resolved by name against the target project). + self._build_dense_to_global(configs_meta) + # Print a message if Model Soup is Enabled if self.apply_modelsoup: self._print(f'Model Soup is Enabled : {len(self.models)} models loaded for ensemble inference') + + def _build_dense_to_global(self, configs_meta): + """ + Build ``self.dense_to_global``: a lookup that expands the model's dense ``[0..K]`` + predictions back into canonical copick GLOBAL labels of the TARGET project before writing. + + In copick label space, the model config records ``{name: model_label}`` plus + ``label_space: copick``. Each channel is resolved BY OBJECT NAME against the target + project (``self.root``, from the ``-c `` passed to ``octopi segment``): + ``dense_to_global[model_label] = self.root.get_object(name).label``. This makes a trained + model reusable across copick projects whose global labels differ. + + Legacy (no ``label_space``) or label-less configs fall back to identity — predictions are + written in the model's dense space unchanged, exactly as before. + + Args: + configs_meta: list of ``(labels, label_space)`` tuples, one per ensemble member. + """ + # Ensemble: every model in the soup must share the same label schema + space. + ref = configs_meta[0] + for meta in configs_meta[1:]: + if meta != ref: + raise ValueError( + f"Ensemble model configs disagree on labels/label_space: {ref} vs {meta}. " + f"All models in a soup must share the same label schema." ) + labels, label_space = ref + + # Legacy / label-less -> identity: predictions written in dense space, unchanged. + if label_space != 'copick' or not labels: + self.dense_to_global = np.arange(self.Nclass, dtype=np.uint8) + return + + if len(labels) != self.Nclass - 1: + raise ValueError( + f"Model config 'labels' has {len(labels)} objects but num_classes={self.Nclass} " + f"implies {self.Nclass - 1}. The label schema is inconsistent." ) + + dense_to_global = np.zeros(self.Nclass, dtype=np.uint8) # background 0 -> 0 + for name, model_label in labels.items(): + model_label = int(model_label) + if not 0 < model_label < self.Nclass: + raise ValueError( + f"Model label {model_label} for '{name}' is out of range [1, {self.Nclass - 1}]." ) + obj = self.root.get_object(name) + if obj is None: + raise ValueError( + f"Model predicts object '{name}' but it is not defined in the target copick " + f"project ({self.config}); cannot assign it a global label. Run inference " + f"against a project that defines '{name}', or add it to the project config." ) + dense_to_global[model_label] = obj.label + + self.dense_to_global = dense_to_global @torch.inference_mode() def _run_single_model_inference(self, model, input_data, out_device: torch.device): @@ -399,6 +456,9 @@ def batch_predict( # Write the Prediction to the corresponding runs for i in range(len(preds)): pred = preds[i].cpu().numpy() + # Expand dense [0..K] -> canonical copick global labels for the target + # project, so the stored `predict` segmentation is ordinary multilabel copick. + pred = self.dense_to_global[pred] run = self.root.get_run(data['runid'][i]) writers.segmentation( run, pred, @@ -456,14 +516,17 @@ def save_parameters(self, # Load the model config model_config = io.load_yaml(self.model_config[0]) - # Create parameters dictionary + # Mirror the model's label schema into segment-*.yaml: `labels` ({name: model_label}) + # and, in copick label space, the `label_space` marker. Downstream `localize` reads + # these to interpret the predict seg (copick globals vs legacy model space). params = { "inputs": { "config": self.config, "tomo_alg": tomo_algorithm, "voxel_size": voxel_size }, - 'labels': model_config['labels'], + 'label_space': model_config.get('label_space'), + 'labels': model_config.get('labels'), 'model': { 'configs': self.model_config, 'weights': self.model_weights diff --git a/octopi/utils/io.py b/octopi/utils/io.py index 9347124..12e9c58 100644 --- a/octopi/utils/io.py +++ b/octopi/utils/io.py @@ -153,16 +153,20 @@ def save_parameters_to_yaml(model, trainer, dataloader, filename: str): # Check for the target configuration file for model labels target_config = check_target_config_path(dataloader) - # Extract and flatten parameters + # Extract and flatten parameters. Carry the `label_space` marker (when present) so the + # model config records that its `labels` are model channels over copick-global targets, + # which inference/fine-tuning/localize use to remap copick<->model by name. parameters = { 'model': model.get_model_parameters(), 'labels': target_config['input']['labels'], 'optimizer': get_optimizer_parameters(trainer), 'dataloader': dataloader.get_dataloader_parameters() } + if target_config['input'].get('label_space'): + parameters['label_space'] = target_config['input']['label_space'] save_parameters_yaml(parameters, filename) - print(f"⚙️ Training Parameters saved to {filename}") + print(f"⚙️ Training Parameters saved to {filename}") def flatten_params(params, parent_key=''): """ diff --git a/octopi/workflows.py b/octopi/workflows.py index e765a4d..645b9c2 100644 --- a/octopi/workflows.py +++ b/octopi/workflows.py @@ -183,14 +183,19 @@ def localize(config, voxel_size, seg_info, pick_user_id, pick_session_id, n_proc # Load the Model Output Configuration seg_config = io.get_config(config, seg_info[0], 'segment', seg_info[1], seg_info[2]) - # sync labels from the model config and remove objects not in model labels + # Sync labels with the model config and drop objects the model does not predict. + # In copick label space the predict seg already holds copick GLOBAL labels (== obj.label), + # so no value remap is needed. In legacy space it holds the model's dense labels, so remap + # each object's global label to its model channel (old behavior). label_map = seg_config.get('labels', {}) + label_space = seg_config.get('label_space') for row in objects.copy(): # avoid modifying the list while iterating name, label, radius = row - if name in label_map and label != label_map[name]: - row[1] = int(label_map[name]) # mutate in place - elif name not in label_map: # remove this entry from objects + if name not in label_map: # object not predicted by the model -> drop it objects.remove(row) + continue + if label_space != 'copick' and label != label_map[name]: + row[1] = int(label_map[name]) # legacy: remap global -> model dense label # Filter objects based on the provided list if pick_objects is not None: