diff --git a/development/check_apg_tiled_refinement.py b/development/check_apg_tiled_refinement.py index 8360633f2..115f131c4 100644 --- a/development/check_apg_tiled_refinement.py +++ b/development/check_apg_tiled_refinement.py @@ -112,9 +112,10 @@ def main(): tiled = _build(model, decoder, args.device, is_tiled=True) # One tile covering the image. Its outer block is clipped to the image, so the halo is irrelevant. tiled.initialize(image, ndim=2, tile_shape=tuple(image.shape[:2]), halo=(0, 0)) - tiled_plain = tiled.generate() + tiled_proposals = tiled.propose() + tiled_plain = tiled.select(tiled_proposals) tiled._last_generation_stats = {} - tiled_refined = tiled.generate(**generate_kwargs) + tiled_refined = tiled.select(tiled_proposals, **generate_kwargs) _report("tiled, one tile", tiled_refined, labels, tiled._last_generation_stats) tiled.clear_state() @@ -129,9 +130,11 @@ def main(): print(f"\nSmoke run with tiles {tuple(args.tile_shape)} and halo {tuple(args.halo)}:") tiled.initialize(image, ndim=2, tile_shape=tuple(args.tile_shape), halo=tuple(args.halo)) + # One round of prompting for both, as the screening harness does: only the selection differs. + proposals = tiled.propose() for name, kwargs in (("tiled, no refinement", {}), ("tiled, refined", generate_kwargs)): tiled._last_generation_stats = {} - segmentation = tiled.generate(**kwargs) + segmentation = tiled.select(proposals, **kwargs) _report(name, segmentation, labels, tiled._last_generation_stats) tiled.clear_state() diff --git a/finetuning/v2/evaluation/common.py b/finetuning/v2/evaluation/common.py index 4c5c12b38..49d21448b 100644 --- a/finetuning/v2/evaluation/common.py +++ b/finetuning/v2/evaluation/common.py @@ -1,4 +1,5 @@ import os +from pathlib import Path import re import ast import csv @@ -1708,36 +1709,29 @@ def resolve_params(overrides=None, ndim=2, model_type=None): return params -def load_apg_overrides(path, dataset_name): - """Read one APG configuration file and return its name and the overrides for one dataset. +def load_apg_overrides(path): + """Read one APG configuration file and return its name and raw 2d parameter overrides. - The file has the format of the optimization benchmark: ``{"name": ..., "params_2d": {...}, - "params_3d": {...}}``, with an optional ``params_dense``. Images use 'params_2d' and volumes use - 'params_3d'. The dense-neuron EM volumes use 'params_dense' if the file has it. The function returns - the overrides unresolved, so that they can go on top of tuned parameters. `resolve_params` fills in - the defaults. + The file has the shape the optimization benchmark uses, ``{"name": ..., "params_2d": {...}}`` + (``params_3d`` may be present and is ignored here). The overrides are returned unresolved, so + they can be layered over tuned parameters; `resolve_params` fills in the defaults. Args: path: The JSON configuration file. - dataset_name: The dataset that the overrides are for. It selects the section. Returns: - The configuration name and the overrides, keyed as `generate` takes them. + The configuration name and the 2d overrides, keyed as `generate` takes them. """ import json with open(path) as f: config = json.load(f) - unknown_top_level = set(config) - {"name", "params_2d", "params_3d", "params_dense"} + unknown_top_level = set(config) - {"name", "params_2d", "params_3d"} if unknown_top_level: raise ValueError(f"Unknown configuration fields in '{path}': {sorted(unknown_top_level)}.") - if dataset_name in DATASETS_DENSE and "params_dense" in config: - section = "params_dense" - else: - section = "params_3d" if dataset_name in DATASETS_3D else "params_2d" - overrides = config.get(section, {}) + overrides = config.get("params_2d", {}) if not isinstance(overrides, dict): - raise TypeError(f"'{section}' in '{path}' must be an object.") + raise TypeError(f"'params_2d' in '{path}' must be an object.") unknown = set(overrides) - set(GENERATE_PARAM_KEYS) if unknown: raise ValueError(f"Unknown APG parameters in '{path}': {sorted(unknown)}.") @@ -2023,7 +2017,7 @@ def predict_unisam2(model, raw, ndim, device, normalization=None, devices=None): def postprocess_unisam2(out, dataset_name, model_type, params=None): - """Turn a (4, *spatial) prediction into an instance segmentation. + """Turn a (4, *spatial) prediction (or (5, *spatial) with a contact channel) into an instance segmentation. EM datasets use the dense (multicut) mode, all others the sparse (flow) mode. 'params' overrides the postprocessing defaults, e.g. with the best combination found by grid_search_automatic_cells. @@ -2039,7 +2033,8 @@ def postprocess_unisam2(out, dataset_name, model_type, params=None): seg = run_multicut(boundary_map, distances, model_type=model_type, **params) else: spacing = DATASET_SPACING.get(dataset_name, None) - seg = flow_instance_segmentation(fg, out[1:], model_type=model_type, spacing=spacing, **params) + contact = {"contact": out[4]} if out.shape[0] > 4 else {} + seg = flow_instance_segmentation(fg, out[1:4], model_type=model_type, spacing=spacing, **contact, **params) return seg.astype("uint32") diff --git a/finetuning/v2/evaluation/evaluate_automatic_segmentation.py b/finetuning/v2/evaluation/evaluate_automatic_segmentation.py index 31d4cb7eb..8347d4764 100644 --- a/finetuning/v2/evaluation/evaluate_automatic_segmentation.py +++ b/finetuning/v2/evaluation/evaluate_automatic_segmentation.py @@ -26,27 +26,51 @@ import torch from common import ( - DATA_ROOT, DATASETS_2D, DATASETS_3D, DATASET_SPACING, MODEL_TYPES, MODES, VOLUME_SPEED_OPTIONS, build_model, - check_data_download, evaluate_samples, has_val_split, load_apg_overrides, postprocess_unisam2, predict_unisam2, - read_tuned_params, resolve_checkpoint_identity, + DATA_ROOT, DATASETS_2D, DATASETS_3D, DATASETS_DENSE, DATASET_SPACING, GT_MIN_SIZE_2D, MODEL_TYPES, MODES, + VOLUME_SPEED_OPTIONS, build_model, check_data_download, drop_severed_objects, genuine_misses, + has_val_split, load_apg_overrides, load_data, n_samples, postprocess_unisam2, predict_unisam2, + read_tuned_params, resolve_checkpoint_identity, run_dataset_evaluation, ) def segment(model, mode, raw, ndim, dataset_name, model_type, params, device, spacing=None, devices=None): - """Segment one sample with the tuned parameters of a mode.""" + """Segment one sample with the tuned parameters of a mode. + + For 'ais' the parameters may be the nested form ``{"sparse": {...}, "dense": {...}}`` of an AIS + benchmark configuration (see `load_ais_params`); the dataset's pipeline picks its own dict. + """ if mode == "apg": model.clear_state() model.initialize(raw, ndim=ndim, **(VOLUME_SPEED_OPTIONS if ndim == 3 else {})) volume_params = {"spacing": spacing} if ndim == 3 else {} return model.generate(**{**volume_params, **params}).astype("uint32") + if set(params) & {"sparse", "dense"}: + params = params["dense" if dataset_name in DATASETS_DENSE else "sparse"] prediction = predict_unisam2(model, raw, ndim=ndim, device=device, devices=devices) return postprocess_unisam2(prediction, dataset_name, model_type=model_type, params=params) +def load_ais_params(path, model_type, ndim): + """Read an AIS benchmark configuration and resolve its parameters for images or volumes. + + The file has the shape `benchmark_ais_optimization.py` uses (``{"name", "mode", "params_2d", + "params_3d"}``); the result is ``{"sparse": {...}, "dense": {...}}`` with every post-processing + keyword resolved against the library defaults, so the evaluation runs exactly the benchmarked + configuration. + """ + import sys + from pathlib import Path + sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "optimization")) + from benchmark_ais_optimization import load_config + + name, _, params_2d, params_3d = load_config(Path(path), model_type) + return name, (params_3d if ndim == 3 else params_2d) + + def run_evaluation( - model, mode, dataset_name, data_root, experiment_folder, model_type, params, device, limit, - crop_shape=None, checkpoint_id=None, devices=None, tuned=None, result_tag=None, config_name=None, sample_index=None, + model, mode, dataset_name, data_root, experiment_folder, model_type, params, device, + crop_shape=None, checkpoint_id=None, devices=None, tuned=None, result_tag=None, config_name=None, ): """Score the test split with the given parameters and write the result CSV. @@ -72,7 +96,6 @@ def run_evaluation( result_tag: Optional tag appended to the result file name, so that a run with explicit parameter overrides does not collide with the plain evaluation. config_name: The name of the configuration the overrides came from, stored in the results. - sample_index: The index of the only sample to score, for one array task. See `common.evaluate_samples`. Returns: The results as a DataFrame, or None while the rows of other samples are missing. @@ -82,8 +105,6 @@ def run_evaluation( tag = "tuned" if tuned else "default" if result_tag: tag = f"{tag}_{result_tag}" - if limit is not None: - tag = f"{tag}_n{limit}" legacy_path = os.path.join( experiment_folder, "results", f"{dataset_name}_micro_sam2_{model_type}_{mode}_{tag}.csv" ) @@ -100,16 +121,40 @@ def run_evaluation( ndim = 3 if dataset_name in DATASETS_3D else 2 spacing = DATASET_SPACING.get(dataset_name) - extra_columns = {"parameters": json.dumps(params, sort_keys=True, default=str) if params else "default"} + border_min_size = GT_MIN_SIZE_2D.get(dataset_name, 0) if ndim == 2 else 0 + total = n_samples(dataset_name, data_root) + samples = load_data(dataset_name, data_root, ndim, crop_shape=crop_shape) + + all_gt, all_seg, misses = [], [], [] + for raw, labels, valid_roi in tqdm(samples, total=total, desc=f"{mode}-{model_type}"): + if labels.max() == 0: # Nothing to score without ground-truth. + continue + seg = segment( + model, mode, raw, ndim, dataset_name, model_type, params or {}, device, spacing=spacing, + devices=devices, + ) + if valid_roi is not None: + seg[~valid_roi] = 0 + if ndim == 2: + # The ground truth has no severed objects either, so predicting one is not a false positive. + seg = drop_severed_objects(seg, border_min_size) + else: + misses.append(genuine_misses(labels, seg)) + all_gt.append(labels) + all_seg.append(seg) + + os.makedirs(os.path.dirname(save_path), exist_ok=True) + results = run_dataset_evaluation(all_gt, all_seg, dataset_name, save_path) + if misses: + # The aggregate metric hides which objects went missing. + results["unmatched"] = sum(count[0] for count in misses) + results["genuine_misses"] = sum(count[1] for count in misses) + results["parameters"] = json.dumps(params, sort_keys=True, default=str) if params else "default" if config_name is not None: - extra_columns["config_name"] = config_name - return evaluate_samples( - lambda raw: segment( - model, mode, raw, ndim, dataset_name, model_type, params or {}, device, spacing=spacing, devices=devices, - ), - dataset_name, data_root, save_path, desc=f"{mode}-{model_type}", limit=limit, crop_shape=crop_shape, - sample_index=sample_index, extra_columns=extra_columns, - ) + results["config_name"] = config_name + results.to_csv(save_path, index=False) + print(results) + return results def main(): @@ -138,18 +183,25 @@ def main(): parser.add_argument("--devices", nargs="*", default=None, help="Inference devices. All visible GPUs by default.") parser.add_argument( "--apg_params", type=str, default=None, - help="APG only. A JSON configuration in the benchmark format. Its section for the dataset ('params_2d', " - "'params_3d' or 'params_dense') overrides the tuned parameters, or the defaults with --skip_tuning.", + help="APG only. A benchmark-style JSON configuration whose 'params_2d' are layered over the tuned " + "parameters (or the defaults with --skip_tuning).", + ) + parser.add_argument( + "--ais_params", type=str, default=None, + help="AIS only. An AIS benchmark configuration ('params_2d' / 'params_3d', flat or " + "{'sparse', 'dense'}) whose resolved post-processing parameters replace the tuned ones.", ) parser.add_argument( "--result_tag", type=str, default=None, - help="Tag appended to the result file name. Defaults to the --apg_params configuration name.", + help="Tag appended to the result file name. Defaults to the --apg_params / --ais_params configuration name.", ) args = parser.parse_args() check_data_download(args.dataset_name, args.input_path) if args.apg_params is not None and args.mode != "apg": parser.error("--apg_params applies to --mode apg only.") + if args.ais_params is not None and args.mode != "ais": + parser.error("--ais_params applies to --mode ais only.") print("Device:", torch.cuda.get_device_name() if torch.cuda.is_available() else "CPU") device = "cuda" if torch.cuda.is_available() else "cpu" @@ -190,16 +242,21 @@ def main(): config_name, result_tag = None, args.result_tag if args.apg_params is not None: - config_name, overrides = load_apg_overrides(args.apg_params, args.dataset_name) + config_name, overrides = load_apg_overrides(args.apg_params) params = {**(params or {}), **overrides} if result_tag is None: result_tag = config_name + if args.ais_params is not None: + # The configuration is complete (every keyword resolved), so it replaces rather than layers. + config_name, params = load_ais_params(args.ais_params, args.model_type, ndim) + tuned = False + if result_tag is None: + result_tag = config_name run_evaluation( model, args.mode, args.dataset_name, args.input_path, args.experiment_folder, args.model_type, params, device, crop_shape=crop_shape, checkpoint_id=checkpoint_id, devices=args.devices or None, tuned=tuned, result_tag=result_tag, config_name=config_name, - limit=args.n_samples, sample_index=args.sample_index, ) diff --git a/finetuning/v2/evaluation/optimization/ais_campaign_tasks.py b/finetuning/v2/evaluation/optimization/ais_campaign_tasks.py new file mode 100644 index 000000000..162792c09 --- /dev/null +++ b/finetuning/v2/evaluation/optimization/ais_campaign_tasks.py @@ -0,0 +1,144 @@ +"""Build the task lists of the AIS optimization campaign and hand them to the submitter. + +Each subcommand turns a few arguments into '(tag, command)' pairs for `benchmark_ais_optimization.py` +and submits them through `submit_optimization_jobs.submit_tasks`. '--extra' appends verbatim arguments +to every command. + +Usage examples: + # Cache the predictions of two subsets, one task per subset, on the session GPU. + python ais_campaign_tasks.py predict --name ais_predict --subsets primary training_extra --local + + # One CPU task per (subset, configuration): the screen of a candidate family against the baseline. + python ais_campaign_tasks.py screen --name s0_travel --preset cpu --subsets primary training_extra \\ + --configs configs/ais_control_registry_defaults.json configs/ais_s0_*.json + + # A parameter sweep, one task per (subset, dataset, shard). + python ais_campaign_tasks.py sweep --name lm_grid --preset cpu --subsets primary --grid configs/ais_grid_lm.json \\ + --datasets livecell tissuenet --num-shards 4 +""" + +from __future__ import annotations + +import argparse +import glob +import shlex +import sys +from pathlib import Path +from typing import Iterable, List, Optional, Sequence, Tuple + +OPTIMIZATION_ROOT = Path(__file__).resolve().parent +sys.path.insert(0, str(OPTIMIZATION_ROOT)) + +from submit_optimization_jobs import add_submit_arguments, sanitize, submit_from_args # noqa + +SCRIPT = OPTIMIZATION_ROOT / "benchmark_ais_optimization.py" + +Task = Tuple[str, str] + + +def _command(*args: object) -> str: + return shlex.join(["python", str(SCRIPT), *[str(arg) for arg in args]]) + + +def _config_stem(path: Optional[Path]) -> str: + if path is None: + return "defaults" + stem = Path(path).stem + return sanitize(stem[4:] if stem.startswith("ais_") else stem) + + +def _expand(patterns: Iterable[str]) -> List[Path]: + paths: List[Path] = [] + for pattern in patterns: + matches = sorted(glob.glob(pattern)) + if not matches: + raise FileNotFoundError(f"No configuration matches '{pattern}'.") + paths.extend(Path(match).resolve() for match in matches) + return paths + + +def predict_tasks(kind: str, subsets: Sequence[str], extra: Sequence[str] = ()) -> List[Task]: + """One `predict` task per subset.""" + return [ + (f"predict_{kind}_{sanitize(subset)}", _command("predict", "--kind", kind, "--subset", subset, *extra)) + for subset in subsets + ] + + +def run_tasks( + kind: str, subsets: Sequence[str], configs: Sequence[Optional[Path]], trial_ids: Sequence[str], + extra: Sequence[str] = (), +) -> List[Task]: + """One `run` task per (subset, configuration, trial).""" + tasks = [] + for subset in subsets: + for config in configs: + for trial in trial_ids: + args: List[object] = ["run", "--kind", kind, "--subset", subset, "--trial-id", trial] + if config is not None: + args.extend(["--config", config]) + args.extend(extra) + tag = f"run_{kind}_{sanitize(subset)}_{_config_stem(config)}_{sanitize(trial)}" + tasks.append((tag, _command(*args))) + return tasks + + +def sweep_tasks( + kind: str, subsets: Sequence[str], grid: Path, datasets: Sequence[str], num_shards: int, extra: Sequence[str] = (), +) -> List[Task]: + """One `sweep` task per (subset, dataset, shard).""" + tasks = [] + for subset in subsets: + for dataset in datasets: + for shard in range(num_shards): + args: List[object] = [ + "sweep", "--kind", kind, "--subset", subset, "--grid", grid, "--datasets", dataset, + "--shard-index", shard, "--num-shards", num_shards, *extra, + ] + tag = f"sweep_{kind}_{sanitize(subset)}_{sanitize(dataset)}_{shard}of{num_shards}" + tasks.append((tag, _command(*args))) + return tasks + + +def main(argv: Optional[Iterable[str]] = None) -> int: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + subparsers = parser.add_subparsers(dest="command", required=True) + + predict = subparsers.add_parser("predict", help="Cache the decoder predictions of subsets.") + screen = subparsers.add_parser("screen", help="Run configurations on the cache, one task each.") + screen.add_argument("--configs", nargs="*", default=[], help="Configuration files or globs.") + screen.add_argument("--no-defaults", action="store_true", help="Do not add the library-defaults baseline.") + screen.add_argument("--trial-ids", nargs="*", default=["trial-1"]) + sweep = subparsers.add_parser("sweep", help="Sweep a grid on the cache, one task per dataset and shard.") + sweep.add_argument("--grid", type=Path, required=True) + sweep.add_argument("--datasets", nargs="+", required=True) + sweep.add_argument("--num-shards", type=int, default=1) + + for sub in (predict, screen, sweep): + sub.add_argument("--kind", choices=("v5", "apg3d"), default="v5") + sub.add_argument("--subsets", nargs="+", default=["primary"]) + sub.add_argument("--extra", default="", help="Arguments appended verbatim to every command.") + sub.add_argument("--print-only", action="store_true", help="Print the tasks and stop.") + add_submit_arguments(sub) + + args = parser.parse_args(list(argv) if argv is not None else None) + extra = shlex.split(args.extra) if args.extra else [] + if args.command == "predict": + tasks = predict_tasks(args.kind, args.subsets, extra) + elif args.command == "screen": + configs: List[Optional[Path]] = list(_expand(args.configs)) + if not args.no_defaults: + configs = [None, *configs] + tasks = run_tasks(args.kind, args.subsets, configs, args.trial_ids, extra) + else: + tasks = sweep_tasks(args.kind, args.subsets, args.grid.resolve(), args.datasets, args.num_shards, extra) + for tag, command in tasks: + print(f"{tag}\t{command}") + if args.print_only: + return 0 + submit_from_args(tasks, args) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/finetuning/v2/evaluation/optimization/benchmark_ais_optimization.py b/finetuning/v2/evaluation/optimization/benchmark_ais_optimization.py new file mode 100644 index 000000000..77049bf50 --- /dev/null +++ b/finetuning/v2/evaluation/optimization/benchmark_ais_optimization.py @@ -0,0 +1,1639 @@ +"""Benchmark AIS (decoder-based automatic instance segmentation) post-processing on cached predictions. + +The UniSAM2 decoder prediction of a sample is a (4, *spatial) array of foreground probability and three +directed-distance channels, optionally followed by a fifth boundary channel. It does not depend on any +post-processing choice. This benchmark therefore +predicts every sample of a manifest once (`predict`, GPU), caches the prediction, and runs every +post-processing configuration, diagnostic and parameter sweep on the cache (CPU). A configuration run +still writes a canonical run directory in the layout of `benchmark_apg_optimization.py`, so +`compare_apg_optimization.py` reads it unchanged. + +Manifests are reused from the APG campaigns: the 2d subset manifests (`--kind v5`: primary, holdout, +training_extra and the sealed 180-image ood_extended set) and the deep 3d crop manifests +(`--kind apg3d`: primary, holdout, test). Nothing is rebuilt and the data root is read-only. + +Usage examples: + # Cache the predictions of the primary subset on the session GPU. + python benchmark_ais_optimization.py predict --kind v5 --subset primary + + # Run the library defaults on the cache (the baseline) and a candidate. + python benchmark_ais_optimization.py run --kind v5 --subset primary + python benchmark_ais_optimization.py run --kind v5 --subset primary --config configs/ais_travel_200.json + + # Screen several configurations on two subsets, then report them against the baseline. + python benchmark_ais_optimization.py screen --kind v5 --subset primary training_extra \\ + --configs configs/ais_control_registry_defaults.json configs/ais_s0_*.json --name s0_screen + python benchmark_ais_optimization.py report --index /ais/screens/_s0_screen.json + + # Sweep a parameter grid on the cache, one shard of the grid per task. + python benchmark_ais_optimization.py sweep --kind v5 --subset primary --grid configs/ais_grid_lm.json \\ + --datasets livecell --shard-index 0 --num-shards 4 + +The configuration file has this shape (a flat parameter dict is read as sparse overrides): + { + "name": "travel-200", + "mode": "auto", + "params_2d": {"sparse": {"n_iter": 200, "dt": 1.0}, "dense": {"beta": 0.6}}, + "params_3d": {"n_iter": 200} + } +""" + +from __future__ import annotations + +import argparse +import datetime +import glob +import itertools +import json +import platform +import sys +import time +from concurrent import futures +from pathlib import Path +from typing import Any, Dict, List, Optional, Sequence, Tuple + +import numpy as np +import pandas as pd +import torch +import xxhash + +EVALUATION_ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(EVALUATION_ROOT)) + +import common # noqa +from common import ( # noqa + DATASETS_3D_EM, DATASET_SPACING, GT_MIN_SIZE_2D, build_model, checkpoint_checksum, drop_severed_objects, + export_joint_checkpoint, get_joint_checkpoint, predict_unisam2, +) +from parameter_search import ( # noqa + compute_metrics, dense_boundary_and_distances, deduplicate_flow_travel, score_image_dense_cached, + score_image_sparse_cached, +) +from optimization import apg3d_manifest # noqa +from optimization.benchmark_apg_optimization import ( # noqa + DEFAULT_DATA_ROOT, DEFAULT_OUTPUT_ROOT, MANIFEST_SUBSETS, _atomic_write_csv, _atomic_write_json, + _content_checksum, _default_manifest_path, _git_revision, _hardware_identity, _load_2d_sample, + _load_3d_sample, _load_normalized_3d_source, prepare_manifest, +) +from optimization.benchmark_apg_3d import _bootstrap_ci # noqa + +from micro_sam.v2.postprocessing import ( # noqa + _compute_flow_density, default_postprocessing, drop_instances_without_boundary_dip, flow_instance_segmentation, + lower_height_under_seeds, run_multicut, watershed_heightmap, +) +from bioimage_cpp.segmentation import label as connected_components, watershed # noqa + +REPOSITORY_ROOT = EVALUATION_ROOT.parents[2] +CAMPAIGN = "ais" +KINDS = ("v5", "apg3d") +MODES = ("auto", "sparse", "dense") +BALANCED_ROW = "__dataset_balanced__" + +# The keywords of the two post-processing functions, i.e. what a configuration may override. +SPARSE_KEYS = ( + "foreground_threshold", "n_iter", "dt", "sigma", "density_threshold", "min_size", "foreground_weight", + "boundary_magnitude_max", "seed_floor", "contact_weight", "contact_mask_threshold", +) +DENSE_KEYS = ("beta", "density_threshold", "n_iter", "dt", "sigma") +EXPLICIT_OFF = "off" +# Metric columns of a sample row; means and standard deviations are reported per dataset. +METRIC_COLUMNS = ("msa", "cremi", "vi_split", "vi_merge", "adapted_rand", "fg_iou", "fg_area_ratio", "matched_iou") +# Count columns; sums are reported per dataset. +COUNT_COLUMNS = ( + "gt_objects", "predicted_objects", "matched", "unmatched", "severed_objects", "genuine_misses", + "matched_before_min_size", "n_seeds", "gt_with_0_seeds", "gt_with_1_seed", "gt_with_2plus_seeds", + "background_seeds", "seeded_unmatched", "seeded_split", "seeded_merged", "seeded_undersized", + "seeded_oversized", "unseeded_absorbed", "unseeded_missing", "pipeline_mismatch", +) +# The generalization gate of the 2026-09 screens (EXPERIMENTAL_SETUP.md, section 9). +GATE = {"max_down": 2, "max_relative_loss": -0.02, "max_absolute_loss": -0.005, "min_balanced_gain": 0.02} +SWEEP_CACHE_KEYS = { + "sparse": ("foreground_threshold", "sigma", "n_iter", "dt"), + "dense": ("density_threshold", "sigma", "n_iter", "dt"), +} + +IMPLEMENTATION_FILES = ( + Path(__file__), + Path(common.__file__), + EVALUATION_ROOT / "optimization/benchmark_apg_optimization.py", + EVALUATION_ROOT / "parameter_search.py", + REPOSITORY_ROOT / "micro_sam/v2/instance_segmentation.py", + REPOSITORY_ROOT / "micro_sam/v2/postprocessing.py", +) + + +def implementation_checksum() -> str: + """Hash the code that determines the prediction, the post-processing and the scoring.""" + checksum = xxhash.xxh128() + for path in IMPLEMENTATION_FILES: + with open(path, "rb") as f: + for block in iter(lambda: f.read(1024 * 1024), b""): + checksum.update(block) + checksum.update(b"\0") + return checksum.hexdigest() + + +# ---------------------------------------------------------------------------------------------- +# configurations + + +def resolve_postprocessing( + overrides: Optional[Dict[str, Any]], model_type: str, ndim: int = 2, +) -> Dict[str, Dict[str, Any]]: + """The sparse and dense parameters a run uses, with 'overrides' on top of the library defaults. + + A flat dict is read as sparse overrides; the nested form ``{"sparse": {...}, "dense": {...}}`` sets + both. 'ndim' selects the image or volume defaults. The result is what `flow_instance_segmentation` / + `run_multicut` receive, so a run without overrides is exactly the library default and shares its run + directory with an explicit copy of it. + """ + overrides = dict(overrides or {}) + if set(overrides) & {"sparse", "dense"}: + unknown = set(overrides) - {"sparse", "dense"} + if unknown: + raise ValueError(f"A nested configuration may only contain 'sparse' and 'dense', got {sorted(unknown)}.") + sparse, dense = dict(overrides.get("sparse", {})), dict(overrides.get("dense", {})) + else: + sparse, dense = overrides, {} + unknown_sparse, unknown_dense = set(sparse) - set(SPARSE_KEYS), set(dense) - set(DENSE_KEYS) + if unknown_sparse or unknown_dense: + raise ValueError(f"Unknown AIS parameters: sparse={sorted(unknown_sparse)}, dense={sorted(unknown_dense)}.") + + sparse_defaults = default_postprocessing(model_type, "sparse", ndim=ndim) + dense_defaults = default_postprocessing(model_type, "dense", ndim=ndim) + + def normalize(values: Dict[str, Any], defaults: Dict[str, Any], mode: str) -> Dict[str, Any]: + normalized = dict(defaults) + for key, value in values.items(): + # Every post-processing keyword uses None to request its model default. Preserve that + # convention in JSON too; otherwise a sweep row and the same row evaluated through + # `run` can silently execute different pipelines. + if value is None: + continue + if value == EXPLICIT_OFF: + if mode != "sparse" or key != "boundary_magnitude_max": + raise ValueError( + f"The explicit value '{EXPLICIT_OFF}' is only valid for boundary_magnitude_max." + ) + value = float("inf") + normalized[key] = value + return normalized + + return { + "sparse": normalize(sparse, sparse_defaults, "sparse"), + "dense": normalize(dense, dense_defaults, "dense"), + } + + +def load_config(path: Optional[Path], model_type: str) -> Tuple[str, str, Dict[str, Any], Dict[str, Any]]: + """Read one configuration file: its name, mode and the resolved 2d and 3d parameters.""" + if path is None: + config: Dict[str, Any] = {"name": "current-defaults"} + else: + with open(path) as f: + config = json.load(f) + unknown = set(config) - {"name", "mode", "params_2d", "params_3d"} + if unknown: + raise ValueError(f"Unknown configuration fields: {sorted(unknown)}.") + mode = config.get("mode", "auto") + if mode not in MODES: + raise ValueError(f"Unknown mode '{mode}'; expected one of {MODES}.") + name = config.get("name", path.stem if path is not None else "current-defaults") + params_2d = resolve_postprocessing(config.get("params_2d", {}), model_type, ndim=2) + # Without its own overrides a volume takes the image overrides, over the library's volume defaults. + params_3d = resolve_postprocessing(config.get("params_3d", config.get("params_2d", {})), model_type, ndim=3) + return str(name), mode, params_2d, params_3d + + +# ---------------------------------------------------------------------------------------------- +# manifests and samples + + +def load_campaign_manifest(kind: str, subset: str, output_root: Path, data_root: Path, campaign_root: Path) -> Dict: + """The 2d subset manifest (`v5`) or the deep 3d crop manifest (`apg3d`) of one subset, validated.""" + if kind == "v5": + if subset not in MANIFEST_SUBSETS: + raise ValueError(f"Unknown v5 subset '{subset}'; expected one of {MANIFEST_SUBSETS}.") + manifest_path = _default_manifest_path(output_root, "standard", subset) + if not manifest_path.exists(): + raise FileNotFoundError(f"The manifest does not exist and is not rebuilt here: '{manifest_path}'.") + manifest = prepare_manifest(data_root, manifest_path, "standard", subset=subset) + manifest["kind"], manifest["subset"] = kind, subset + return manifest + if kind == "apg3d": + manifest = apg3d_manifest.load_manifest(subset, campaign_root, data_root) + manifest["kind"] = kind + return manifest + raise ValueError(f"Unknown manifest kind '{kind}'; expected one of {KINDS}.") + + +def sample_context(sample: Dict[str, Any], kind: str, mode: str) -> Dict[str, Any]: + """Metric mode, post-processing mode, spacing and border size floor of one sample.""" + ndim = int(sample["ndim"]) + if kind == "apg3d": + metric_mode = sample["metric_mode"] + spacing = tuple(sample["spacing"]) if sample.get("spacing") else None + else: + metric_mode = "dense" if sample["dataset"] in DATASETS_3D_EM else "sparse" + spacing = DATASET_SPACING.get(sample["dataset"]) if ndim == 3 else None + if spacing is not None and tuple(spacing) == (1, 1, 1): + spacing = None + dense = (metric_mode == "dense") if mode == "auto" else (mode == "dense") + return { + "ndim": ndim, + "metric_mode": metric_mode, + "postprocessing_mode": "dense" if dense else "sparse", + "spacing": spacing, + "border_min_size": GT_MIN_SIZE_2D.get(sample["dataset"], 0) if ndim == 2 else 0, + } + + +def sample_file_stem(sample: Dict[str, Any]) -> str: + return sample["sample_id"].replace(":", "_") + + +class SampleLoader: + """Loads raw data and labels of manifest samples, caching the normalized 3d source volume.""" + + def __init__(self, kind: str, data_root: Path) -> None: + self.kind, self.data_root = kind, data_root + self._source_key: Optional[tuple] = None + self._source: Optional[np.ndarray] = None + + def _normalized_source(self, sample: Dict[str, Any]) -> np.ndarray: + key = (sample["raw_path"], tuple(sample["normalization_z_range"])) + if key != self._source_key: + self._source = None + if self.kind == "apg3d": + self._source = apg3d_manifest.load_normalized_source(sample, self.data_root) + else: + self._source = _load_normalized_3d_source(sample, self.data_root) + self._source_key = key + return self._source + + def load(self, sample: Dict[str, Any]) -> Tuple[np.ndarray, np.ndarray, Optional[np.ndarray]]: + """The sample's raw data, connected-component labels and valid mask (None unless partially annotated).""" + if self.kind == "apg3d": + return apg3d_manifest.load_sample(sample, self.data_root, self._normalized_source(sample)) + if int(sample["ndim"]) == 2: + raw, labels = _load_2d_sample(sample, self.data_root) + else: + raw, labels = _load_3d_sample(sample, self.data_root, self._normalized_source(sample)) + return raw, labels, None + + +# ---------------------------------------------------------------------------------------------- +# prediction cache + + +class PredictionCache: + """Decoder predictions of one manifest, one file per sample, below + '/ais/predictions///'.""" + + def __init__(self, output_root: Path, checkpoint_id: str, manifest_checksum: str) -> None: + self.root = output_root / CAMPAIGN / "predictions" / checkpoint_id / manifest_checksum + self.checkpoint_id = checkpoint_id + + def paths(self, sample: Dict[str, Any]) -> Tuple[Path, Path]: + stem = sample_file_stem(sample) + return self.root / f"{stem}.npz", self.root / f"{stem}.json" + + def has(self, sample: Dict[str, Any]) -> bool: + return all(path.exists() for path in self.paths(sample)) + + def load(self, sample: Dict[str, Any]) -> Tuple[np.ndarray, np.ndarray, Optional[np.ndarray], Dict[str, Any]]: + """The cached prediction, labels, valid mask and the prediction record of one sample. + + Every array is read once in full; indexing a compressed archive per row decompresses it again. + """ + array_path, record_path = self.paths(sample) + with np.load(array_path) as data: + prediction = np.ascontiguousarray(data["prediction"], dtype="float32") + labels = np.ascontiguousarray(data["labels"], dtype="uint32") + valid = np.ascontiguousarray(data["valid"], dtype=bool) if "valid" in data.files else None + with open(record_path) as f: + record = json.load(f) + if record.get("checkpoint_checksum") != self.checkpoint_id: + raise RuntimeError(f"Cached prediction '{array_path}' belongs to a different checkpoint.") + if record.get("sample_id") != sample["sample_id"]: + raise RuntimeError(f"Cached prediction '{array_path}' belongs to a different sample.") + if record.get("shape") != list(prediction.shape): + raise RuntimeError(f"Cached prediction '{array_path}' does not match its recorded shape.") + if prediction.shape[0] < 4 or prediction.shape[1:] != labels.shape: + raise RuntimeError( + f"Cached prediction / label shape mismatch for '{sample['sample_id']}': " + f"{prediction.shape} and {labels.shape}." + ) + if valid is not None and valid.shape != labels.shape: + raise RuntimeError(f"Cached validity mask for '{sample['sample_id']}' has the wrong shape.") + return prediction, labels, valid, record + + def store( + self, sample: Dict[str, Any], prediction: np.ndarray, labels: np.ndarray, valid: Optional[np.ndarray], + record: Dict[str, Any], + ) -> None: + self.root.mkdir(parents=True, exist_ok=True) + array_path, record_path = self.paths(sample) + arrays = { + "prediction": prediction.astype("float32", copy=False), "labels": labels.astype("uint32", copy=False), + } + if valid is not None: + arrays["valid"] = valid.astype(bool, copy=False) + tmp = array_path.with_suffix(".tmp.npz") + # Uncompressed: a screen reads every file many times, and float32 predictions compress poorly anyway. + np.savez(tmp, **arrays) + tmp.replace(array_path) + _atomic_write_json(record_path, record) + + def records(self, samples: Sequence[Dict[str, Any]]) -> List[Dict[str, Any]]: + records = [] + for sample in samples: + _, record_path = self.paths(sample) + if record_path.exists(): + with open(record_path) as f: + records.append(json.load(f)) + return records + + +class Predictor: + """Builds the UniSAM2 decoder on first use and predicts one sample at a time. + + The decoder half of the joint checkpoint is exported below '/model_exports', keyed by + the checkpoint checksum, like the APG benchmark does (the library's default export root is not + writable for every user). + """ + + def __init__( + self, model_type: str, joint_checkpoint: str, checkpoint_id: str, device: str, output_root: Path, + ) -> None: + self.model_type, self.joint_checkpoint, self.checkpoint_id, self.device = ( + model_type, joint_checkpoint, checkpoint_id, device, + ) + self.export_root = output_root / "model_exports" + self._model = None + + @property + def model(self): + if self._model is None: + _, decoder_path = export_joint_checkpoint( + self.model_type, self.joint_checkpoint, source_checksum=self.checkpoint_id, + export_root=str(self.export_root), + ) + self._model = build_model( + mode="ais", model_type=self.model_type, device=self.device, ndim=2, checkpoint_path=decoder_path, + ) + return self._model + + def predict(self, raw: np.ndarray, ndim: int) -> Tuple[np.ndarray, Dict[str, Any]]: + cuda_device = torch.device(self.device) if self.device.startswith("cuda") else None + if cuda_device is not None: + torch.cuda.reset_peak_memory_stats(cuda_device) + started = time.perf_counter() + prediction = predict_unisam2(self.model, raw, ndim=ndim, device=self.device) + seconds = time.perf_counter() - started + record = { + "predict_seconds": seconds, + "peak_cuda_memory_bytes": int(torch.cuda.max_memory_allocated(cuda_device)) if cuda_device else None, + "device": self.device, + "hardware": _hardware_identity(self.device), + "checkpoint_checksum": self.checkpoint_id, + "checkpoint_name": self.joint_checkpoint, + "model_type": self.model_type, + "implementation_checksum": implementation_checksum(), + "git_revision": _git_revision(), + "torch": torch.__version__, + "shape": list(prediction.shape), + "created": datetime.datetime.now().isoformat(timespec="seconds"), + } + return np.ascontiguousarray(prediction, dtype="float32"), record + + +def ensure_prediction( + cache: PredictionCache, sample: Dict[str, Any], loader: SampleLoader, predictor: Optional[Predictor], +) -> Tuple[np.ndarray, np.ndarray, Optional[np.ndarray], Dict[str, Any]]: + """Read a sample from the cache, predicting and caching it first when it is missing.""" + if cache.has(sample): + return cache.load(sample) + if predictor is None: + raise FileNotFoundError( + f"No cached prediction for '{sample['sample_id']}' under '{cache.root}'. Run 'predict' first, or pass " + "--predict-missing." + ) + raw, labels, valid = loader.load(sample) + prediction, record = predictor.predict(raw, int(sample["ndim"])) + record["sample_id"] = sample["sample_id"] + cache.store(sample, prediction, labels, valid, record) + return prediction, labels, valid, record + + +# ---------------------------------------------------------------------------------------------- +# post-processing, mirrored pipeline and diagnostics + + +def segment_prediction( + prediction: np.ndarray, params: Dict[str, Any], dense: bool, spacing: Optional[tuple], model_type: str, + n_threads: int, +) -> np.ndarray: + """Post-process one prediction exactly like `common.postprocess_unisam2` does in production.""" + if dense: + boundary_map, distances = dense_boundary_and_distances(prediction) + if boundary_map.ndim == 2: + seg = run_multicut( + boundary_map[None], distances[:, None], model_type=model_type, n_threads=n_threads, **params, + )[0] + else: + seg = run_multicut(boundary_map, distances, model_type=model_type, n_threads=n_threads, **params) + else: + contact = {"contact": prediction[4]} if prediction.shape[0] > 4 else {} + seg = flow_instance_segmentation( + prediction[0], prediction[1:4], model_type=model_type, spacing=spacing, n_threads=n_threads, **contact, + **params, + ) + return seg.astype("uint32") + + +def sparse_pipeline( + prediction: np.ndarray, params: Dict[str, Any], spacing: Optional[tuple], n_threads: int, +) -> Dict[str, np.ndarray]: + """`flow_instance_segmentation` step by step, keeping the intermediates the diagnostics read. + + 'params' must be fully resolved (see `resolve_postprocessing`). The segmentation must equal the + library's; `score_sample` records a mismatch per sample, which is the bit-identity check of an epoch. + """ + foreground, directed = prediction[0], prediction[1:4] + contact = prediction[4] if prediction.shape[0] > 4 else None + ndim = foreground.ndim + if directed.shape[0] > ndim: + directed = directed[-ndim:] + fg_mask = foreground > params["foreground_threshold"] + density = _compute_flow_density( + directed, fg_mask, n_iter=int(params["n_iter"]), dt=params["dt"], sigma=params["sigma"], spacing=spacing, + n_threads=n_threads, + ) + seeds = connected_components(density > params["density_threshold"]) + hmap = watershed_heightmap(foreground, directed, params["foreground_weight"]) + contact_weight = params.get("contact_weight") + if contact is not None and contact_weight is not None and contact_weight != 0: + hmap = np.ascontiguousarray(hmap + np.float32(contact_weight) * np.clip(contact, 0, 1), dtype="float32") + hmap = lower_height_under_seeds(hmap, seeds, params.get("seed_floor", "none")) + contact_mask_threshold = params.get("contact_mask_threshold") + if contact is not None and contact_mask_threshold is not None: + open_mask = fg_mask & ~(contact > contact_mask_threshold) + first = watershed(hmap, markers=np.where(open_mask, seeds, 0).astype(seeds.dtype), mask=open_mask) + before = watershed(hmap, markers=first, mask=fg_mask) + else: + before = watershed(hmap, markers=seeds, mask=fg_mask) + seg = before + min_size = int(params["min_size"]) + if min_size > 0: + ids, sizes = np.unique(before, return_counts=True) + discard = ids[(sizes < min_size) & (ids > 0)] + seg = before.copy() + seg[np.isin(seg, discard)] = 0 + seg = watershed(hmap, markers=seg, mask=fg_mask) + max_median = params.get("boundary_magnitude_max") + if max_median is not None and np.isfinite(max_median): + seg = drop_instances_without_boundary_dip(seg, directed, max_median) + return { + "segmentation": seg.astype("uint32"), "before_min_size": before.astype("uint32"), "seeds": seeds, + "fg_mask": fg_mask, "density": density, "heightmap": hmap, + } + + +def contingency(a: np.ndarray, b: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + """Overlap counts of every (a, b) label pair with at least one non-zero label.""" + mask = (a != 0) | (b != 0) + av, bv = a[mask].astype("int64", copy=False), b[mask].astype("int64", copy=False) + if av.size == 0: + empty = np.array([], dtype="int64") + return empty, empty, empty + stride = int(bv.max()) + 1 + keys, counts = np.unique(av * stride + bv, return_counts=True) + return keys // stride, keys % stride, counts + + +def matched_ids(labels: np.ndarray, segmentation: np.ndarray, iou_threshold: float = 0.5) -> np.ndarray: + """The ground-truth ids that some predicted instance matches at the IoU threshold. + + At a threshold of 0.5 or more at most one instance can match an object, and it is the instance with + the largest overlap, so this equals the complement of `common.unmatched_objects`. + """ + gt, seg, inter = contingency(labels, segmentation) + keep = (gt != 0) & (seg != 0) + gt, seg, inter = gt[keep], seg[keep], inter[keep] + if gt.size == 0: + return np.array([], dtype="int64") + gt_sizes = np.bincount(labels.ravel().astype("int64")) + seg_sizes = np.bincount(segmentation.ravel().astype("int64")) + iou = inter / (gt_sizes[gt] + seg_sizes[seg] - inter) + return np.unique(gt[iou >= iou_threshold]) + + +def object_counts(labels: np.ndarray, segmentation: np.ndarray, max_span: int = 2) -> Dict[str, int]: + """Ground-truth object counts of one sample: all, crop-severed (volumes), matched, unmatched, genuine misses. + + The same numbers as `benchmark_apg_3d.object_counts`, computed from one contingency table instead of + one pass over the volume per object. Severed objects (those spanning at most 'max_span' slices) are + only defined for a volume; an image reports 0 severed objects and every miss as genuine. 'matched' + counts matched objects among the unsevered ones (the reference calls it 'merged'). + """ + gt_ids = np.unique(labels) + gt_ids = gt_ids[gt_ids != 0] + matched = matched_ids(labels, segmentation) + if labels.ndim == 3: + spans = np.zeros(int(labels.max()) + 1, dtype="int64") + for plane in labels: + spans[np.unique(plane)] += 1 + severed = gt_ids[spans[gt_ids] <= max_span] + else: + severed = np.array([], dtype=gt_ids.dtype) + unmatched = np.setdiff1d(gt_ids, matched, assume_unique=True) + # Like the reference, 'matched' counts the objects the crop did not sever, so that + # gt_objects = severed_objects + matched + genuine_misses. + return { + "gt_objects": int(len(gt_ids)), + "severed_objects": int(len(severed)), + "matched": int(len(np.setdiff1d(matched, severed, assume_unique=True))), + "unmatched": int(len(unmatched)), + "genuine_misses": int((~np.isin(unmatched, severed)).sum()), + "predicted_objects": int(len(np.unique(segmentation)) - 1), + } + + +def object_fates(labels: np.ndarray, segmentation: np.ndarray) -> Dict[str, np.ndarray]: + """What became of every ground-truth object: its majority instance, the IoU with it, and flags. + + Returns arrays over the ground-truth ids ('ids'): 'iou' (with the instance overlapping most of the + object, 0 without any), 'absorbed' (that instance covers at least half of the object), 'merged' (that + instance covers at least half of two or more objects) and 'undersized' (the instance is smaller than + the object). + """ + ids = np.unique(labels) + ids = ids[ids != 0] + gt, seg, inter = contingency(labels, segmentation) + keep = (gt != 0) & (seg != 0) + gt, seg, inter = gt[keep], seg[keep], inter[keep] + n = int(labels.max()) + 1 + iou = np.zeros(n, dtype="float64") + absorbed = np.zeros(n, dtype=bool) + merged = np.zeros(n, dtype=bool) + undersized = np.zeros(n, dtype=bool) + if gt.size: + gt_sizes = np.bincount(labels.ravel().astype("int64"), minlength=n) + seg_sizes = np.bincount(segmentation.ravel().astype("int64")) + order = np.lexsort((-inter, gt)) + first = np.ones(len(order), dtype=bool) + first[1:] = gt[order][1:] != gt[order][:-1] + major_gt, major_seg, major_inter = gt[order][first], seg[order][first], inter[order][first] + iou[major_gt] = major_inter / (gt_sizes[major_gt] + seg_sizes[major_seg] - major_inter) + strong = major_inter >= 0.5 * gt_sizes[major_gt] + absorbed[major_gt] = strong + claims = np.bincount(major_seg[strong], minlength=len(seg_sizes)) + merged[major_gt] = strong & (claims[major_seg] >= 2) + undersized[major_gt] = seg_sizes[major_seg] < gt_sizes[major_gt] + return { + "ids": ids, "iou": iou[ids], "absorbed": absorbed[ids], "merged": merged[ids], "undersized": undersized[ids], + } + + +def seed_diagnostics( + intermediates: Dict[str, np.ndarray], labels: np.ndarray, segmentation: np.ndarray, +) -> Dict[str, Any]: + """Where the sparse pipeline loses objects: the seeds, the size filter, or the assignment. + + Per ground-truth object the number of seed components inside it (0 = a miss before any + assignment, 2+ = a split), seeds whose majority pixel is background, objects matched before the + size filter, the IoU of the thresholded foreground with the ground-truth foreground and its area ratio + ('fg_area_ratio', the extent calibration), and the fate of + the objects the result lost (IoU below 0.5): seeded ones are 'split' (two or more seeds), 'merged' + (their instance also covers another object), 'undersized' or 'oversized' (an extent error); + unseeded ones are 'absorbed' (mostly covered by a neighbour's instance) or 'missing'. 'matched_iou' + is the mean IoU of the matched objects, a boundary-precision figure. + """ + matched = matched_ids(labels, segmentation) + fates = object_fates(labels, segmentation) + seeds = intermediates["seeds"] + seed_ids, gt_ids, counts = contingency(seeds, labels) + n_seeds = int(seeds.max()) + seeds_per_object = np.zeros(int(labels.max()) + 1, dtype="int64") + inside = (seed_ids != 0) & (gt_ids != 0) + np.add.at(seeds_per_object, gt_ids[inside], 1) + gt_present = np.unique(labels) + gt_present = gt_present[gt_present != 0] + per_object = seeds_per_object[gt_present] + # A seed belongs to the label most of its pixels fall on; background seeds are false starts. + background_seeds = 0 + if n_seeds > 0: + order = np.lexsort((-counts, seed_ids)) + first = np.ones(len(order), dtype=bool) + first[1:] = seed_ids[order][1:] != seed_ids[order][:-1] + majority_label = gt_ids[order][first] + majority_seed = seed_ids[order][first] + background_seeds = int(((majority_label == 0) & (majority_seed != 0)).sum()) + is_matched = np.isin(gt_present, matched) + seeded, lost = per_object >= 1, ~is_matched + seeded_lost = seeded & lost + split = seeded_lost & (per_object >= 2) + merged = seeded_lost & ~split & fates["merged"] + extent = seeded_lost & ~split & ~merged + fg_mask, gt_fg = intermediates["fg_mask"], labels != 0 + union = int((fg_mask | gt_fg).sum()) + gt_area = int(gt_fg.sum()) + return { + "n_seeds": n_seeds, + "gt_with_0_seeds": int((per_object == 0).sum()), + "gt_with_1_seed": int((per_object == 1).sum()), + "gt_with_2plus_seeds": int((per_object >= 2).sum()), + "background_seeds": background_seeds, + "seeded_unmatched": int(seeded_lost.sum()), + "seeded_split": int(split.sum()), + "seeded_merged": int(merged.sum()), + "seeded_undersized": int((extent & fates["undersized"]).sum()), + "seeded_oversized": int((extent & ~fates["undersized"]).sum()), + "unseeded_absorbed": int((~seeded & lost & fates["absorbed"]).sum()), + "unseeded_missing": int((~seeded & lost & ~fates["absorbed"]).sum()), + "matched_before_min_size": int(len(matched_ids(labels, intermediates["before_min_size"]))), + "fg_iou": float((fg_mask & gt_fg).sum() / union) if union else float("nan"), + "fg_area_ratio": float(fg_mask.sum() / gt_area) if gt_area else float("nan"), + "matched_iou": float(fates["iou"][is_matched].mean()) if is_matched.any() else float("nan"), + } + + +# ---------------------------------------------------------------------------------------------- +# running one configuration + + +def run_identity( + params_2d: Dict[str, Any], params_3d: Dict[str, Any], mode: str, dimensions: Sequence[int], trial_id: str, + device: str, hardware: Dict[str, Any], datasets: Optional[Sequence[str]] = None, +) -> str: + identity = { + "params_2d": params_2d, "params_3d": params_3d, "mode": mode, "dimensions": list(dimensions), + "trial_id": trial_id, "device": device, "hardware": hardware, + } + if datasets: + # A run restricted to some datasets is a different (partial) result, not the manifest's. + identity["datasets"] = sorted(datasets) + return _content_checksum(identity) + + +def _prediction_identity(records: Sequence[Dict[str, Any]]) -> Tuple[str, Dict[str, Any]]: + """The device and hardware the cached predictions were made on ('mixed' where they differ).""" + if not records: + return "cache", {} + devices = sorted({str(record.get("device")) for record in records}) + accelerators = sorted({str((record.get("hardware") or {}).get("accelerator")) for record in records}) + hardware = dict(records[0].get("hardware") or {}) + if len(accelerators) > 1: + hardware["accelerator"] = "mixed:" + "|".join(accelerators) + return devices[0] if len(devices) == 1 else "mixed:" + "|".join(devices), hardware + + +def score_sample( + sample: Dict[str, Any], context: Dict[str, Any], prediction: np.ndarray, labels: np.ndarray, + valid: Optional[np.ndarray], record: Dict[str, Any], params: Dict[str, Any], model_type: str, n_threads: int, + diagnostics: bool, +) -> Dict[str, Any]: + """Post-process one cached prediction and score it; the row of `samples.csv`.""" + dense = context["postprocessing_mode"] == "dense" + active = params["dense" if dense else "sparse"] + started = time.perf_counter() + segmentation = segment_prediction(prediction, active, dense, context["spacing"], model_type, n_threads) + generation_seconds = time.perf_counter() - started + if valid is not None: + segmentation[~valid] = 0 + if context["ndim"] == 2: + # Symmetric with the ground truth, which the loader filtered the same way. + segmentation = drop_severed_objects(segmentation, context["border_min_size"]) + metrics = compute_metrics(segmentation, labels, context["metric_mode"], border_min_size=0) + counts = object_counts(labels, segmentation) + predict_seconds = float(record.get("predict_seconds", float("nan"))) + row = { + "sample_id": sample["sample_id"], + "dataset": sample["dataset"], + "ndim": context["ndim"], + "family": sample.get("family", sample["dataset"]), + "stratum": sample.get("stratum", ""), + "seen_in_training": str(sample.get("seen_in_training", "")), + "metric_mode": context["metric_mode"], + "postprocessing_mode": context["postprocessing_mode"], + "initialization_seconds": predict_seconds, + "generation_seconds": generation_seconds, + "total_seconds": (predict_seconds if np.isfinite(predict_seconds) else 0.0) + generation_seconds, + "peak_cuda_memory_bytes": record.get("peak_cuda_memory_bytes"), + **metrics, + **counts, + } + if diagnostics and not dense: + intermediates = sparse_pipeline(prediction, active, context["spacing"], n_threads) + mirrored = intermediates["segmentation"] + if valid is not None: + mirrored[~valid] = 0 + if context["ndim"] == 2: + mirrored = drop_severed_objects(mirrored, context["border_min_size"]) + row["pipeline_mismatch"] = int(not np.array_equal(mirrored, segmentation)) + row.update(seed_diagnostics(intermediates, labels, segmentation)) + return row + + +def summarize(samples: pd.DataFrame) -> pd.DataFrame: + """Per-dataset means (metrics) and sums (counts, seconds), a balanced row and, for the 3d crop + manifests, family and unseen macros in the style of `benchmark_apg_3d.summarize`.""" + metric_columns = [column for column in METRIC_COLUMNS if column in samples.columns] + count_columns = [column for column in COUNT_COLUMNS if column in samples.columns] + second_columns = ["initialization_seconds", "generation_seconds", "total_seconds"] + rows = [] + for dataset, group in samples.groupby("dataset", sort=True): + row: Dict[str, Any] = { + "dataset": dataset, + "family": group["family"].iloc[0] if "family" in group else dataset, + "seen_in_training": str(group["seen_in_training"].iloc[0]) if "seen_in_training" in group else "", + "n_samples": int(len(group)), + } + for column in second_columns: + row[column] = float(group[column].sum()) + if "peak_cuda_memory_bytes" in group: + values = group["peak_cuda_memory_bytes"].dropna() + row["peak_cuda_memory_bytes"] = int(values.max()) if len(values) else np.nan + for metric in metric_columns: + values = group[metric].dropna().to_numpy(dtype="float64") + row[f"{metric}_mean"] = float(values.mean()) if len(values) else np.nan + row[f"{metric}_std"] = float(values.std(ddof=0)) if len(values) else np.nan + if "msa" in group: + row["msa_ci_low"], row["msa_ci_high"] = _bootstrap_ci(group["msa"].dropna().to_numpy()) + for column in count_columns: + values = group[column].dropna() + row[column] = int(values.sum()) if len(values) else np.nan + rows.append(row) + summary = pd.DataFrame(rows) + + def macro(name: str, selected: pd.DataFrame, by: str) -> Dict[str, Any]: + row = {"dataset": name, "n_samples": int(selected["n_samples"].sum()) if len(selected) else 0} + if selected.empty: + return row + groups = selected.groupby(by) + for metric in metric_columns: + means = groups[f"{metric}_mean"].mean().dropna() + row[f"{metric}_mean"] = float(means.mean()) if len(means) else np.nan + row[f"{metric}_std"] = float(means.std(ddof=0)) if len(means) else np.nan + row["n_groups"] = int(groups.ngroups) + for column in second_columns + [c for c in count_columns if c in selected]: + row[column] = selected[column].sum() + if "peak_cuda_memory_bytes" in selected: + values = selected["peak_cuda_memory_bytes"].dropna() + row["peak_cuda_memory_bytes"] = int(values.max()) if len(values) else np.nan + return row + + macros = [macro(BALANCED_ROW, summary, "dataset")] + if (summary["family"] != summary["dataset"]).any(): + macros.append(macro("__family_macro__", summary, "family")) + macros.append(macro("__unseen_macro__", summary[summary["seen_in_training"] == "False"], "family")) + return pd.concat([summary, pd.DataFrame(macros)], ignore_index=True) + + +def run_config( + manifest: Dict[str, Any], output_root: Path, data_root: Path, model_type: str, joint_checkpoint: str, + checkpoint_id: str, config_name: str, mode: str, params_2d: Dict[str, Any], params_3d: Dict[str, Any], + dimensions: Sequence[int], trial_id: str, workers: int, n_threads: int, diagnostics: bool, + predictor: Optional[Predictor] = None, datasets: Optional[Sequence[str]] = None, force: bool = False, +) -> Tuple[Path, pd.DataFrame, Dict[str, Any]]: + """Run one configuration on the cached predictions of a manifest and write its run directory.""" + cache = PredictionCache(output_root, checkpoint_id, manifest["manifest_checksum"]) + loader = SampleLoader(manifest["kind"], data_root) + samples = [sample for sample in manifest["samples"] if int(sample["ndim"]) in dimensions] + if datasets: + samples = [sample for sample in samples if sample["dataset"] in datasets] + if not samples: + raise ValueError("No samples selected.") + for sample in samples: + if not cache.has(sample) and predictor is None: + raise FileNotFoundError( + f"No cached prediction for '{sample['sample_id']}' under '{cache.root}'. Run 'predict' first, or pass " + "--predict-missing." + ) + device, hardware = _prediction_identity(cache.records(samples)) + config_checksum = run_identity(params_2d, params_3d, mode, dimensions, trial_id, device, hardware, datasets) + epoch = implementation_checksum() + run_dir = output_root / CAMPAIGN / model_type / checkpoint_id / ( + f"{manifest['manifest_checksum']}-{config_checksum}-{epoch}" + ) + samples_path, summary_path = run_dir / "samples.csv", run_dir / "summary.csv" + metadata_path = run_dir / "metadata.json" + run_dir.mkdir(parents=True, exist_ok=True) + if metadata_path.exists() and not force: + with open(metadata_path) as f: + metadata = json.load(f) + if metadata.get("status") == "complete" and samples_path.exists() and summary_path.exists(): + print(f"Completed result already exists at '{run_dir}'.") + return run_dir, pd.read_csv(summary_path), metadata + + completed = pd.read_csv(samples_path) if samples_path.exists() and not force else pd.DataFrame() + done = set(completed["sample_id"]) if not completed.empty else set() + pending = [sample for sample in samples if sample["sample_id"] not in done] + metadata = { + "campaign": CAMPAIGN, + "status": "running", + "config_name": config_name, + "config_checksum": config_checksum, + "mode": mode, + "manifest_kind": manifest["kind"], + "subset": manifest.get("subset"), + "manifest_checksum": manifest["manifest_checksum"], + "implementation_checksum": epoch, + "checkpoint_checksum": checkpoint_id, + "checkpoint_name": joint_checkpoint, + "model_type": model_type, + "params_2d": params_2d, + "params_3d": params_3d, + "dimensions": list(dimensions), + "datasets": sorted({sample["dataset"] for sample in samples}), + "trial_id": trial_id, + # The device and hardware of the cached predictions: the identity the comparator pairs runs by. + "device": device, + "hardware": hardware, + "postprocessing_hardware": _hardware_identity("cpu"), + "workers": workers, + "n_threads": n_threads, + "diagnostics": diagnostics, + "prediction_cache": str(cache.root), + "platform": platform.platform(), + "python": sys.version, + "torch": torch.__version__, + "git_revision": _git_revision(), + } + _atomic_write_json(metadata_path, metadata) + params_by_dimension = {2: params_2d, 3: params_3d} + started = time.perf_counter() + + def process(sample: Dict[str, Any]) -> Dict[str, Any]: + context = sample_context(sample, manifest["kind"], mode) + prediction, labels, valid, record = ensure_prediction(cache, sample, loader, predictor) + row = score_sample( + sample, context, prediction, labels, valid, record, params_by_dimension[context["ndim"]], model_type, + n_threads, diagnostics, + ) + row["trial_id"] = trial_id + return row + + rows: List[Dict[str, Any]] = [] + + def flush() -> None: + nonlocal completed, rows + if rows: + completed = pd.concat([completed, pd.DataFrame(rows)], ignore_index=True) + rows = [] + _atomic_write_csv(samples_path, completed) + + try: + flush_every = 1 if any(int(s["ndim"]) == 3 for s in pending) else 20 + if workers <= 1 or predictor is not None: + # Prediction needs the GPU and the source cache in one thread; a plain loop keeps it simple. + results = map(process, pending) + pool = None + else: + pool = futures.ThreadPoolExecutor(workers) + results = pool.map(process, pending) + try: + for index, row in enumerate(results, start=1): + rows.append(row) + if len(rows) >= flush_every: + flush() + print(f"{config_name}: {index}/{len(pending)} samples, {time.perf_counter() - started:.0f} s") + finally: + if pool is not None: + pool.shutdown() + flush() + expected = {sample["sample_id"] for sample in samples} + if set(completed["sample_id"]) != expected: + raise RuntimeError(f"Run finished with {len(completed)} of {len(expected)} samples.") + summary = summarize(completed) + _atomic_write_csv(summary_path, summary) + metadata.update({ + "status": "complete", "wall_seconds": time.perf_counter() - started, "n_samples": int(len(completed)), + }) + _atomic_write_json(metadata_path, metadata) + except Exception as error: + metadata.update({"status": "failed", "error": f"{type(error).__name__}: {error}"}) + _atomic_write_json(metadata_path, metadata) + raise + return run_dir, summary, metadata + + +# ---------------------------------------------------------------------------------------------- +# reports and the generalization gate + + +def dataset_scores(samples: pd.DataFrame) -> pd.Series: + """Per-dataset quality: mean mSA, or the mean CREMI score (negated, so higher is better) on dense data.""" + scores = {} + for dataset, group in samples.groupby("dataset"): + dense = "metric_mode" in group and group["metric_mode"].iloc[0] == "dense" + if dense and "cremi" in group and group["cremi"].notna().any(): + scores[dataset] = -float(group["cremi"].mean()) + else: + scores[dataset] = float(group["msa"].mean()) + return pd.Series(scores).sort_index() + + +def gate_table(baseline: pd.Series, candidate: pd.Series, gate: Dict[str, float] = GATE) -> Dict[str, Any]: + """The generalization gate: up on all but 'max_down' datasets, no dataset below both loss limits, + balanced gain at least 'min_balanced_gain'. 'baseline' and 'candidate' are per-dataset scores.""" + datasets = sorted(set(baseline.index) & set(candidate.index)) + base = baseline[datasets].to_numpy(dtype="float64") + cand = candidate[datasets].to_numpy(dtype="float64") + with np.errstate(divide="ignore", invalid="ignore"): + relative = np.where(base != 0, cand / np.where(base != 0, base, 1.0) - 1.0, np.nan) + absolute = cand - base + up = int((absolute > 0).sum()) + violates = (relative < gate["max_relative_loss"]) & (absolute < gate["max_absolute_loss"]) + balanced_gain = float(cand.mean() / base.mean() - 1.0) if base.mean() else float("nan") + checks = { + "up_on_all_but_two": bool(up >= len(datasets) - gate["max_down"]), + "no_dataset_below_loss_limits": bool(not violates.any()), + "balanced_gain_at_least_2_percent": bool(balanced_gain >= gate["min_balanced_gain"]), + } + return { + "datasets": datasets, "n_up": up, "n_datasets": len(datasets), + "relative": dict(zip(datasets, relative.tolist())), + "balanced_baseline": float(base.mean()), "balanced_candidate": float(cand.mean()), + "balanced_gain": balanced_gain, + "worst_relative": ( + float(np.nanmin(relative)) if len(relative) and np.isfinite(relative).any() else float("nan") + ), + "checks": checks, "passed": bool(all(checks.values())), + } + + +def load_run(run_dir: Path) -> Tuple[Dict[str, Any], pd.DataFrame]: + with open(run_dir / "metadata.json") as f: + metadata = json.load(f) + if metadata.get("status") != "complete": + raise RuntimeError(f"Run is not complete: '{run_dir}'.") + return metadata, pd.read_csv(run_dir / "samples.csv") + + +def report( + run_dirs_by_config: Dict[str, List[Path]], baseline_name: str, ndim: Optional[int] = None, + datasets: Optional[Sequence[str]] = None, +) -> Tuple[pd.DataFrame, pd.DataFrame]: + """Join the sample tables of every configuration over its subsets and compare with the baseline. + + 'ndim' and 'datasets' restrict the samples (the 2d screens read the eleven image datasets; a + manifest's single volumes are too few to compare). Returns the per-configuration table (balanced + score, gain, gate verdict, count sums) and the per-(configuration, dataset) table of relative changes. + """ + joined: Dict[str, pd.DataFrame] = {} + for name, run_dirs in run_dirs_by_config.items(): + samples = pd.concat([load_run(run_dir)[1] for run_dir in run_dirs], ignore_index=True) + if ndim is not None: + samples = samples[samples["ndim"] == ndim] + if datasets: + samples = samples[samples["dataset"].isin(datasets)] + joined[name] = samples.reset_index(drop=True) + if baseline_name not in joined: + raise ValueError(f"Baseline '{baseline_name}' is not among the configurations {sorted(joined)}.") + baseline_scores = dataset_scores(joined[baseline_name]) + baseline_counts = joined[baseline_name][[c for c in COUNT_COLUMNS if c in joined[baseline_name]]].sum() + rows, details = [], [] + for name, samples in joined.items(): + scores = dataset_scores(samples) + verdict = gate_table(baseline_scores, scores) + counts = samples[[c for c in COUNT_COLUMNS if c in samples]].sum() + row = { + "config": name, "n_samples": int(len(samples)), "balanced": verdict["balanced_candidate"], + "balanced_gain": verdict["balanced_gain"], "n_up": verdict["n_up"], "n_datasets": verdict["n_datasets"], + "worst_relative": verdict["worst_relative"], "passed": verdict["passed"], + "generation_seconds": float(samples["generation_seconds"].sum()), + } + for column in ("matched", "unmatched", "predicted_objects", "gt_with_0_seeds", "gt_with_2plus_seeds", + "background_seeds", "seeded_unmatched", "seeded_split", "seeded_merged", "seeded_undersized", + "seeded_oversized", "unseeded_absorbed", "unseeded_missing", "pipeline_mismatch"): + if column in counts: + row[column] = int(counts[column]) + row[f"{column}_delta"] = int(counts[column] - baseline_counts.get(column, 0)) + rows.append(row) + for dataset in verdict["datasets"]: + details.append({ + "config": name, "dataset": dataset, "baseline": float(baseline_scores[dataset]), + "candidate": float(scores[dataset]), "relative": verdict["relative"][dataset], + }) + table = pd.DataFrame(rows).sort_values("balanced", ascending=False).reset_index(drop=True) + return table, pd.DataFrame(details) + + +def _format_relative(value: float) -> str: + return "n/a" if value is None or not np.isfinite(value) else f"{100 * value:+.1f}%" + + +def print_report(table: pd.DataFrame, details: pd.DataFrame) -> None: + pivot = details.pivot(index="config", columns="dataset", values="relative").loc[table["config"]] + columns = ["config", "balanced", "balanced_gain", "n_up", "n_datasets", "worst_relative", "passed"] + columns += [c for c in ("matched_delta", "gt_with_0_seeds_delta", "gt_with_2plus_seeds_delta", + "background_seeds_delta", "seeded_split_delta", "seeded_merged_delta", + "seeded_undersized_delta", "seeded_oversized_delta", "pipeline_mismatch") if c in table] + shown = table[columns].copy() + for column in ("balanced_gain", "worst_relative"): + shown[column] = shown[column].map(_format_relative) + shown["balanced"] = shown["balanced"].map(lambda v: f"{v:.4f}") + print(shown.to_string(index=False)) + print() + print("Relative change per dataset:") + print(pivot.map(_format_relative).to_string()) + + +# ---------------------------------------------------------------------------------------------- +# parameter sweeps on the cache + + +def grid_combinations(grid: Dict[str, Any], mode: str) -> List[Dict[str, Any]]: + """Expand a Cartesian grid, explicit candidates, or a shared grid with named mechanism families.""" + if set(grid) == {"shared", "families"}: + shared, families = grid["shared"], grid["families"] + if not isinstance(shared, dict) or not isinstance(families, dict) or not families: + raise TypeError("A family grid needs 'shared' and a non-empty 'families' parameter mapping.") + combinations = [] + for family, overrides in families.items(): + if not isinstance(family, str) or not family or not isinstance(overrides, dict): + raise TypeError("Every grid family needs a non-empty string name and a parameter mapping.") + overlap = set(shared) & set(overrides) + if overlap: + raise ValueError(f"Grid family '{family}' redefines shared parameters: {sorted(overlap)}.") + combinations.extend( + {"mechanism_family": family, **combo} + for combo in grid_combinations({**shared, **overrides}, mode) + ) + return combinations + + if set(grid) == {"combinations"}: + combinations = grid["combinations"] + if not isinstance(combinations, list) or not all(isinstance(combo, dict) for combo in combinations): + raise TypeError("An explicit grid needs a list of parameter dictionaries in 'combinations'.") + if not combinations: + raise ValueError("An explicit grid needs at least one parameter combination.") + allowed = SPARSE_KEYS if mode == "sparse" else DENSE_KEYS + unknown = set().union(*(set(combo) for combo in combinations)) - set(allowed) + if unknown: + raise ValueError(f"Unknown {mode} grid parameters: {sorted(unknown)}.") + unique = {json.dumps(combo, sort_keys=True): dict(combo) for combo in combinations} + combinations = list(unique.values()) + return deduplicate_flow_travel(combinations) if mode == "sparse" else combinations + + keys = list(grid) + allowed = SPARSE_KEYS if mode == "sparse" else DENSE_KEYS + unknown = set(keys) - set(allowed) + if unknown: + raise ValueError(f"Unknown {mode} grid parameters: {sorted(unknown)}.") + combinations = [dict(zip(keys, combo)) for combo in itertools.product(*[grid[key] for key in keys])] + if mode == "sparse": + combinations = deduplicate_flow_travel(combinations) + return combinations + + +def sweep_dir( + output_root: Path, checkpoint_id: str, manifest_checksum: str, grid_name: str, grid: Dict[str, Any], +) -> Path: + identity = f"{grid_name}-{_content_checksum(grid)[:12]}-{implementation_checksum()[:12]}" + return output_root / CAMPAIGN / "sweeps" / checkpoint_id / manifest_checksum / identity + + +def shard_combinations( + combinations: Sequence[Dict[str, Any]], mode: str, shard_index: int, num_shards: int, +) -> List[Dict[str, Any]]: + """Partition without splitting an expensive cached flow/oversegmentation group across shards.""" + if num_shards < 1 or not 0 <= shard_index < num_shards: + raise ValueError(f"Invalid shard {shard_index} of {num_shards}.") + if num_shards == 1: + return list(combinations) + keys = SWEEP_CACHE_KEYS[mode] + + def identity(combo: Dict[str, Any]) -> str: + return json.dumps([combo[key] for key in keys], separators=(",", ":")) + + groups = {identity(combo): combo for combo in combinations} + if num_shards > len(groups): + raise ValueError( + f"Requested {num_shards} shards for only {len(groups)} distinct {mode} cache groups." + ) + # Flow integration cost is approximately linear in n_iter. Greedy longest-first assignment avoids + # round-robin shards made entirely of the 1,600-step groups, which otherwise leave most processes on + # a packed CPU node idle while a small slow tail finishes. + loads = [0] * num_shards + group_counts = [0] * num_shards + assignment = {} + ordered = sorted(groups.items(), key=lambda item: (-int(item[1].get("n_iter", 1)), item[0])) + for key, combo in ordered: + shard = min(range(num_shards), key=lambda index: (loads[index], group_counts[index], index)) + assignment[key] = shard + loads[shard] += int(combo.get("n_iter", 1)) + group_counts[shard] += 1 + return [combo for combo in combinations if assignment[identity(combo)] == shard_index] + + +def sweep_dataset( + manifest: Dict[str, Any], cache: PredictionCache, dataset: str, mode: str, grid: Dict[str, Any], + model_type: str, n_threads: int, shard_index: int, num_shards: int, out_dir: Path, +) -> Path: + """Score every grid combination of one dataset on the cache; writes the `parameter_search` CSV layout.""" + samples = [sample for sample in manifest["samples"] if sample["dataset"] == dataset] + if not samples: + raise ValueError(f"No samples of '{dataset}' in the manifest.") + contexts = [sample_context(sample, manifest["kind"], mode) for sample in samples] + postproc_mode = contexts[0]["postprocessing_mode"] + # The grid keys the sweep did not name stay at the library defaults, and the row records them. + combinations = [] + for candidate in grid_combinations(grid, postproc_mode): + candidate = dict(candidate) + family = candidate.pop("mechanism_family", None) + resolved = resolve_postprocessing( + {postproc_mode: candidate}, model_type, ndim=contexts[0]["ndim"], + )[postproc_mode] + if family is not None: + resolved["mechanism_family"] = family + combinations.append(resolved) + combinations = shard_combinations(combinations, postproc_mode, shard_index, num_shards) + suffix = "" if num_shards <= 1 else f".shard{shard_index}of{num_shards}" + out_path = out_dir / f"{dataset}{suffix}.csv" + if out_path.exists(): + print(f"Sweep result exists: {out_path}") + return out_path + metric_lists: List[List[Dict[str, float]]] = [[] for _ in combinations] + started = time.perf_counter() + for index, (sample, context) in enumerate(zip(samples, contexts), start=1): + prediction, labels, _, _ = cache.load(sample) + # The scorers see no valid mask: invalid voxels are background in the labels, so a prediction there + # costs precision the same way in every combination. + if postproc_mode == "sparse": + scores = score_image_sparse_cached( + prediction, labels, combinations, n_threads=n_threads, spacing=context["spacing"], + border_min_size=context["border_min_size"], + ) + else: + scores = score_image_dense_cached(prediction, labels, combinations, n_threads=n_threads, border_min_size=0) + for metrics, collected in zip(scores, metric_lists): + if metrics is not None: + collected.append(metrics) + print(f"{dataset}: {index}/{len(samples)} samples, {len(combinations)} combinations, " + f"{time.perf_counter() - started:.0f} s") + rows = [] + for combo, per_sample in zip(combinations, metric_lists): + if not per_sample: + continue + row = {**combo, "n_images": len(per_sample)} + for key in per_sample[0]: + values = np.asarray([m[key] for m in per_sample], dtype="float64") + row[f"{key}_mean"], row[f"{key}_std"] = float(values.mean()), float(values.std()) + rows.append(row) + out_dir.mkdir(parents=True, exist_ok=True) + _atomic_write_csv(out_path, pd.DataFrame(rows)) + print(f"Saved {out_path} ({time.perf_counter() - started:.0f} s).") + return out_path + + +def merge_sweep(out_dir: Path, dataset: str, num_shards: int) -> Path: + out_path = out_dir / f"{dataset}.csv" + if num_shards <= 1: + if not out_path.exists(): + raise FileNotFoundError(f"Missing sweep result: {out_path}") + return out_path + paths = [out_dir / f"{dataset}.shard{i}of{num_shards}.csv" for i in range(num_shards)] + missing = [str(p) for p in paths if not p.exists()] + if missing: + raise FileNotFoundError(f"Missing shards: {missing}") + _atomic_write_csv(out_path, pd.concat([pd.read_csv(p) for p in paths], ignore_index=True)) + return out_path + + +def shared_configuration(out_dir: Path, datasets: Sequence[str], criterion: str = "msa") -> pd.DataFrame: + """Rank the combinations every dataset scored by how close they come to each dataset's own optimum. + + Columns: the parameters, per-dataset scores and relative-to-optimum ratios, 'mean_relative' (the + selection criterion), 'min_relative' (the worst dataset) and 'balanced' (the equal-weight mean). + """ + tables = [] + keys: Optional[List[str]] = None + for dataset in datasets: + table = pd.read_csv(out_dir / f"{dataset}.csv") + params = [c for c in table.columns if not c.endswith(("_mean", "_std")) and c != "n_images"] + keys = params if keys is None else keys + column = f"{criterion}_mean" + if criterion == "cremi": + table[column] = -table[column] + tables.append(table[params + [column]].rename(columns={column: dataset})) + merged = tables[0] + for table in tables[1:]: + merged = merged.merge(table, on=keys, how="inner") + for dataset in datasets: + best = merged[dataset].max() + merged[f"{dataset}_relative"] = merged[dataset] / best if best else np.nan + relative = merged[[f"{d}_relative" for d in datasets]] + merged["mean_relative"] = relative.mean(axis=1) + merged["min_relative"] = relative.min(axis=1) + merged["balanced"] = merged[list(datasets)].mean(axis=1) + merged = merged.sort_values(["mean_relative", "min_relative"], ascending=False).reset_index(drop=True) + _atomic_write_csv(out_dir / "shared_config.csv", merged) + return merged + + +# ---------------------------------------------------------------------------------------------- +# oracles: what the seeds, the height map and the foreground each cost + + +ORACLES = ("baseline", "gt_seeds", "gt_seeds_gt_fg", "gt_heightmap", "gt_fg", "gt_seeds_gt_heightmap") + + +def gt_seed_markers(labels: np.ndarray) -> np.ndarray: + """One marker per ground-truth object around its deepest interior point, carrying the object's id. + + The marker is the point's 3-neighbourhood clipped to the object. A single pixel would not do: the + geodesic field's magnitude is zero at the object's centre (the gradient vanishes at its source), so + the inverted-magnitude height map has a one-pixel spike there, and the monotone flooding of + `bioimage_cpp.segmentation.watershed` lets a seed sitting on a spike flood last. + """ + from scipy.ndimage import grey_dilation + from micro_sam.v2.automatic_prompt_generation import interior_points + + points = np.zeros(labels.shape, dtype="uint64") + ids = np.unique(labels) + ids = ids[ids != 0] + for index, point in zip(ids, interior_points(labels)): + points[tuple(int(c) for c in point)] = index + dilated = grey_dilation(points, size=(3,) * labels.ndim) + return np.where(labels.astype("uint64") == dilated, dilated, 0).astype("uint64") + + +def gt_ridge_heightmap(labels: np.ndarray) -> np.ndarray: + """A height map whose only ridges are the ground-truth object boundaries.""" + from skimage.segmentation import find_boundaries + + return np.ascontiguousarray(find_boundaries(labels, mode="inner"), dtype="float32") + + +def _finish_watershed(before: np.ndarray, hmap: np.ndarray, fg_mask: np.ndarray, min_size: int) -> np.ndarray: + """The size filter and refill of `flow_instance_segmentation`, applied to an oracle's watershed.""" + seg = before + if min_size > 0: + ids, sizes = np.unique(before, return_counts=True) + discard = ids[(sizes < min_size) & (ids > 0)] + seg = before.copy() + seg[np.isin(seg, discard)] = 0 + seg = watershed(hmap, markers=seg, mask=fg_mask) + return seg.astype("uint32") + + +def oracle_sample( + sample: Dict[str, Any], context: Dict[str, Any], prediction: np.ndarray, labels: np.ndarray, + valid: Optional[np.ndarray], params: Dict[str, Any], n_threads: int, +) -> Dict[str, Any]: + """Score the sparse pipeline with parts of it replaced by the ground truth. + + 'gt_seeds': ground-truth seeds, predicted height map and foreground (ceiling of any seed logic); + 'gt_heightmap': predicted seeds and foreground, ridges at the ground-truth boundaries (ceiling of + any height-map / assignment logic); 'gt_fg': predicted seeds and height map inside the ground-truth + foreground (ceiling of the foreground); and the two-part combinations. + """ + active = params["sparse"] + intermediates = sparse_pipeline(prediction, active, context["spacing"], n_threads) + fg_pred, hmap_pred, seeds_pred = intermediates["fg_mask"], intermediates["heightmap"], intermediates["seeds"] + gt_fg, gt_markers, gt_hmap = labels != 0, gt_seed_markers(labels), gt_ridge_heightmap(labels) + min_size = int(active["min_size"]) + + def finish(hmap: np.ndarray, markers: np.ndarray, mask: np.ndarray) -> np.ndarray: + return _finish_watershed(watershed(hmap, markers=markers, mask=mask), hmap, mask, min_size) + + variants = { + "baseline": intermediates["segmentation"], + "gt_seeds": finish(hmap_pred, gt_markers, fg_pred), + "gt_seeds_gt_fg": finish(hmap_pred, gt_markers, gt_fg), + "gt_heightmap": finish(gt_hmap, seeds_pred, fg_pred), + "gt_fg": finish(hmap_pred, seeds_pred, gt_fg), + "gt_seeds_gt_heightmap": finish(gt_hmap, gt_markers, fg_pred), + } + row = { + "sample_id": sample["sample_id"], "dataset": sample["dataset"], "ndim": context["ndim"], + "family": sample.get("family", sample["dataset"]), "stratum": sample.get("stratum", ""), + "metric_mode": context["metric_mode"], + "gt_objects": int(len(np.unique(labels)) - 1), + } + for name, segmentation in variants.items(): + segmentation = segmentation.astype("uint32") + if valid is not None: + segmentation[~valid] = 0 + if context["ndim"] == 2: + segmentation = drop_severed_objects(segmentation, context["border_min_size"]) + counts = object_counts(labels, segmentation) + row[f"msa_{name}"] = compute_metrics(segmentation, labels, "sparse", border_min_size=0)["msa"] + row[f"matched_{name}"] = counts["matched"] + row[f"predicted_{name}"] = counts["predicted_objects"] + return row + + +def summarize_oracles(samples: pd.DataFrame) -> pd.DataFrame: + """Per-dataset means of every oracle, their gain over the baseline, and the balanced row.""" + rows = [] + for dataset, group in samples.groupby("dataset", sort=True): + row: Dict[str, Any] = { + "dataset": dataset, "n_samples": int(len(group)), "gt_objects": int(group["gt_objects"].sum()), + } + for name in ORACLES: + row[f"msa_{name}"] = float(group[f"msa_{name}"].mean()) + row[f"matched_{name}"] = int(group[f"matched_{name}"].sum()) + rows.append(row) + summary = pd.DataFrame(rows) + balanced = { + "dataset": BALANCED_ROW, "n_samples": int(summary["n_samples"].sum()), + "gt_objects": int(summary["gt_objects"].sum()), + } + for name in ORACLES: + balanced[f"msa_{name}"] = float(summary[f"msa_{name}"].mean()) + balanced[f"matched_{name}"] = int(summary[f"matched_{name}"].sum()) + summary = pd.concat([summary, pd.DataFrame([balanced])], ignore_index=True) + for name in ORACLES[1:]: + summary[f"gain_{name}"] = summary[f"msa_{name}"] / summary["msa_baseline"] - 1.0 + return summary + + +def cmd_oracle(args: argparse.Namespace) -> None: + checkpoint_id = _checkpoint_identity(args.model_type, args.joint_checkpoint) + name, _, params_2d, params_3d = load_config(args.config, args.model_type) + params_by_dimension = {2: params_2d, 3: params_3d} + for manifest in _manifests(args): + cache = PredictionCache(args.output_root, checkpoint_id, manifest["manifest_checksum"]) + samples = [sample for sample in manifest["samples"] if int(sample["ndim"]) in _dimensions(args)] + if args.datasets: + samples = [sample for sample in samples if sample["dataset"] in args.datasets] + identity = _content_checksum( + {"params_2d": params_2d, "params_3d": params_3d, "datasets": sorted(args.datasets or [])} + ) + out_dir = args.output_root / CAMPAIGN / "oracles" / checkpoint_id / manifest["manifest_checksum"] / ( + f"{name}-{identity[:12]}-{implementation_checksum()[:12]}" + ) + out_dir.mkdir(parents=True, exist_ok=True) + samples_path = out_dir / "samples.csv" + completed = pd.read_csv(samples_path) if samples_path.exists() else pd.DataFrame() + done = set(completed["sample_id"]) if not completed.empty else set() + pending = [sample for sample in samples if sample["sample_id"] not in done] + + def process(sample: Dict[str, Any]) -> Dict[str, Any]: + # The oracles are about the sparse pipeline; every sample runs through it. + context = sample_context(sample, manifest["kind"], "sparse") + prediction, labels, valid, _ = cache.load(sample) + return oracle_sample( + sample, context, prediction, labels, valid, params_by_dimension[context["ndim"]], args.threads, + ) + + started = time.perf_counter() + with futures.ThreadPoolExecutor(max(1, args.workers)) as pool: + rows = [] + for index, row in enumerate(pool.map(process, pending), start=1): + rows.append(row) + if len(rows) >= 20: + completed = pd.concat([completed, pd.DataFrame(rows)], ignore_index=True) + rows = [] + _atomic_write_csv(samples_path, completed) + elapsed = time.perf_counter() - started + print(f"oracle {manifest.get('subset')}: {index}/{len(pending)}, {elapsed:.0f} s") + if rows: + completed = pd.concat([completed, pd.DataFrame(rows)], ignore_index=True) + _atomic_write_csv(samples_path, completed) + summary = summarize_oracles(completed) + _atomic_write_csv(out_dir / "summary.csv", summary) + _atomic_write_json(out_dir / "metadata.json", { + "campaign": CAMPAIGN, "kind": "oracle", "config_name": name, "params_2d": params_2d, + "params_3d": params_3d, "manifest_checksum": manifest["manifest_checksum"], + "subset": manifest.get("subset"), + "checkpoint_checksum": checkpoint_id, "implementation_checksum": implementation_checksum(), + "n_samples": int(len(completed)), "git_revision": _git_revision(), + }) + shown = ["dataset", "n_samples"] + [f"msa_{n}" for n in ORACLES] + print(f"\nOracles on {manifest['kind']}/{manifest.get('subset')}: {out_dir}") + print(summary[shown].to_string(index=False, float_format=lambda v: f"{v:.4f}")) + print(summary[["dataset"] + [f"gain_{n}" for n in ORACLES[1:]]].to_string( + index=False, float_format=lambda v: f"{100 * v:+.1f}%")) + + +# ---------------------------------------------------------------------------------------------- +# commands + + +def _checkpoint_identity(model_type: str, joint_checkpoint: str) -> str: + return checkpoint_checksum(get_joint_checkpoint(model_type, joint_checkpoint)) + + +def _manifests(args: argparse.Namespace) -> List[Dict[str, Any]]: + return [ + load_campaign_manifest(args.kind, subset, args.output_root, args.data_root, args.campaign_root) + for subset in args.subset + ] + + +def _dimensions(args: argparse.Namespace) -> Tuple[int, ...]: + return (2, 3) if args.ndim == "both" else (int(args.ndim),) + + +def cmd_predict(args: argparse.Namespace) -> None: + checkpoint_id = _checkpoint_identity(args.model_type, args.joint_checkpoint) + predictor = Predictor(args.model_type, args.joint_checkpoint, checkpoint_id, args.device, args.output_root) + for manifest in _manifests(args): + cache = PredictionCache(args.output_root, checkpoint_id, manifest["manifest_checksum"]) + loader = SampleLoader(manifest["kind"], args.data_root) + samples = [sample for sample in manifest["samples"] if int(sample["ndim"]) in _dimensions(args)] + if args.datasets: + samples = [sample for sample in samples if sample["dataset"] in args.datasets] + if args.sample_index is not None: + samples = [samples[args.sample_index]] + pending = [sample for sample in samples if args.force or not cache.has(sample)] + print(f"{manifest['kind']}/{manifest.get('subset')}: {len(pending)} of {len(samples)} samples to predict " + f"-> {cache.root}") + started = time.perf_counter() + for index, sample in enumerate(pending, start=1): + raw, labels, valid = loader.load(sample) + prediction, record = predictor.predict(raw, int(sample["ndim"])) + record["sample_id"] = sample["sample_id"] + cache.store(sample, prediction, labels, valid, record) + print(f" {sample['sample_id']:40s} {str(prediction.shape):24s} {record['predict_seconds']:6.2f} s " + f"({index}/{len(pending)}, {time.perf_counter() - started:.0f} s)") + + +def _run_configs(args: argparse.Namespace, config_paths: Sequence[Optional[Path]]) -> Dict[str, Dict[str, str]]: + checkpoint_id = _checkpoint_identity(args.model_type, args.joint_checkpoint) + predictor = None + if args.predict_missing: + predictor = Predictor(args.model_type, args.joint_checkpoint, checkpoint_id, args.device, args.output_root) + index: Dict[str, Dict[str, str]] = {} + for manifest in _manifests(args): + for config_path in config_paths: + name, mode, params_2d, params_3d = load_config(config_path, args.model_type) + run_dir, summary, _ = run_config( + manifest, args.output_root, args.data_root, args.model_type, args.joint_checkpoint, checkpoint_id, + name, mode, params_2d, params_3d, _dimensions(args), args.trial_id, args.workers, args.threads, + not args.no_diagnostics, predictor=predictor, datasets=args.datasets, force=args.force, + ) + index.setdefault(name, {})[str(manifest.get("subset"))] = str(run_dir) + shown = [c for c in ("dataset", "n_samples", "msa_mean", "cremi_mean", "matched", "unmatched", + "gt_with_0_seeds", "gt_with_2plus_seeds", "background_seeds", "pipeline_mismatch", + "generation_seconds") if c in summary] + print(f"\n{name} on {manifest['kind']}/{manifest.get('subset')}: {run_dir}") + print(summary[shown].to_string(index=False)) + return index + + +def cmd_run(args: argparse.Namespace) -> None: + _run_configs(args, [args.config]) + + +def cmd_screen(args: argparse.Namespace) -> None: + config_paths: List[Path] = [] + for pattern in args.configs: + matches = sorted(glob.glob(pattern)) + if not matches: + raise FileNotFoundError(f"No configuration matches '{pattern}'.") + config_paths.extend(Path(match) for match in matches) + index = _run_configs(args, config_paths) + screens = args.output_root / CAMPAIGN / "screens" + screens.mkdir(parents=True, exist_ok=True) + stamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") + index_path = screens / f"{stamp}_{args.name}.json" + _atomic_write_json(index_path, { + "name": args.name, "kind": args.kind, "subsets": list(args.subset), "runs": index, + "implementation_checksum": implementation_checksum(), "created": stamp, + }) + print(f"\nScreen index: {index_path}") + if args.baseline in index: + table, details = report( + {name: [Path(p) for p in runs.values()] for name, runs in index.items()}, args.baseline, + ndim=None if args.ndim == "both" else int(args.ndim), + ) + print_report(table, details) + + +def cmd_report(args: argparse.Namespace) -> None: + runs_by_config: Dict[str, List[Path]] = {} + for index_path in args.index or []: + with open(index_path) as f: + index = json.load(f) + for name, runs in index["runs"].items(): + runs_by_config.setdefault(name, []).extend(Path(p) for p in runs.values()) + for run_dir in args.runs or []: + metadata, _ = load_run(Path(run_dir)) + runs_by_config.setdefault(metadata["config_name"], []).append(Path(run_dir)) + if not runs_by_config: + raise SystemExit("Pass --index and/or --runs.") + table, details = report( + runs_by_config, args.baseline, ndim=None if args.ndim == "both" else int(args.ndim), datasets=args.datasets, + ) + print_report(table, details) + if args.output is not None: + args.output.parent.mkdir(parents=True, exist_ok=True) + _atomic_write_csv(args.output, table) + _atomic_write_csv(args.output.with_name(args.output.stem + "_datasets.csv"), details) + print(f"\nReport: {args.output}") + + +def cmd_sweep(args: argparse.Namespace) -> None: + checkpoint_id = _checkpoint_identity(args.model_type, args.joint_checkpoint) + with open(args.grid) as f: + grid = json.load(f) + grid_name = args.grid.stem + for manifest in _manifests(args): + cache = PredictionCache(args.output_root, checkpoint_id, manifest["manifest_checksum"]) + out_dir = sweep_dir(args.output_root, checkpoint_id, manifest["manifest_checksum"], grid_name, grid) + datasets = args.datasets or sorted({sample["dataset"] for sample in manifest["samples"]}) + if args.merge: + for dataset in datasets: + print(f"Merged: {merge_sweep(out_dir, dataset, args.num_shards)}") + shared = shared_configuration(out_dir, datasets) + print(shared.head(args.top).to_string(index=False)) + continue + for dataset in datasets: + sweep_dataset( + manifest, cache, dataset, args.mode, grid, args.model_type, args.threads, args.shard_index, + args.num_shards, out_dir, + ) + print(f"Sweep directory: {out_dir}") + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + sub = parser.add_subparsers(dest="command", required=True) + + def common_arguments(p: argparse.ArgumentParser) -> None: + p.add_argument("--kind", choices=KINDS, default="v5", help="Manifest family: 2d subsets or deep 3d crops.") + p.add_argument("--subset", nargs="+", default=["primary"]) + p.add_argument("--data-root", type=Path, default=DEFAULT_DATA_ROOT) + p.add_argument("--output-root", type=Path, default=DEFAULT_OUTPUT_ROOT) + p.add_argument("--campaign-root", type=Path, default=apg3d_manifest.CAMPAIGN_ROOT, + help="Where the deep 3d manifests live (--kind apg3d).") + p.add_argument("--model-type", default="hvit_t", choices=common.MODEL_TYPES) + p.add_argument("--joint-checkpoint", default="best") + p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") + p.add_argument("--ndim", choices=("2", "3", "both"), default="both") + p.add_argument("--datasets", nargs="*", default=None, help="Restrict to these datasets.") + + predict = sub.add_parser("predict", help="Cache the decoder predictions of a manifest.") + common_arguments(predict) + predict.add_argument("--sample-index", type=int, default=None) + predict.add_argument("--force", action="store_true", help="Re-predict cached samples.") + + def run_arguments(p: argparse.ArgumentParser) -> None: + common_arguments(p) + p.add_argument("--trial-id", default="trial-1") + p.add_argument("--workers", type=int, default=1, help="Samples post-processed concurrently.") + p.add_argument("--threads", type=int, default=4, help="Threads per post-processing call.") + p.add_argument("--no-diagnostics", action="store_true", help="Skip the mirrored pipeline and seed columns.") + p.add_argument("--predict-missing", action="store_true", help="Predict samples missing from the cache.") + p.add_argument("--force", action="store_true", help="Recompute a finished run.") + + run = sub.add_parser("run", help="Run one configuration on the cache.") + run_arguments(run) + run.add_argument("--config", type=Path, default=None) + + screen = sub.add_parser("screen", help="Run several configurations on the cache and report them.") + run_arguments(screen) + screen.add_argument("--configs", nargs="+", required=True, help="Configuration files or globs.") + screen.add_argument("--name", required=True, help="Names the screen index file.") + screen.add_argument("--baseline", default="current-defaults", help="Configuration name the report compares to.") + + rep = sub.add_parser("report", help="Compare finished runs with a baseline under the generalization gate.") + rep.add_argument("--index", type=Path, nargs="*", default=None, help="Screen index files.") + rep.add_argument("--runs", type=Path, nargs="*", default=None, help="Run directories.") + rep.add_argument("--baseline", default="current-defaults") + rep.add_argument("--ndim", choices=("2", "3", "both"), default="both", help="Restrict to images or volumes.") + rep.add_argument("--datasets", nargs="*", default=None, help="Restrict to these datasets.") + rep.add_argument("--output", type=Path, default=None, help="CSV path for the tables.") + + oracle = sub.add_parser("oracle", help="Score the pipeline with ground-truth seeds, height map or foreground.") + common_arguments(oracle) + oracle.add_argument("--config", type=Path, default=None) + oracle.add_argument("--workers", type=int, default=1) + oracle.add_argument("--threads", type=int, default=4) + + sweep = sub.add_parser("sweep", help="Score a parameter grid on the cache, one dataset at a time.") + common_arguments(sweep) + sweep.add_argument("--grid", type=Path, required=True, help="JSON dict of parameter lists.") + sweep.add_argument("--mode", choices=MODES, default="auto") + sweep.add_argument("--threads", type=int, default=4) + sweep.add_argument("--shard-index", type=int, default=0) + sweep.add_argument("--num-shards", type=int, default=1) + sweep.add_argument("--merge", action="store_true", help="Merge the shards and rank the shared configuration.") + sweep.add_argument("--top", type=int, default=20) + return parser + + +def main(argv: Optional[Sequence[str]] = None) -> int: + parser = build_parser() + args = parser.parse_args(argv) + if hasattr(args, "data_root"): + args.data_root = args.data_root.expanduser().resolve(strict=True) + args.output_root = args.output_root.expanduser().resolve() + if args.output_root == args.data_root or args.data_root in args.output_root.parents: + parser.error("The output root must not be inside the read-only data root.") + commands = { + "predict": cmd_predict, "run": cmd_run, "screen": cmd_screen, "report": cmd_report, "sweep": cmd_sweep, + "oracle": cmd_oracle, + } + commands[args.command](args) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/finetuning/v2/evaluation/optimization/benchmark_apg_3d.py b/finetuning/v2/evaluation/optimization/benchmark_apg_3d.py index aa8323e03..a19e19e48 100644 --- a/finetuning/v2/evaluation/optimization/benchmark_apg_3d.py +++ b/finetuning/v2/evaluation/optimization/benchmark_apg_3d.py @@ -92,22 +92,16 @@ def load_volume_config(path: Optional[Path], model_type: str = "hvit_t") -> Tupl return str(config.get("name", path.stem)), resolve_volume_params(config.get("params_3d", {}), model_type) -def run_identity( - config_name: str, params_3d: Dict[str, Any], checkpoint_id: str, manifest_checksum: str, trial_id: str, -) -> str: - identity = { - "params_3d": params_3d, "checkpoint_checksum": checkpoint_id, - "manifest_checksum": manifest_checksum, "trial_id": trial_id, - } +def run_identity(config_name: str, params_3d: Dict[str, Any]) -> str: + # 'artifacts' is a frozen empty field: the learned artifacts it once recorded are gone, but keeping + # the key leaves the run-directory family of a configuration intact, so the historical crops of the + # campaign still aggregate with the current ones. + identity = {"params_3d": params_3d, "artifacts": {}} return f"{config_name}-{_content_checksum(identity)[:12]}-{_implementation_checksum()[:12]}" -def run_dir( - campaign_root: Path, subset: str, config_name: str, params_3d: Dict[str, Any], - checkpoint_id: str, manifest_checksum: str, trial_id: str, -) -> Path: - identity = run_identity(config_name, params_3d, checkpoint_id, manifest_checksum, trial_id) - return campaign_root / "runs" / subset / identity +def run_dir(campaign_root: Path, subset: str, config_name: str, params_3d: Dict[str, Any]) -> Path: + return campaign_root / "runs" / subset / run_identity(config_name, params_3d) def sibling_run_dirs(run_path: Path) -> List[Path]: @@ -132,10 +126,7 @@ def object_counts(labels: np.ndarray, segmentation: np.ndarray) -> Dict[str, Any severed = set(int(value) for value in severed_ids) genuine = gt_ids - severed unmatched = set(int(value) for value in np.unique(unmatched_objects(labels, segmentation)) if value != 0) - result = { - "gt_objects": len(gt_ids), "severed_objects": len(severed), - "merged": len(gt_ids - unmatched), "non_severed_matches": len(genuine - unmatched), - } + result = {"gt_objects": len(gt_ids), "severed_objects": len(severed), "merged": len(genuine - unmatched)} result["unmatched"], result["genuine_misses"] = genuine_misses(labels, segmentation) return result @@ -144,12 +135,13 @@ def object_counts(labels: np.ndarray, segmentation: np.ndarray) -> Dict[str, Any # running -def _build(model_type: str, joint_checkpoint: str, checkpoint_id: str, device: str, export_root: Path): +def _build(model_type: str, joint_checkpoint: str, device: str, export_root: Path): + checkpoint_id = checkpoint_checksum(get_joint_checkpoint(model_type, joint_checkpoint)) segmenter = build_apg_segmenter( model_type, 3, device, joint_checkpoint=joint_checkpoint, joint_checksum=checkpoint_id, export_root=str(export_root), ) - return segmenter + return segmenter, checkpoint_id def _save_outputs(path: Path, segmentation: np.ndarray) -> None: @@ -235,11 +227,7 @@ def _write_metadata(run_path: Path, manifest: Dict[str, Any], config_name: str, def run(args: argparse.Namespace) -> None: manifest = load_manifest(args.subset, args.campaign_root, args.data_root) config_name, params_3d = load_volume_config(args.config, args.model_type) - checkpoint_id = checkpoint_checksum(get_joint_checkpoint(args.model_type, args.joint_checkpoint)) - run_path = run_dir( - args.campaign_root, args.subset, config_name, params_3d, - checkpoint_id, manifest["manifest_checksum"], args.trial_id, - ) + run_path = run_dir(args.campaign_root, args.subset, config_name, params_3d) samples = manifest["samples"] if args.sample_index is not None: samples = [samples[args.sample_index]] @@ -256,13 +244,13 @@ def run(args: argparse.Namespace) -> None: if not pending: print(f"All {len(samples)} crop(s) already done in {run_path}.") return - segmenter = _build( - args.model_type, args.joint_checkpoint, checkpoint_id, args.device, DEFAULT_OUTPUT_ROOT / "model_exports", + segmenter, checkpoint_id = _build( + args.model_type, args.joint_checkpoint, args.device, DEFAULT_OUTPUT_ROOT / "model_exports", ) if not (run_path / "metadata.json").exists(): _write_metadata( run_path, manifest, config_name, params_3d, args.model_type, args.joint_checkpoint, - checkpoint_id, args.device, status="running", extra={"trial_id": args.trial_id}, + checkpoint_id, args.device, status="running", ) source_cache: Dict[tuple, np.ndarray] = {} started = time.perf_counter() @@ -309,8 +297,7 @@ def summarize(samples: pd.DataFrame) -> pd.DataFrame: rows = [] numeric = [column for column in samples.columns if pd.api.types.is_numeric_dtype(samples[column])] sums = [column for column in numeric if column in ( - "gt_objects", "severed_objects", "merged", "non_severed_matches", "unmatched", "genuine_misses", - "predicted_objects", *STATS_KEYS, + "gt_objects", "severed_objects", "merged", "unmatched", "genuine_misses", "predicted_objects", *STATS_KEYS, )] per_dataset = {} for dataset, group in samples.groupby("dataset", sort=True): @@ -357,11 +344,7 @@ def macro(name: str, selected: pd.DataFrame) -> Dict[str, Any]: def aggregate(args: argparse.Namespace) -> None: manifest = load_manifest(args.subset, args.campaign_root, args.data_root) config_name, params_3d = load_volume_config(args.config, args.model_type) - checkpoint_id = checkpoint_checksum(get_joint_checkpoint(args.model_type, args.joint_checkpoint)) - run_path = run_dir( - args.campaign_root, args.subset, config_name, params_3d, - checkpoint_id, manifest["manifest_checksum"], args.trial_id, - ) + run_path = run_dir(args.campaign_root, args.subset, config_name, params_3d) by_sample: Dict[str, Dict[str, Any]] = {} implementations = [] for sibling in sibling_run_dirs(run_path): @@ -386,8 +369,6 @@ def aggregate(args: argparse.Namespace) -> None: metadata_path = run_path / "metadata.json" metadata = json.load(open(metadata_path)) if metadata_path.exists() else {} metadata.update({ - "checkpoint_checksum": checkpoint_id, "manifest_checksum": manifest["manifest_checksum"], - "trial_id": args.trial_id, "status": "complete" if done == expected else "partial", "n_crops": len(rows), "n_expected": len(expected), "missing": sorted(expected - done), "implementation_checksums": sorted(set(implementations)), diff --git a/finetuning/v2/evaluation/optimization/benchmark_apg_optimization.py b/finetuning/v2/evaluation/optimization/benchmark_apg_optimization.py index f2da32799..2c0faeb2c 100644 --- a/finetuning/v2/evaluation/optimization/benchmark_apg_optimization.py +++ b/finetuning/v2/evaluation/optimization/benchmark_apg_optimization.py @@ -40,6 +40,7 @@ import subprocess import sys import time +import warnings from collections import defaultdict from pathlib import Path from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple @@ -115,7 +116,26 @@ "covid_if": 5, "deepseas": 40, } -MANIFEST_SUBSETS = ("primary", "holdout", "training_extra") +# The sealed 2d-only confirmation set for the Dice-foreground AIS decoder comparison. These domains are +# absent from the decoder fine-tuning manifest. The large heterogeneous datasets are sampled within +# their acquisition/stain strata; the small official test sets are kept in full. +OOD_EXTENDED_DATASETS = ("arvidsson", "bitdepth_nucseg", "cellbindb", "microbeseg", "vicar") +SAMPLE_COUNTS_2D_OOD_EXTENDED = { + "arvidsson": 10, + "bitdepth_nucseg": 70, + "cellbindb": 48, + "microbeseg": 2, + "vicar": 50, +} +OOD_EXTENDED_STRATUM_COUNTS = { + "bitdepth_nucseg": {"20x": 9, "40x air": 19, "40x oil": 20, "63x oil": 22}, + "cellbindb": { + "10×Genomics_DAPI": 8, "10×Genomics_HE": 8, "DAPI": 8, + "HE": 8, "mIF": 8, "ssDNA": 8, + }, + "vicar": {"A2058": 10, "G361": 10, "HOB": 10, "PC3": 10, "PNT1A": 10}, +} +MANIFEST_SUBSETS = ("primary", "holdout", "training_extra", "ood_extended") TARGETS_3D = (0.5,) # Match the 512 x 512 training field of view and use enough depth to contain representative 3d # structure. C. elegans keeps the deeper crop needed to contain its 11-13-slice nuclei; its source @@ -412,16 +432,34 @@ def _center_crop_roi(shape: Sequence[int], crop_shape: Sequence[int]) -> Tuple[s return tuple(roi) -def _scan_2d_dataset(dataset: str, data_root: Path) -> List[Dict[str, Any]]: +def _scan_2d_dataset( + dataset: str, data_root: Path, split: str = "val", validate_raw: bool = False, + skip_read_errors: bool = False, +) -> List[Dict[str, Any]]: raw_paths, label_paths, raw_key, label_key = get_data_paths( - dataset, str(data_root), download=False, split="val" + dataset, str(data_root), download=False, split=split ) candidates = [] pairs = sorted_path_pairs(raw_paths, label_paths) for raw_path, label_path in tqdm(pairs, desc=f"select-{dataset}", leave=False): raw_relative = _relative_data_path(raw_path, data_root) label_relative = _relative_data_path(label_path, data_root) - labels = read_2d(str(_source_path(label_relative, data_root)), label_key) + try: + labels = read_2d(str(_source_path(label_relative, data_root)), label_key) + if validate_raw: + # CellBinDB contains a handful of corrupt files. Validate both halves before a + # sealed sample can enter the manifest, rather than failing much later at inference. + raw = read_2d(str(_source_path(raw_relative, data_root)), raw_key) + if raw.shape[:2] != labels.shape[:2]: + raise RuntimeError(f"raw shape {raw.shape} does not match labels {labels.shape}") + except Exception as error: + if not skip_read_errors: + raise + warnings.warn( + f"Skipping unreadable {dataset} pair '{raw_relative}': {type(error).__name__}: {error}", + stacklevel=2, + ) + continue roi = _center_crop_roi(labels.shape[:2], CROP_SHAPE_2D) labels = connected_components(labels[roi]).astype("uint32") labels = drop_severed_objects(labels, GT_MIN_SIZE_2D.get(dataset, 0)) @@ -440,7 +478,7 @@ def _scan_2d_dataset(dataset: str, data_root: Path) -> List[Dict[str, Any]]: "foreground_fraction": foreground_fraction, }) if not candidates: - raise RuntimeError(f"No non-empty validation images found for '{dataset}'.") + raise RuntimeError(f"No readable non-empty '{split}' images found for '{dataset}'.") return candidates @@ -522,6 +560,74 @@ def _select_2d_samples( return samples +def _ood_stratum(sample: Dict[str, Any]) -> Optional[str]: + """Return the acquisition stratum encoded in an OOD sample's path.""" + parts = Path(sample["raw_path"]).parts + dataset = sample["dataset"] + offsets = {"bitdepth_nucseg": ("data", 1), "cellbindb": ("Other", 1), "vicar": ("labelled", 1)} + if dataset not in offsets: + return None + anchor, offset = offsets[dataset] + try: + return parts[parts.index(anchor) + offset] + except (ValueError, IndexError) as error: + raise RuntimeError(f"Cannot derive the OOD stratum from '{sample['raw_path']}'.") from error + + +def _select_ood_extended_samples(data_root: Path) -> List[Dict[str, Any]]: + """Build the sealed AIS OOD sample list, stratifying heterogeneous sources deterministically.""" + samples = [] + for dataset in OOD_EXTENDED_DATASETS: + candidates = _scan_2d_dataset( + dataset, data_root, split="test", validate_raw=True, skip_read_errors=True, + ) + for candidate in candidates: + stratum = _ood_stratum(candidate) + if stratum is not None: + candidate["stratum"] = stratum + + stratum_counts = OOD_EXTENDED_STRATUM_COUNTS.get(dataset) + if stratum_counts is None: + requested = SAMPLE_COUNTS_2D_OOD_EXTENDED[dataset] + if len(candidates) != requested: + raise RuntimeError( + f"The sealed '{dataset}' pool changed: expected {requested} non-empty images, " + f"found {len(candidates)}. Refuse to silently change the manifest." + ) + _add_complexity(candidates) + selected = sorted( + candidates, key=lambda entry: (entry.get("stratum", ""), entry["raw_path"]), + ) + else: + selected = [] + by_stratum = defaultdict(list) + for candidate in candidates: + by_stratum[candidate["stratum"]].append(candidate) + if set(by_stratum) != set(stratum_counts): + raise RuntimeError( + f"The sealed '{dataset}' strata changed: expected {sorted(stratum_counts)}, " + f"found {sorted(by_stratum)}." + ) + for stratum, requested in stratum_counts.items(): + group = by_stratum[stratum] + if len(group) < requested: + raise RuntimeError( + f"The sealed '{dataset}/{stratum}' pool has {len(group)} images, needs {requested}." + ) + _add_complexity(group) + selected.extend(_select_nearest(group, _quantile_targets(requested))) + + if len(selected) != SAMPLE_COUNTS_2D_OOD_EXTENDED[dataset]: + raise RuntimeError( + f"Selected {len(selected)} '{dataset}' images, expected " + f"{SAMPLE_COUNTS_2D_OOD_EXTENDED[dataset]}." + ) + for sample in selected: + sample["sample_id"] = _sample_identity(sample) + samples.append(sample) + return samples + + def _read_array(path: Path, key: Optional[str], roi: Optional[Tuple[slice, ...]] = None) -> np.ndarray: if key is None: array = np.asarray(common.load_image(str(path))) @@ -686,6 +792,8 @@ def _sample_counts_2d(subset: str) -> Dict[str, int]: return SAMPLE_COUNTS_2D_HOLDOUT if subset == "training_extra": return SAMPLE_COUNTS_2D_TRAINING_EXTRA + if subset == "ood_extended": + return SAMPLE_COUNTS_2D_OOD_EXTENDED return SAMPLE_COUNTS_2D @@ -737,6 +845,37 @@ def _validate_manifest(manifest: Dict[str, Any], data_root: Path, variant: str, if set(counts) != set(expected) or short: raise RuntimeError(f"Unexpected training_extra sample counts: got {dict(counts)}, caps {expected}.") return + if subset == "ood_extended": + expected_policy = { + "subset": "ood_extended", + "role": "sealed-2d-confirmation-only", + "datasets": list(OOD_EXTENDED_DATASETS), + "stratum_counts": OOD_EXTENDED_STRATUM_COUNTS, + "source_split": "test", + "unreadable_source_policy": "validate raw and label; deterministically exclude unreadable pairs", + } + stored_policy = {key: policy.get(key) for key in expected_policy} + if json.loads(_json_bytes(stored_policy)) != json.loads(_json_bytes(expected_policy)): + raise RuntimeError( + f"The ood_extended selection policy changed: got {stored_policy}, expected {expected_policy}." + ) + expected = {(dataset, 2): sample_counts[dataset] for dataset in OOD_EXTENDED_DATASETS} + if dict(counts) != expected: + raise RuntimeError(f"Unexpected ood_extended sample counts: got {dict(counts)}, expected {expected}.") + expected_strata = { + (dataset, stratum): count + for dataset, strata in OOD_EXTENDED_STRATUM_COUNTS.items() + for stratum, count in strata.items() + } + actual_strata = defaultdict(int) + for sample in samples: + if sample["dataset"] in OOD_EXTENDED_STRATUM_COUNTS: + actual_strata[(sample["dataset"], sample.get("stratum"))] += 1 + if dict(actual_strata) != expected_strata: + raise RuntimeError( + f"Unexpected ood_extended strata: got {dict(actual_strata)}, expected {expected_strata}." + ) + return expected = {(dataset, 2): sample_counts[dataset] for dataset in DATASETS_2D} expected.update({(dataset, 3): 1 for dataset in DATASETS_3D}) if dict(counts) != expected: @@ -812,6 +951,16 @@ def prepare_manifest( data_root, counts=SAMPLE_COUNTS_2D_TRAINING_EXTRA, datasets=TRAINING_EXTRA_DATASETS, allow_fewer=True, ) subset_policy = {"subset": "training_extra", "role": "selector-training-only"} + elif subset == "ood_extended": + samples = _select_ood_extended_samples(data_root) + subset_policy = { + "subset": "ood_extended", + "role": "sealed-2d-confirmation-only", + "datasets": list(OOD_EXTENDED_DATASETS), + "stratum_counts": OOD_EXTENDED_STRATUM_COUNTS, + "source_split": "test", + "unreadable_source_policy": "validate raw and label; deterministically exclude unreadable pairs", + } else: samples = _select_2d_samples(data_root) + _select_3d_samples(data_root, variant) @@ -821,7 +970,11 @@ def prepare_manifest( "selection_policy": { "2d_crop_shape": list(CROP_SHAPE_2D), "2d_sample_counts": sample_counts, - "2d_complexity_targets": "even quantile midpoints within each dataset and LIVECell cell type", + "2d_complexity_targets": ( + "even quantile midpoints within each dataset and declared stratum; full small OOD test pools" + if subset == "ood_extended" + else "even quantile midpoints within each dataset and LIVECell cell type" + ), "3d_complexity_targets": list(TARGETS_3D), "complexity": "mean percentile rank of object count and foreground fraction", **subset_policy, diff --git a/finetuning/v2/evaluation/optimization/configs/ais_c1_sigma1_ms50_filter0p4.json b/finetuning/v2/evaluation/optimization/configs/ais_c1_sigma1_ms50_filter0p4.json new file mode 100644 index 000000000..6e77c1f23 --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_c1_sigma1_ms50_filter0p4.json @@ -0,0 +1,9 @@ +{ + "name": "c1-sigma1-ms50-filter0p4", + "mode": "auto", + "params_2d": { + "sigma": 1.0, + "min_size": 50, + "boundary_magnitude_max": 0.4 + } +} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_c1_t400.json b/finetuning/v2/evaluation/optimization/configs/ais_c1_t400.json new file mode 100644 index 000000000..65393a17c --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_c1_t400.json @@ -0,0 +1,11 @@ +{ + "name": "c1-t400", + "mode": "auto", + "params_2d": { + "sigma": 1.0, + "min_size": 50, + "boundary_magnitude_max": 0.4, + "n_iter": 800, + "dt": 0.5 + } +} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_c1v_ms200.json b/finetuning/v2/evaluation/optimization/configs/ais_c1v_ms200.json new file mode 100644 index 000000000..29a600640 --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_c1v_ms200.json @@ -0,0 +1,14 @@ +{ + "name": "c1v-ms200", + "mode": "auto", + "params_2d": { + "sigma": 1.0, + "min_size": 50, + "boundary_magnitude_max": 0.4 + }, + "params_3d": { + "sigma": 1.0, + "min_size": 200, + "boundary_magnitude_max": 0.4 + } +} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_c1v_volume.json b/finetuning/v2/evaluation/optimization/configs/ais_c1v_volume.json new file mode 100644 index 000000000..e1f78c856 --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_c1v_volume.json @@ -0,0 +1,15 @@ +{ + "name": "c1v-volume", + "mode": "auto", + "params_2d": { + "sigma": 1.0, + "min_size": 50, + "boundary_magnitude_max": 0.4 + }, + "params_3d": { + "sigma": 1.0, + "min_size": 200, + "boundary_magnitude_max": 0.4, + "foreground_threshold": 0.6 + } +} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_c3_sigma1_ms50.json b/finetuning/v2/evaluation/optimization/configs/ais_c3_sigma1_ms50.json new file mode 100644 index 000000000..d18eaaded --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_c3_sigma1_ms50.json @@ -0,0 +1,8 @@ +{ + "name": "c3-sigma1-ms50", + "mode": "auto", + "params_2d": { + "sigma": 1.0, + "min_size": 50 + } +} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_contact_mask.json b/finetuning/v2/evaluation/optimization/configs/ais_contact_mask.json new file mode 100644 index 000000000..7918e05c4 --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_contact_mask.json @@ -0,0 +1 @@ +{"name": "contact-mask", "params_2d": {"contact_mask_threshold": 0.5}, "params_3d": {"contact_mask_threshold": 0.5}} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_contact_mask_t0.3.json b/finetuning/v2/evaluation/optimization/configs/ais_contact_mask_t0.3.json new file mode 100644 index 000000000..5d14939f0 --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_contact_mask_t0.3.json @@ -0,0 +1 @@ +{"name": "contact-mask-t0.3", "params_2d": {"contact_mask_threshold": 0.3}, "params_3d": {"contact_mask_threshold": 0.3}} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_contact_mask_t0.7.json b/finetuning/v2/evaluation/optimization/configs/ais_contact_mask_t0.7.json new file mode 100644 index 000000000..5b7070c07 --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_contact_mask_t0.7.json @@ -0,0 +1 @@ +{"name": "contact-mask-t0.7", "params_2d": {"contact_mask_threshold": 0.7}, "params_3d": {"contact_mask_threshold": 0.7}} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_contact_ridge.json b/finetuning/v2/evaluation/optimization/configs/ais_contact_ridge.json new file mode 100644 index 000000000..f16600b52 --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_contact_ridge.json @@ -0,0 +1 @@ +{"name": "contact-ridge", "params_2d": {"contact_weight": 1.0}, "params_3d": {"contact_weight": 1.0}} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_contact_ridge1_mask0.5.json b/finetuning/v2/evaluation/optimization/configs/ais_contact_ridge1_mask0.5.json new file mode 100644 index 000000000..a56376aca --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_contact_ridge1_mask0.5.json @@ -0,0 +1 @@ +{"name": "contact-ridge1-mask0.5", "params_2d": {"contact_weight": 1.0, "contact_mask_threshold": 0.5}, "params_3d": {"contact_weight": 1.0, "contact_mask_threshold": 0.5}} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_contact_ridge_w0.5.json b/finetuning/v2/evaluation/optimization/configs/ais_contact_ridge_w0.5.json new file mode 100644 index 000000000..ff66212ad --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_contact_ridge_w0.5.json @@ -0,0 +1 @@ +{"name": "contact-ridge-w0.5", "params_2d": {"contact_weight": 0.5}, "params_3d": {"contact_weight": 0.5}} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_contact_ridge_w2.0.json b/finetuning/v2/evaluation/optimization/configs/ais_contact_ridge_w2.0.json new file mode 100644 index 000000000..80223dc3a --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_contact_ridge_w2.0.json @@ -0,0 +1 @@ +{"name": "contact-ridge-w2.0", "params_2d": {"contact_weight": 2.0}, "params_3d": {"contact_weight": 2.0}} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_contact_ridge_w4.0.json b/finetuning/v2/evaluation/optimization/configs/ais_contact_ridge_w4.0.json new file mode 100644 index 000000000..2b9d6e3c6 --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_contact_ridge_w4.0.json @@ -0,0 +1 @@ +{"name": "contact-ridge-w4.0", "params_2d": {"contact_weight": 4.0}, "params_3d": {"contact_weight": 4.0}} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_control_registry_defaults.json b/finetuning/v2/evaluation/optimization/configs/ais_control_registry_defaults.json new file mode 100644 index 000000000..3d222841a --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_control_registry_defaults.json @@ -0,0 +1,6 @@ +{ + "name": "current-defaults", + "mode": "auto", + "params_2d": {}, + "params_3d": {} +} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_control_v4_old_defaults.json b/finetuning/v2/evaluation/optimization/configs/ais_control_v4_old_defaults.json new file mode 100644 index 000000000..9c5903790 --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_control_v4_old_defaults.json @@ -0,0 +1,24 @@ +{ + "name": "v4-old-defaults", + "mode": "auto", + "params_2d": { + "foreground_threshold": 0.5, + "density_threshold": 10.0, + "min_size": 100, + "sigma": 0.5, + "n_iter": 50, + "dt": 0.5, + "foreground_weight": 0.5, + "boundary_magnitude_max": Infinity + }, + "params_3d": { + "foreground_threshold": 0.5, + "density_threshold": 10.0, + "min_size": 100, + "sigma": 0.5, + "n_iter": 50, + "dt": 0.5, + "foreground_weight": 0.5, + "boundary_magnitude_max": Infinity + } +} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_d_filter_only.json b/finetuning/v2/evaluation/optimization/configs/ais_d_filter_only.json new file mode 100644 index 000000000..15d35bd07 --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_d_filter_only.json @@ -0,0 +1,24 @@ +{ + "name": "d-filter-only", + "mode": "auto", + "params_2d": { + "foreground_threshold": 0.5, + "density_threshold": 10.0, + "min_size": 100, + "sigma": 0.5, + "n_iter": 50, + "dt": 0.5, + "foreground_weight": 0.5, + "boundary_magnitude_max": 0.4 + }, + "params_3d": { + "foreground_threshold": 0.5, + "density_threshold": 10.0, + "min_size": 100, + "sigma": 0.5, + "n_iter": 50, + "dt": 0.5, + "foreground_weight": 0.5, + "boundary_magnitude_max": 0.4 + } +} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_d_ms50_only.json b/finetuning/v2/evaluation/optimization/configs/ais_d_ms50_only.json new file mode 100644 index 000000000..e6ff987e1 --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_d_ms50_only.json @@ -0,0 +1,24 @@ +{ + "name": "d-ms50-only", + "mode": "auto", + "params_2d": { + "foreground_threshold": 0.5, + "density_threshold": 10.0, + "min_size": 50, + "sigma": 0.5, + "n_iter": 50, + "dt": 0.5, + "foreground_weight": 0.5, + "boundary_magnitude_max": Infinity + }, + "params_3d": { + "foreground_threshold": 0.5, + "density_threshold": 10.0, + "min_size": 50, + "sigma": 0.5, + "n_iter": 50, + "dt": 0.5, + "foreground_weight": 0.5, + "boundary_magnitude_max": Infinity + } +} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_d_sigma_only.json b/finetuning/v2/evaluation/optimization/configs/ais_d_sigma_only.json new file mode 100644 index 000000000..b02b5b9db --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_d_sigma_only.json @@ -0,0 +1,24 @@ +{ + "name": "d-sigma-only", + "mode": "auto", + "params_2d": { + "foreground_threshold": 0.5, + "density_threshold": 10.0, + "min_size": 100, + "sigma": 1.0, + "n_iter": 50, + "dt": 0.5, + "foreground_weight": 0.5, + "boundary_magnitude_max": Infinity + }, + "params_3d": { + "foreground_threshold": 0.5, + "density_threshold": 10.0, + "min_size": 100, + "sigma": 1.0, + "n_iter": 50, + "dt": 0.5, + "foreground_weight": 0.5, + "boundary_magnitude_max": Infinity + } +} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_dec_base_top1.json b/finetuning/v2/evaluation/optimization/configs/ais_dec_base_top1.json new file mode 100644 index 000000000..84fbd1ed1 --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_dec_base_top1.json @@ -0,0 +1 @@ +{"name": "dec-base-top1", "params_2d": {"foreground_threshold": 0.4, "density_threshold": 10.0, "min_size": 50, "sigma": 1.0, "n_iter": 800, "dt": 0.5, "foreground_weight": 0.75, "boundary_magnitude_max": 0.4}} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_dec_bnd_top1.json b/finetuning/v2/evaluation/optimization/configs/ais_dec_bnd_top1.json new file mode 100644 index 000000000..bb10a9b18 --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_dec_bnd_top1.json @@ -0,0 +1 @@ +{"name": "dec-bnd-top1", "params_2d": {"foreground_threshold": 0.5, "density_threshold": 10.0, "min_size": 50, "sigma": 1.0, "n_iter": 800, "dt": 0.5, "foreground_weight": 0.75, "boundary_magnitude_max": 0.4}} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_dec_bnd_top1_ridge1.json b/finetuning/v2/evaluation/optimization/configs/ais_dec_bnd_top1_ridge1.json new file mode 100644 index 000000000..c581545c7 --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_dec_bnd_top1_ridge1.json @@ -0,0 +1 @@ +{"name": "dec-bnd-top1-ridge1", "params_2d": {"foreground_threshold": 0.5, "density_threshold": 10.0, "min_size": 50, "sigma": 1.0, "n_iter": 800, "dt": 0.5, "foreground_weight": 0.75, "boundary_magnitude_max": 0.4, "contact_weight": 1.0}} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_dec_fgcal_top1.json b/finetuning/v2/evaluation/optimization/configs/ais_dec_fgcal_top1.json new file mode 100644 index 000000000..05165fc3e --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_dec_fgcal_top1.json @@ -0,0 +1 @@ +{"name": "dec-fgcal-top1", "params_2d": {"foreground_threshold": 0.5, "density_threshold": 50.0, "min_size": 50, "sigma": 0.5, "n_iter": 800, "dt": 0.5, "foreground_weight": 0.75, "boundary_magnitude_max": 0.4}} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_dec_fgcal_top10.json b/finetuning/v2/evaluation/optimization/configs/ais_dec_fgcal_top10.json new file mode 100644 index 000000000..775c73a3f --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_dec_fgcal_top10.json @@ -0,0 +1 @@ +{"name": "dec-fgcal-top10", "params_2d": {"foreground_threshold": 0.5, "density_threshold": 20.0, "min_size": 50, "sigma": 0.5, "n_iter": 800, "dt": 0.5, "foreground_weight": 0.75, "boundary_magnitude_max": 0.4}} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_dec_top1.json b/finetuning/v2/evaluation/optimization/configs/ais_dec_top1.json new file mode 100644 index 000000000..aa41cc4e8 --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_dec_top1.json @@ -0,0 +1 @@ +{"name": "dec-top1", "params_2d": {"foreground_threshold": 0.5, "density_threshold": 50.0, "min_size": 50, "sigma": 0.5, "n_iter": 800, "dt": 0.5, "foreground_weight": 0.75, "boundary_magnitude_max": 0.4}} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_dec_top1_ridge1.json b/finetuning/v2/evaluation/optimization/configs/ais_dec_top1_ridge1.json new file mode 100644 index 000000000..798a5084c --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_dec_top1_ridge1.json @@ -0,0 +1 @@ +{"name": "dec-top1-ridge1", "params_2d": {"foreground_threshold": 0.5, "density_threshold": 50.0, "min_size": 50, "sigma": 0.5, "n_iter": 800, "dt": 0.5, "foreground_weight": 0.75, "boundary_magnitude_max": 0.4, "contact_weight": 1.0}} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_dec_top1_ridge2_mask0.3.json b/finetuning/v2/evaluation/optimization/configs/ais_dec_top1_ridge2_mask0.3.json new file mode 100644 index 000000000..f83a81ea2 --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_dec_top1_ridge2_mask0.3.json @@ -0,0 +1 @@ +{"name": "dec-top1-ridge2-mask0.3", "params_2d": {"foreground_threshold": 0.5, "density_threshold": 50.0, "min_size": 50, "sigma": 0.5, "n_iter": 800, "dt": 0.5, "foreground_weight": 0.75, "boundary_magnitude_max": 0.4, "contact_weight": 2.0, "contact_mask_threshold": 0.3}} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_dice_reopt_base.json b/finetuning/v2/evaluation/optimization/configs/ais_dice_reopt_base.json new file mode 100644 index 000000000..fad6d0319 --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_dice_reopt_base.json @@ -0,0 +1,11 @@ +{ + "foreground_threshold": [0.3, 0.35, 0.4, 0.45, 0.5, 0.55, 0.6], + "density_threshold": [5.0, 10.0, 20.0, 50.0], + "min_size": [25, 50], + "sigma": [0.5, 1.0], + "n_iter": [400, 800, 1200, 1600], + "dt": [0.5], + "foreground_weight": [0.5, 0.75, 1.0], + "boundary_magnitude_max": [0.4], + "seed_floor": ["none"] +} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_dice_reopt_baseline_polish.json b/finetuning/v2/evaluation/optimization/configs/ais_dice_reopt_baseline_polish.json new file mode 100644 index 000000000..742faf0cc --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_dice_reopt_baseline_polish.json @@ -0,0 +1,494 @@ +{ + "combinations": [ + { + "boundary_magnitude_max": 0.4, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.4, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1600, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.375, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1600, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.425, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1600, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.4, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1600, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.4, + "foreground_weight": 0.625, + "min_size": 50, + "n_iter": 1600, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.4, + "foreground_weight": 0.875, + "min_size": 50, + "n_iter": 1600, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.4, + "foreground_weight": 1.0, + "min_size": 50, + "n_iter": 1600, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.4, + "foreground_weight": 0.75, + "min_size": 0, + "n_iter": 1600, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.4, + "foreground_weight": 0.75, + "min_size": 10, + "n_iter": 1600, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.4, + "foreground_weight": 0.75, + "min_size": 25, + "n_iter": 1600, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.4, + "foreground_weight": 0.75, + "min_size": 75, + "n_iter": 1600, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.4, + "foreground_weight": 0.75, + "min_size": 100, + "n_iter": 1600, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": "off", + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.4, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1600, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.25, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.4, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1600, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.3, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.4, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1600, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.35, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.4, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1600, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.5, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.4, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1600, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.6, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.4, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1600, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.4, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1200, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.375, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1200, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.425, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1200, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.4, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.4, + "foreground_weight": 0.625, + "min_size": 50, + "n_iter": 1200, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.4, + "foreground_weight": 0.875, + "min_size": 50, + "n_iter": 1200, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.4, + "foreground_weight": 1.0, + "min_size": 50, + "n_iter": 1200, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.4, + "foreground_weight": 0.75, + "min_size": 0, + "n_iter": 1200, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.4, + "foreground_weight": 0.75, + "min_size": 10, + "n_iter": 1200, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.4, + "foreground_weight": 0.75, + "min_size": 25, + "n_iter": 1200, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.4, + "foreground_weight": 0.75, + "min_size": 75, + "n_iter": 1200, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.4, + "foreground_weight": 0.75, + "min_size": 100, + "n_iter": 1200, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": "off", + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.4, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1200, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.25, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.4, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1200, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.3, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.4, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1200, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.35, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.4, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1200, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.5, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.4, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1200, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.6, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.4, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1200, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.375, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.425, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.4, + "foreground_weight": 0.5, + "min_size": 0, + "n_iter": 1200, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.4, + "foreground_weight": 0.5, + "min_size": 10, + "n_iter": 1200, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.4, + "foreground_weight": 0.5, + "min_size": 25, + "n_iter": 1200, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.4, + "foreground_weight": 0.5, + "min_size": 75, + "n_iter": 1200, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.4, + "foreground_weight": 0.5, + "min_size": 100, + "n_iter": 1200, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": "off", + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.4, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.25, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.4, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.3, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.4, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.35, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.4, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.5, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.4, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 1.0 + }, + { + "boundary_magnitude_max": 0.6, + "density_threshold": 10.0, + "dt": 0.5, + "foreground_threshold": 0.4, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 1.0 + } + ] +} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_dice_reopt_boundary.json b/finetuning/v2/evaluation/optimization/configs/ais_dice_reopt_boundary.json new file mode 100644 index 000000000..4be781cc8 --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_dice_reopt_boundary.json @@ -0,0 +1,26 @@ +{ + "shared": { + "foreground_threshold": [0.3, 0.35, 0.4, 0.45, 0.5, 0.55, 0.6], + "density_threshold": [5.0, 10.0, 20.0, 50.0], + "min_size": [25, 50], + "sigma": [0.5, 1.0], + "n_iter": [400, 800, 1200, 1600], + "dt": [0.5], + "foreground_weight": [0.5, 0.75, 1.0], + "boundary_magnitude_max": [0.4], + "seed_floor": ["none"] + }, + "families": { + "base": {}, + "boundary_ridge": { + "contact_weight": [0.5, 1.0, 2.0] + }, + "boundary_mask": { + "contact_mask_threshold": [0.3, 0.5, 0.7] + }, + "boundary_combined": { + "contact_weight": [1.0], + "contact_mask_threshold": [0.5] + } + } +} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_dice_reopt_boundary_combined.json b/finetuning/v2/evaluation/optimization/configs/ais_dice_reopt_boundary_combined.json new file mode 100644 index 000000000..45d499a82 --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_dice_reopt_boundary_combined.json @@ -0,0 +1,13 @@ +{ + "foreground_threshold": [0.3, 0.35, 0.4, 0.45, 0.5, 0.55, 0.6], + "density_threshold": [5.0, 10.0, 20.0, 50.0], + "min_size": [25, 50], + "sigma": [0.5, 1.0], + "n_iter": [400, 800, 1200, 1600], + "dt": [0.5], + "foreground_weight": [0.5, 0.75, 1.0], + "boundary_magnitude_max": [0.4], + "seed_floor": ["none"], + "contact_weight": [1.0], + "contact_mask_threshold": [0.5] +} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_dice_reopt_boundary_mask.json b/finetuning/v2/evaluation/optimization/configs/ais_dice_reopt_boundary_mask.json new file mode 100644 index 000000000..e9c9fd382 --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_dice_reopt_boundary_mask.json @@ -0,0 +1,12 @@ +{ + "foreground_threshold": [0.3, 0.35, 0.4, 0.45, 0.5, 0.55, 0.6], + "density_threshold": [5.0, 10.0, 20.0, 50.0], + "min_size": [25, 50], + "sigma": [0.5, 1.0], + "n_iter": [400, 800, 1200, 1600], + "dt": [0.5], + "foreground_weight": [0.5, 0.75, 1.0], + "boundary_magnitude_max": [0.4], + "seed_floor": ["none"], + "contact_mask_threshold": [0.3, 0.5, 0.7] +} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_dice_reopt_boundary_polish.json b/finetuning/v2/evaluation/optimization/configs/ais_dice_reopt_boundary_polish.json new file mode 100644 index 000000000..b85963ccb --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_dice_reopt_boundary_polish.json @@ -0,0 +1,2988 @@ +{ + "combinations": [ + { + "boundary_magnitude_max": 0.4, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.425, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.475, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.625, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.875, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 1.0, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 0, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 10, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 25, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 75, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 100, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": "off", + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.25, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.3, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.35, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.6, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 800, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.425, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 800, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.475, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 800, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 800, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.625, + "min_size": 50, + "n_iter": 800, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.875, + "min_size": 50, + "n_iter": 800, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 1.0, + "min_size": 50, + "n_iter": 800, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 0, + "n_iter": 800, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 10, + "n_iter": 800, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 25, + "n_iter": 800, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 75, + "n_iter": 800, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 100, + "n_iter": 800, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": "off", + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 800, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.25, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 800, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.3, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 800, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.35, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 800, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 800, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.6, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 800, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.425, + "foreground_weight": 0.75, + "min_size": 25, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.475, + "foreground_weight": 0.75, + "min_size": 25, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 25, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.625, + "min_size": 25, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.875, + "min_size": 25, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 1.0, + "min_size": 25, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": "off", + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 25, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.25, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 25, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.3, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 25, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.35, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 25, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 25, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.6, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 25, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.425, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.475, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.625, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.875, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 1.0, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 0, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 10, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 25, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 75, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 100, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": "off", + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.25, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.3, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.35, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.5, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.6, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 0.25, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 0.75, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 1.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 2.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 3.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.2, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.3, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.4, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.6, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.7, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.8, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.425, + "foreground_weight": 0.5, + "min_size": 25, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.475, + "foreground_weight": 0.5, + "min_size": 25, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.625, + "min_size": 25, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 25, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.875, + "min_size": 25, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 1.0, + "min_size": 25, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": "off", + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 25, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.25, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 25, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.3, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 25, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.35, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 25, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.5, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 25, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.6, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 25, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 0.25, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 25, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 25, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 0.75, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 25, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 1.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 25, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 2.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 25, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 3.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 25, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.2, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 25, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.3, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 25, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.4, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 25, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.6, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 25, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.7, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 25, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.8, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 25, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.425, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.475, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.625, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.875, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 1.0, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 0, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 10, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 25, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 75, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 100, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": "off", + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.25, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.3, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.35, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.5, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.6, + "contact_mask_threshold": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 0.25, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 0.75, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 1.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 2.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "contact_weight": 3.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.2, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.3, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.4, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.6, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.7, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.8, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.425, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.475, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.625, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.875, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 1.0, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 0, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 10, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 25, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 75, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 100, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": "off", + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.25, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.3, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.35, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.5, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.6, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.2, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.3, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.4, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.6, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.7, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.8, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.425, + "foreground_weight": 0.5, + "min_size": 25, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.475, + "foreground_weight": 0.5, + "min_size": 25, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.625, + "min_size": 25, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 25, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.875, + "min_size": 25, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 1.0, + "min_size": 25, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": "off", + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 25, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.25, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 25, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.3, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 25, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.35, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 25, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.5, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 25, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.6, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 25, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.2, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 25, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.3, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 25, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.4, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 25, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.6, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 25, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.7, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 25, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.8, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 25, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 800, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.425, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 800, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.475, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 800, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.625, + "min_size": 50, + "n_iter": 800, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 800, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.875, + "min_size": 50, + "n_iter": 800, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 1.0, + "min_size": 50, + "n_iter": 800, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 0, + "n_iter": 800, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 10, + "n_iter": 800, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 25, + "n_iter": 800, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 75, + "n_iter": 800, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 100, + "n_iter": 800, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": "off", + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 800, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.25, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 800, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.3, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 800, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.35, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 800, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.5, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 800, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.6, + "contact_mask_threshold": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 800, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.2, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 800, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.3, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 800, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.4, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 800, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.6, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 800, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.7, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 800, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_mask_threshold": 0.8, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 800, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.425, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.475, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.625, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.875, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 1.0, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 0, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 10, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 25, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 75, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 100, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": "off", + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.25, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.3, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.35, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.5, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.6, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 0.25, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 0.75, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 1.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 2.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 3.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.425, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.475, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.625, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.875, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 1.0, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 0, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 10, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 25, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 75, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 100, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": "off", + "contact_weight": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.25, + "contact_weight": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.3, + "contact_weight": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.35, + "contact_weight": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.5, + "contact_weight": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.6, + "contact_weight": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 0.25, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 0.75, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 1.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 1.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 2.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 3.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.5, + "min_size": 50, + "n_iter": 1200, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.425, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.475, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.625, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.875, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 1.0, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 0, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 10, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 25, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 75, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 100, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": "off", + "contact_weight": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.25, + "contact_weight": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.3, + "contact_weight": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.35, + "contact_weight": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.5, + "contact_weight": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.6, + "contact_weight": 0.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 0.25, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 0.75, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 1.5, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 2.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + }, + { + "boundary_magnitude_max": 0.4, + "contact_weight": 3.0, + "density_threshold": 20.0, + "dt": 0.5, + "foreground_threshold": 0.45, + "foreground_weight": 0.75, + "min_size": 50, + "n_iter": 1600, + "sigma": 0.5 + } + ] +} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_dice_reopt_boundary_ridge.json b/finetuning/v2/evaluation/optimization/configs/ais_dice_reopt_boundary_ridge.json new file mode 100644 index 000000000..f2667aac7 --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_dice_reopt_boundary_ridge.json @@ -0,0 +1,12 @@ +{ + "foreground_threshold": [0.3, 0.35, 0.4, 0.45, 0.5, 0.55, 0.6], + "density_threshold": [5.0, 10.0, 20.0, 50.0], + "min_size": [25, 50], + "sigma": [0.5, 1.0], + "n_iter": [400, 800, 1200, 1600], + "dt": [0.5], + "foreground_weight": [0.5, 0.75, 1.0], + "boundary_magnitude_max": [0.4], + "seed_floor": ["none"], + "contact_weight": [0.5, 1.0, 2.0] +} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_f2_floor_ring.json b/finetuning/v2/evaluation/optimization/configs/ais_f2_floor_ring.json new file mode 100644 index 000000000..a761d8c83 --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_f2_floor_ring.json @@ -0,0 +1,7 @@ +{ + "name": "f2-floor-ring", + "mode": "auto", + "params_2d": { + "seed_floor": "ring" + } +} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_f2_floor_ring_ms100.json b/finetuning/v2/evaluation/optimization/configs/ais_f2_floor_ring_ms100.json new file mode 100644 index 000000000..6bae56abd --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_f2_floor_ring_ms100.json @@ -0,0 +1,8 @@ +{ + "name": "f2-floor-ring-ms100", + "mode": "auto", + "params_2d": { + "seed_floor": "ring", + "min_size": 100 + } +} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_f2_floor_zero.json b/finetuning/v2/evaluation/optimization/configs/ais_f2_floor_zero.json new file mode 100644 index 000000000..e6e20a7cd --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_f2_floor_zero.json @@ -0,0 +1,7 @@ +{ + "name": "f2-floor-zero", + "mode": "auto", + "params_2d": { + "seed_floor": "zero" + } +} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_f2_floor_zero_ms100.json b/finetuning/v2/evaluation/optimization/configs/ais_f2_floor_zero_ms100.json new file mode 100644 index 000000000..19121f90a --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_f2_floor_zero_ms100.json @@ -0,0 +1,8 @@ +{ + "name": "f2-floor-zero-ms100", + "mode": "auto", + "params_2d": { + "seed_floor": "zero", + "min_size": 100 + } +} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_f_filter0p3_t400.json b/finetuning/v2/evaluation/optimization/configs/ais_f_filter0p3_t400.json new file mode 100644 index 000000000..b49fc3715 --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_f_filter0p3_t400.json @@ -0,0 +1,9 @@ +{ + "name": "f-filter0p3-t400", + "mode": "auto", + "params_2d": { + "boundary_magnitude_max": 0.3, + "n_iter": 800, + "dt": 0.5 + } +} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_f_filter0p4_t25.json b/finetuning/v2/evaluation/optimization/configs/ais_f_filter0p4_t25.json new file mode 100644 index 000000000..149144ca6 --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_f_filter0p4_t25.json @@ -0,0 +1,7 @@ +{ + "name": "f-filter0p4-t25", + "mode": "auto", + "params_2d": { + "boundary_magnitude_max": 0.4 + } +} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_f_filter0p4_t400.json b/finetuning/v2/evaluation/optimization/configs/ais_f_filter0p4_t400.json new file mode 100644 index 000000000..16066c4c2 --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_f_filter0p4_t400.json @@ -0,0 +1,9 @@ +{ + "name": "f-filter0p4-t400", + "mode": "auto", + "params_2d": { + "boundary_magnitude_max": 0.4, + "n_iter": 800, + "dt": 0.5 + } +} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_f_filter0p4_t400_fg0p6.json b/finetuning/v2/evaluation/optimization/configs/ais_f_filter0p4_t400_fg0p6.json new file mode 100644 index 000000000..f4d0f2aca --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_f_filter0p4_t400_fg0p6.json @@ -0,0 +1,10 @@ +{ + "name": "f-filter0p4-t400-fg0p6", + "mode": "auto", + "params_2d": { + "boundary_magnitude_max": 0.4, + "n_iter": 800, + "dt": 0.5, + "foreground_threshold": 0.6 + } +} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_f_filter0p4_t400_ms25.json b/finetuning/v2/evaluation/optimization/configs/ais_f_filter0p4_t400_ms25.json new file mode 100644 index 000000000..1e94e490d --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_f_filter0p4_t400_ms25.json @@ -0,0 +1,10 @@ +{ + "name": "f-filter0p4-t400-ms25", + "mode": "auto", + "params_2d": { + "boundary_magnitude_max": 0.4, + "n_iter": 800, + "dt": 0.5, + "min_size": 25 + } +} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_f_filter0p5_t400.json b/finetuning/v2/evaluation/optimization/configs/ais_f_filter0p5_t400.json new file mode 100644 index 000000000..d3912c88e --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_f_filter0p5_t400.json @@ -0,0 +1,9 @@ +{ + "name": "f-filter0p5-t400", + "mode": "auto", + "params_2d": { + "boundary_magnitude_max": 0.5, + "n_iter": 800, + "dt": 0.5 + } +} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_f_filter0p6_t25.json b/finetuning/v2/evaluation/optimization/configs/ais_f_filter0p6_t25.json new file mode 100644 index 000000000..9c89170fa --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_f_filter0p6_t25.json @@ -0,0 +1,7 @@ +{ + "name": "f-filter0p6-t25", + "mode": "auto", + "params_2d": { + "boundary_magnitude_max": 0.6 + } +} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_f_filter0p6_t400.json b/finetuning/v2/evaluation/optimization/configs/ais_f_filter0p6_t400.json new file mode 100644 index 000000000..8ba4126c5 --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_f_filter0p6_t400.json @@ -0,0 +1,9 @@ +{ + "name": "f-filter0p6-t400", + "mode": "auto", + "params_2d": { + "boundary_magnitude_max": 0.6, + "n_iter": 800, + "dt": 0.5 + } +} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_f_t400.json b/finetuning/v2/evaluation/optimization/configs/ais_f_t400.json new file mode 100644 index 000000000..37e236f9f --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_f_t400.json @@ -0,0 +1,8 @@ +{ + "name": "f-t400", + "mode": "auto", + "params_2d": { + "n_iter": 800, + "dt": 0.5 + } +} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_grid_lm3d_v4.json b/finetuning/v2/evaluation/optimization/configs/ais_grid_lm3d_v4.json new file mode 100644 index 000000000..5375439b2 --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_grid_lm3d_v4.json @@ -0,0 +1,10 @@ +{ + "foreground_threshold": [0.4, 0.5, 0.6, 0.7], + "density_threshold": [5.0, 10.0, 20.0, 50.0], + "min_size": [50, 100, 200], + "sigma": [0.5, 1.0], + "n_iter": [50, 800], + "dt": [0.5], + "foreground_weight": [0.5], + "boundary_magnitude_max": [null, 0.4, 0.6] +} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_grid_lm_v4.json b/finetuning/v2/evaluation/optimization/configs/ais_grid_lm_v4.json new file mode 100644 index 000000000..dd390cf58 --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_grid_lm_v4.json @@ -0,0 +1,40 @@ +{ + "foreground_threshold": [ + 0.4, + 0.5, + 0.6, + 0.7 + ], + "density_threshold": [ + 5.0, + 10.0, + 20.0, + 50.0 + ], + "min_size": [ + 25, + 50, + 100 + ], + "sigma": [ + 0.5, + 1.0 + ], + "n_iter": [ + 50, + 800 + ], + "dt": [ + 0.5 + ], + "foreground_weight": [ + 0.25, + 0.5, + 0.75 + ], + "boundary_magnitude_max": [ + null, + 0.4, + 0.6 + ] +} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_s0_travel_100.json b/finetuning/v2/evaluation/optimization/configs/ais_s0_travel_100.json new file mode 100644 index 000000000..862805638 --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_s0_travel_100.json @@ -0,0 +1,8 @@ +{ + "name": "s0-travel-100", + "mode": "auto", + "params_2d": { + "n_iter": 200, + "dt": 0.5 + } +} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_s0_travel_12p5.json b/finetuning/v2/evaluation/optimization/configs/ais_s0_travel_12p5.json new file mode 100644 index 000000000..72b637c42 --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_s0_travel_12p5.json @@ -0,0 +1,8 @@ +{ + "name": "s0-travel-12p5", + "mode": "auto", + "params_2d": { + "n_iter": 25, + "dt": 0.5 + } +} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_s0_travel_200.json b/finetuning/v2/evaluation/optimization/configs/ais_s0_travel_200.json new file mode 100644 index 000000000..a02cf0ff4 --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_s0_travel_200.json @@ -0,0 +1,8 @@ +{ + "name": "s0-travel-200", + "mode": "auto", + "params_2d": { + "n_iter": 400, + "dt": 0.5 + } +} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_s0_travel_400.json b/finetuning/v2/evaluation/optimization/configs/ais_s0_travel_400.json new file mode 100644 index 000000000..34b57b255 --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_s0_travel_400.json @@ -0,0 +1,8 @@ +{ + "name": "s0-travel-400", + "mode": "auto", + "params_2d": { + "n_iter": 800, + "dt": 0.5 + } +} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_s0_travel_50.json b/finetuning/v2/evaluation/optimization/configs/ais_s0_travel_50.json new file mode 100644 index 000000000..107d9f851 --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_s0_travel_50.json @@ -0,0 +1,8 @@ +{ + "name": "s0-travel-50", + "mode": "auto", + "params_2d": { + "n_iter": 100, + "dt": 0.5 + } +} diff --git a/finetuning/v2/evaluation/optimization/configs/ais_v_filter_only.json b/finetuning/v2/evaluation/optimization/configs/ais_v_filter_only.json new file mode 100644 index 000000000..1b85c789e --- /dev/null +++ b/finetuning/v2/evaluation/optimization/configs/ais_v_filter_only.json @@ -0,0 +1,19 @@ +{ + "name": "v-filter-only", + "mode": "auto", + "params_2d": { + "sigma": 1.0, + "min_size": 50, + "boundary_magnitude_max": 0.4 + }, + "params_3d": { + "foreground_threshold": 0.5, + "density_threshold": 10.0, + "min_size": 100, + "sigma": 0.5, + "n_iter": 50, + "dt": 0.5, + "foreground_weight": 0.5, + "boundary_magnitude_max": 0.4 + } +} diff --git a/finetuning/v2/evaluation/optimization/diagnose_decoder_fields.py b/finetuning/v2/evaluation/optimization/diagnose_decoder_fields.py new file mode 100644 index 000000000..df67c5a7a --- /dev/null +++ b/finetuning/v2/evaluation/optimization/diagnose_decoder_fields.py @@ -0,0 +1,173 @@ +"""Field diagnostics of cached decoder predictions: contact geometry and foreground extent. + +For every cached sample of a manifest (`benchmark_ais_optimization.py predict` must have run) and per dataset: +the cosine between the predicted flow one pixel on either side of a ground-truth contact pixel (the field of a +well separated pair flips, so the cosine is negative), the same cosine one pixel apart inside objects (a +smooth field gives +1), the median distance magnitude at contacts and inside, the foreground area ratio +`area(fg > threshold) / area(gt)` and, for five channel predictions, the Dice of `contact > 0.5` with the +ground-truth contact target. The proposal's "what would show that it worked" figures. CPU only, reader only. + +`--contact-mode` must match what the fifth channel was trained on ("touching" for the `contact` / `both` +decoders, "all" for the `boundary` / `boundary_fgcal` decoders), otherwise its precision is scored against a +target that calls the head's correct pixels negative. The recall on the touching lines and on the +background-facing boundary lines is reported separately in both modes, so the two targets can be compared. + + export MICRO_SAM2_JOINT_CHECKPOINT_ROOT= + python diagnose_decoder_fields.py --joint-checkpoint contact --subset primary training_extra --output + python diagnose_decoder_fields.py --joint-checkpoint boundary --contact-mode all --output +""" + +import argparse +import os +import sys +from pathlib import Path +from typing import Dict + +import numpy as np +import pandas as pd + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) +import benchmark_ais_optimization as ais # noqa: E402 +from micro_sam.v2.transforms.labels import object_boundaries, touching_boundaries # noqa: E402 + + +def _shift(array: np.ndarray, axis: int, step: int) -> np.ndarray: + """The array shifted by 'step' along 'axis' with edge replication (so a difference at the border is zero).""" + shifted = np.roll(array, -step, axis=axis) + index = [slice(None)] * array.ndim + if step > 0: + index[axis] = slice(-step, None) + shifted[tuple(index)] = np.take(array, [-1], axis=axis) + else: + index[axis] = slice(None, -step) + shifted[tuple(index)] = np.take(array, [0], axis=axis) + return shifted + + +def flow_cosines(directed: np.ndarray, where: np.ndarray, offset: int) -> np.ndarray: + """Cosine between the flow 'offset' pixels before and after every pixel of 'where', along the axis of the + stronger local label change (both in-plane axes are tried and the smaller cosine kept: the flip axis).""" + norms = np.linalg.norm(directed, axis=0) + 1e-6 + unit = directed / norms + cosines = [] + for axis in range(1, unit.ndim): + before = _shift(unit, axis, -offset) + after = _shift(unit, axis, offset) + cosines.append((before * after).sum(axis=0)) + cosine = np.minimum.reduce(cosines) + return cosine[where] + + +def _contact_target(labels: np.ndarray, mode: str, dilation: int) -> np.ndarray: + """The training target of the fifth channel: the touching lines only, or every object boundary.""" + if mode == "all": + return object_boundaries(labels, dilation=dilation) + return touching_boundaries(labels, radius=1, dilation=dilation) + + +def sample_row( + prediction: np.ndarray, labels: np.ndarray, threshold: float, contact_mode: str = "touching", +) -> Dict[str, float]: + ndim = labels.ndim + foreground, directed = prediction[0], prediction[1:4][-ndim:] + contact_gt = touching_boundaries(labels, radius=1, dilation=0) + interior = (labels > 0) & ~touching_boundaries(labels, radius=2, dilation=0) + from skimage.segmentation import find_boundaries + interior &= ~find_boundaries(labels, mode="inner") + magnitude = np.linalg.norm(directed, axis=0) + fg_mask = foreground > threshold + row = { + "gt_objects": int(len(np.unique(labels)) - 1), + "contact_pixels": int(contact_gt.sum()), + "fg_area_ratio": float(fg_mask.sum() / max(1, (labels > 0).sum())), + "fg_iou": float((fg_mask & (labels > 0)).sum() / max(1, (fg_mask | (labels > 0)).sum())), + "magnitude_bg_median": float(np.median(magnitude[labels == 0])) if (labels == 0).any() else float("nan"), + "magnitude_interior_median": float(np.median(magnitude[interior])) if interior.any() else float("nan"), + } + if contact_gt.any(): + row["magnitude_contact_median"] = float(np.median(magnitude[contact_gt])) + for offset in (1, 3): + row[f"cosine_contact_{offset}px"] = float(np.median(flow_cosines(directed, contact_gt, offset))) + else: + row["magnitude_contact_median"] = float("nan") + row["cosine_contact_1px"] = row["cosine_contact_3px"] = float("nan") + if interior.any(): + for offset in (1, 3): + row[f"cosine_interior_{offset}px"] = float(np.median(flow_cosines(directed, interior, offset))) + if prediction.shape[0] > 4: + contact_pred = prediction[4] > 0.5 + target = _contact_target(labels, contact_mode, 1) + denominator = contact_pred.sum() + target.sum() + row["contact_dice"] = float(2 * (contact_pred & target).sum() / denominator) if denominator else float("nan") + row["contact_pred_pixels"] = int(contact_pred.sum()) + row["contact_target_pixels"] = int(target.sum()) + # Share of the predicted contact mass that lies within two pixels of the target. + near = _contact_target(labels, contact_mode, 2) + row["contact_precision_2px"] = float((contact_pred & near).sum() / max(1, contact_pred.sum())) + row["contact_recall"] = float((contact_pred & target).sum() / max(1, target.sum())) + # Mode independent, so that a touching-target and a full-boundary head can be read side by side: + # where the head fires on the lines between objects, and where it fires on the background-facing rim. + touching = touching_boundaries(labels, radius=1, dilation=1) + rim = object_boundaries(labels, dilation=1) & ~touching + row["recall_touching"] = float((contact_pred & touching).sum() / max(1, touching.sum())) + row["recall_bg_boundary"] = float((contact_pred & rim).sum() / max(1, rim.sum())) + return row + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + # The manifest / checkpoint arguments of benchmark_ais_optimization.py, so `_manifests` and friends apply. + parser.add_argument("--kind", choices=ais.KINDS, default="v5") + parser.add_argument("--subset", nargs="+", default=["primary"]) + parser.add_argument("--data-root", type=Path, default=Path("/mnt/vast-nhr/projects/cidas/cca/data")) + parser.add_argument("--output-root", type=Path, + default=Path("/mnt/vast-nhr/projects/cidas/cca/experiments/micro_sam2/apg_optimization")) + parser.add_argument("--campaign-root", type=Path, + default=Path("/mnt/vast-nhr/projects/cidas/cca/experiments/micro_sam2/apg_optimization/3d_v2")) + parser.add_argument("--model-type", default="hvit_t") + parser.add_argument("--joint-checkpoint", default="best") + parser.add_argument("--ndim", choices=["2", "3", "both"], default="2") + parser.add_argument("--datasets", nargs="*", default=None) + parser.add_argument("--foreground-threshold", type=float, default=0.5) + parser.add_argument("--contact-mode", choices=["touching", "all"], default="touching", + help="The training target of the fifth channel; 'all' for the boundary decoders.") + parser.add_argument("--output", default=None) + args = parser.parse_args() + checkpoint_id = ais._checkpoint_identity(args.model_type, args.joint_checkpoint) + dimensions = ais._dimensions(args) + rows = [] + for manifest in ais._manifests(args): + cache = ais.PredictionCache(args.output_root, checkpoint_id, manifest["manifest_checksum"]) + samples = [s for s in manifest["samples"] if int(s["ndim"]) in dimensions] + if args.datasets: + samples = [s for s in samples if s["dataset"] in args.datasets] + for index, sample in enumerate(samples, start=1): + if not cache.has(sample): + raise FileNotFoundError(f"No cached prediction for '{sample['sample_id']}' under '{cache.root}'.") + prediction, labels, valid, _ = cache.load(sample) + if valid is not None: + labels = np.where(valid, labels, 0) + row = sample_row(prediction, labels.astype("int64"), args.foreground_threshold, args.contact_mode) + row.update({ + "sample_id": sample["sample_id"], "dataset": sample["dataset"], "subset": manifest.get("subset"), + }) + rows.append(row) + if index % 50 == 0: + print(f"{manifest.get('subset')}: {index}/{len(samples)}") + table = pd.DataFrame(rows) + numeric = [c for c in table.columns if c not in ("sample_id", "dataset", "subset")] + summary = table.groupby("dataset")[numeric].median(numeric_only=True) + summary["n_samples"] = table.groupby("dataset").size() + pd.set_option("display.width", 250) + print(f"\nCheckpoint {checkpoint_id[:8]}: per-dataset medians") + print(summary.to_string(float_format=lambda v: f"{v:.3f}")) + if args.output: + Path(args.output).parent.mkdir(parents=True, exist_ok=True) + table.to_csv(args.output, index=False) + summary.to_csv(str(Path(args.output).with_name(Path(args.output).stem + "_summary.csv"))) + print(f"written {args.output}") + + +if __name__ == "__main__": + main() diff --git a/finetuning/v2/evaluation/optimization/notes/AIS_DECODER_HANDOVER.md b/finetuning/v2/evaluation/optimization/notes/AIS_DECODER_HANDOVER.md new file mode 100644 index 000000000..737b42de4 --- /dev/null +++ b/finetuning/v2/evaluation/optimization/notes/AIS_DECODER_HANDOVER.md @@ -0,0 +1,147 @@ +# Hand-over: AIS decoder campaign, round 2 (full-boundary channel) - reading the results + +Written 2026-09-07 18:00, state refreshed 21:35 when the session was stopped on purpose so a later one can +watch the trainings land (they finish ~40 min after that session's 12 h job would have ended). Everything is committed on +branch `ais-train-optim`. Read first: `AIS_DECODER_TRAINING.md` - sections 4.0-4.4 hold round 1, **4.5** the +round-1 completions, **4.7** the sweep optima (which revise 4.4), **5.1-5.3** the round-2 launch and the chain. Memory note `ais-decoder-campaign-state`. +`` = `/mnt/vast-nhr/projects/cidas/cca/experiments/micro_sam2/apg_optimization`, +`` = `/ais_decoder_training`, `` = `finetuning/v2/evaluation/optimization`, +`` = `finetuning/v2/generalist/ais_decoder`, `` = `/ais/reports`, +python = `micromamba activate new-stack`. + +## 1. What is left to do + +**Nothing has to be submitted.** Round 2 is chained end to end (section 5.3 of the notes); the successor reads +the tables the chain writes and finishes the write-up: + +1. Read the round-2 tables (section 3 below) and write them into **section 5.6** of `AIS_DECODER_TRAINING.md` (5.4 records what the ridge and mask modes + mean once the channel is a full boundary and the signature to look for; 5.5 the `SBATCH_EXPORT` incident). +2. Decide point 1.1 (the fifth channel) with the boundary target on the evidence, and update section 4.4 point 3 + if the verdict changes. The user's rule: only cross-dataset wins count - balanced mSA plus the gate + (>= 9 / 11 up, worst > -2 %, balanced >= +2 %) against the fine-tuned `baseline` on dev, confirmed on the + holdout. No per-dataset fits. +3. Write the conclusive overview of all six decoders (**section 6**), update the memory note, commit. + +## 2. State when the session was stopped (2026-09-07 21:35) + +**First command on resume** - if the trainings are gone from `squeue`, check how they ended before reading +anything, because `afterany` runs the evaluation on whatever `best.pt` exists: + +```bash +sacct -j 15776831,15776833,15776838,15776839,15777359,15777505 -X \ + -o JobID,JobName%26,Start,Elapsed,State,ExitCode +squeue -u $USER -h -o "%i %j %T %M %R" | sort -k2 +``` +`COMPLETED` after ~12.6 h = the full 48000 iterations. `TIMEOUT` or `PREEMPTED` = a short run; say so in the +write-up and check `best.pt`'s epoch in `/logs/slurm/ais_decoder__.err`. A preempted job +requeues from iteration 0 (`Requeue=1`) and cannot finish inside its window - that needs a decision, not a rerun. + +## 2b. The chain + +| job | what | expected | +|---|---|---| +| 15776831 `boundary` (ggpu158), 15776833 `boundary_fgcal` (ggpu192) | the two trainings; at 21:29 they were at 15841 / 15673 of 48000 iterations, 1.05 it/s including validation | `boundary` ~05:45, `boundary_fgcal` ~05:49; SLURM wall limit 07:15:50 | +| 15776838 / 15776839 `ais_eval_` | `afterany` the training: stage, cache v5 primary / training_extra / holdout and apg3d primary / holdout, then the `current-defaults`, `contact-ridge` and `contact-mask` screens | ~06:00, screens ~07:00 | +| 15777359 `ais_decoder_tuning2` | `afterany` both evaluations: waits for the 2d caches, submits the two grid sweeps (1728 combinations) and the eight-configuration contact screen per new variant, then ranks all six sweeps into `/dec__sweep_dev.csv` | ~06:05, rankings ~11:00 | +| 15777505 `ais_decoder_finalize_r2` | `afterany` both evaluations: submits the `dec-top1` screens of the two new decoders `afterok` their prediction jobs, waits for every round-2 screen (up to 8 h), then writes the overview tables and the field diagnostics | ~06:05, tables ~11:00-13:00 | +| ~~15772853~~ | the round-1 launcher; FAILED at 19:03 on the same edited-script pitfall (notes 5.2), nothing lost - it had submitted its sweeps and all four rankings are written | - | + +**Round 1 is complete**: `/decoders_{defaults,tuned,final}_*`, `decoders_final_3d*`, +`decoder_fields_{production,baseline,contact,fgcal,both}*` and all four `dec__sweep_dev.csv`, written up +in sections 4.1-4.7. The command that ranks a sweep, for the two new decoders should `tuning2` not get to it: + +```bash +cd ; export MICRO_SAM2_JOINT_CHECKPOINT_ROOT=/ais_decoder_training/staged +for v in baseline contact; do $PY report_ais_sweep.py --grid configs/ais_grid_lm_v4.json \ + --subset primary training_extra --joint-checkpoint $v --top 25 --output /dec_${v}_sweep_dev.csv; done +``` + +Check the chain with `squeue -u $USER -h -o "%i %j %T %M %R" | sort -k2`; task markers are `logs/.done` / +`.failed` in the newest `/jobs/_/`; drivers log to `/logs/slurm/`. + +## 3. The tables to read, and the read-outs the user needs + +Reference for every comparison is the fine-tuned `baseline`; the production decoder is the second reference. + +| file under `` | what | +|---|---| +| `decoders_all_defaults_{dev,holdout}{,_datasets,_mechanisms}.csv` | all six plus production under the library defaults, `contact-ridge` and `contact-mask` | +| `decoders_all_tuned_{dev,holdout}*.csv` | all six at the shared tuned `dec-top1` (reference: baseline at `dec-top1`), with the ridge / mask variants | +| `decoders_all_3d*.csv` | apg3d primary + holdout; regression instrument only (see 4.5: even the plain fine-tune loses 25-60 % per LM family, and the round-1 five-channel decoders are 0 because `boundary_magnitude_max` removes every instance) | +| `decoders_{boundary,boundary_fgcal}_contact_dev*.csv` | the eight ridge / mask settings against the decoder's own defaults | +| `dec_{boundary,boundary_fgcal}_sweep_dev.csv` | each new decoder at its own sweep optimum (dev-tuned; read the model comparison on the holdout) | +| `decoder_fields_{boundary,boundary_fgcal}{,_summary}.csv` | field diagnostics, scored with `--contact-mode all` | + +Read-outs: + +1. **The gate.** Balanced mSA and the gate against `baseline` on dev, confirmed on the holdout, at the defaults + *and* at `dec-top1`. Round-1 numbers to beat: `contact` -4.7 % / -5.2 % (defaults), -4.2 % / -3.8 % + (`dec-top1`); `both` -1.3 % / -1.9 % and +1.3 % / +1.8 %; `fgcal` +0.6 % / +1.1 % and +2.4 % / +1.7 %. + Compare each decoder at its **own** sweep optimum too (4.7): baseline 0.4244 at `foreground_threshold` 0.4, + fgcal 0.4298 and both 0.4254 at 0.5, contact 0.4083 at 0.6 - so fgcal is +1.3 %, both +0.2 % and contact + -3.8 % against baseline's own optimum. **Which threshold the two boundary decoders want is the cleanest test + of whether their foreground is calibrated**: 0.6 like `contact` means the extra task still inflates the + foreground, 0.4-0.5 means the full-boundary target does not. +2. **Is the head confident now?** The round-1 contact head was precise but under-confident exactly on the + datasets whose merges motivated it. `recall_touching` of the round-1 `contact` decoder (per-dataset medians): + dynamicnuclearnet 0.72, yeaz 0.76, livecell 0.63, covid_if 0.59, tissuenet **0.19**, neurips **0.13**, puma + 0.12, tnbc 0.03, deepbacs 0.015, dic_hepg2 **0.001**. The boundary head has to lift tissuenet, neurips, + deepbacs and dic_hepg2; `recall_bg_boundary` (0.00-0.15 for `contact`) shows whether it also learned the + background-facing rim, i.e. whether it learned the target at all. +3. **Do the shared-feature losses disappear?** Round 1 lost -12 % deepbacs, -21 % covid_if, -26 % deepseas, + -38 % dic_hepg2 through the shared features, not through the ridge (the head never fired there). Read the + per-dataset columns and the mechanism shares: dic_hepg2 lost seeds (unseeded 43.9 -> 57.1 %), deepbacs split + its rods (1.7 -> 4.1 %). If the boundary target removes these, point 1.1 becomes a candidate again. +4. **The merges it was for.** Seeded-merge share at the defaults and at `dec-top1`, with and without the ridge. + Round 1: baseline 8.1 % / 13.4 %, `contact` + ridge 4.5 % / 4.1 %, `both` + ridge 4.2 % / 4.0 %. +5. **Extent.** `fg_area_ratio` and `matched_iou` per dataset - never the summary CSV's mean (deepseas 12-91 and + neurips 2.3-12 dominate it; see 4.5). +6. **The ridge / mask setting** of a denser head: with a few percent of the pixels positive the mask mode at 0.5 + may finally do something (it was inert in round 1 because the head rarely exceeded 0.5). + +## 4. If something went wrong + +- **A training timed out** (wall 07:16): `afterany` still fires, and the driver stages `best.pt` of the last + finished epoch, so the chain completes on a slightly shorter run. Note the epoch in the write-up. +- **A job was preempted** (everything runs on `grete:preemptible`): resubmit the driver by hand, e.g. + `bash /evaluate_ais_decoder.sh boundary best`, or the frozen copies under `/jobs/frozen/` + (`finalize_round2_.sh`, `launch_tuning_.sh`, `finalize_.sh`). +- **Reports come out empty**: check `--epoch 856a433c4b33348e1d85c4c13278f057` still matches + `implementation_checksum()`. It hashes `benchmark_ais_optimization.py`, `common.py`, `parameter_search.py`, + `micro_sam/v2/instance_segmentation.py` and `micro_sam/v2/postprocessing.py` - **do not touch those five while + the chain is in flight**, or the new runs get a different epoch and every filtered report goes blank. + `micro_sam/v2/transforms/labels.py` and the readers are not hashed, so the round-2 code changes did not move it. +- **A screen or sweep task failed**: `/jobs/_/logs/.failed` holds the reason; rerun the one + command from `tasks.txt`. + +## 5. Pitfalls (met, fixed, listed so they are not re-debugged) + +- **Never submit a repo path for a long-running driver.** Bash re-reads a running script by byte offset, so + editing the file while a job sleeps in a wait loop breaks the parse - that is how the round-1 finalisation died + after waiting 7.4 h (notes 5.2). Copy it to `/jobs/frozen/_.sh` and submit the copy. +- **Our own array can block our own jobs.** An unschedulable job at the head of a partition blocks every + lower-priority job of the same user; the round-2 trainings pended 17 minutes behind our own sweep array with + seven 3g slices free (notes 5.1). Diagnosis: `squeue -u $USER -O "jobid,name,state,reason,priority"` and look + for `TopOfQueue`. Fix: `scontrol update jobid= nice=100`, or `scontrol hold` / `release`. +- **`sbatch --test-only` is worthless on `grete:preemptible`** - it ignores preemption and returned the same + 10-hour-away estimate for every pool while jobs started immediately. +- `diagnose_decoder_fields.py --contact-mode` must match the training target of the fifth channel (`touching` + for `contact` / `both`, `all` for `boundary` / `boundary_fgcal`), otherwise the head's precision is scored + against a target that calls its correct pixels negative. +- **`SBATCH_EXPORT=none` is set on this cluster**, so `sbatch` does not propagate the submitting environment: + campaign parameters passed as environment variables reach the job as empty and the script falls back to its + defaults, silently doing plausible but wrong work (notes 5.5). Pass them on the command line. +- `tasks_done` uses `ls -td | head -1`, so an empty *newer* job directory of the same name shadows a finished + one. If a resubmission has to be cancelled, move its directory to `/jobs/_superseded/`. +- Historical sweep results before the Dice-foreground re-optimization ignored the contact keywords. The current cached + scorer mirrors the production ridge and mask paths; its implementation checksum keeps those corrected sweeps + separate from the invalid old cache. +- Python 3.14 starts DataLoader workers through a fork server (30-60 s each, every epoch); + `train_ais_decoder.py` forces `fork`. Do not remove. +- Files with fewer than three objects (yeaz frames) make torch_em's sampler raise after 500 attempts; the subset + wrappers redraw. `RandomSubsetDataset` exists because torch_em splits `n_samples` over files. +- Trainer checkpoints pickle the datasets: import `ais_decoder_lib` before `torch.load` of a `best.pt`. +- Always export `MICRO_SAM2_JOINT_CHECKPOINT_ROOT=/staged` and + `MICRO_SAM2_JOINT_EXPORT_ROOT=/model_exports` before any benchmark command. +- The session cwd drifts after `cd`; use absolute paths. `.sh` files are git-ignored: `git add -f`. +- The CPU preset takes 16 cores per task and roughly one task runs at a time, so a 22-task sweep needs ~2 h. diff --git a/finetuning/v2/evaluation/optimization/notes/AIS_DECODER_TRAINING.md b/finetuning/v2/evaluation/optimization/notes/AIS_DECODER_TRAINING.md new file mode 100644 index 000000000..936217f35 --- /dev/null +++ b/finetuning/v2/evaluation/optimization/notes/AIS_DECODER_TRAINING.md @@ -0,0 +1,1203 @@ +# AIS decoder training campaign: contact channel and boundary-calibrated foreground + +Decision log of the campaign that trains the changes proposed in `AIS_DECODER_TRAINING_PROPOSAL.md` +(points 1.1 and 4.1) and compares them on the AIS benchmarks. Branch `ais-train-optim`; paths relative to +`finetuning/v2/evaluation/` unless they start with `micro_sam/` or `finetuning/`; `` = +`/mnt/vast-nhr/projects/cidas/cca/experiments/micro_sam2/apg_optimization`. + +## 1. Question and design (2026-09-07) + +Two losses of the v4 geodesic `hvit_t` decoder cannot be recovered by post-processing (`AIS_V4_OPTIMIZATION.md`, +sections D2/D3): merges of touching cells (livecell 36 %, tissuenet 25 %, neurips 27 % of the objects) and an +over-wide foreground (`fg > 0.5` covers 1.1-3.5x the ground-truth area). Four decoders are trained under +identical conditions and compared with the library's current post-processing defaults (no re-tuning): + +| variant | target / loss change | output channels | +|---|---|---| +| `baseline` | none (foreground Dice, three masked-MSE distance channels) | 4 | +| `contact` | fifth channel = touching boundaries (`touching_boundaries`, radius 1, dilation 1), Dice + BCE | 5 | +| `fgcal` | foreground loss = Dice + BCE weighted 5x within +-2 px of every object boundary (`boundary_weight=4`) | 4 | +| `both` | both changes | 5 | + +Decisions (with the user): decoder-only training with the image encoder frozen at the v4 joint weights (the +interactive half is untouched; a single GPU suffices); warm start from the v4 decoder (the baseline is then +"v4 decoder + 12 h fine-tune on the tuning data"); the BCE variant of point 4.1, not the signed distance; +training data = train splits of the AIS tuning datasets; the `both` job is submitted after the other three. + +Library changes are epoch A5 (commit `d85bccb`, implementation checksum `856a433c4b33348e1d85c4c13278f057`, +previous A4 `184eba917bd0cff28b5719b5584f6967`): `micro_sam/v2/transforms/labels.py` (`touching_boundaries`, +`contact=True`), `micro_sam/v2/loss/directed_distance_based.py` (`contact`, `boundary_weight`), +`micro_sam/v2/models/util.py` (sigmoid on channels >= 4), five-channel plumbing in `instance_segmentation.py` +and `batched_inference.py`, `flow_instance_segmentation(contact=, contact_weight=, contact_mask_threshold=)` +(opt-in; the default path is bit-identical, the three-channel `out[1:]` convenience stays, any other channel +count raises), APG and evaluation mirrors read `[1:4]`, the harness gains `fg_area_ratio` and the configs +`configs/ais_contact_ridge.json` (`contact_weight` 1.0) and `configs/ais_contact_mask.json` +(`contact_mask_threshold` 0.5). + +Epoch A5 bit-identity (2026-09-07 02:55): the production decoder's `current-defaults` runs under A5 (jobs +15769317 / 15769318, `--ndim 2` for the images) reproduce the A4 runs sample by sample on v5 primary (240), +training_extra (157), holdout (233) and the apg3d primary (57) and holdout (18) crops (mSA, matched, predicted +objects, merged / absorbed counts, fg and matched IoU all identical; balanced 0.2457 / 0.4253 / 0.2437). These +A5 run directories are the production reference for the decoder comparison (`fg_area_ratio` included). + +## 2. Data + +Train splits of nine of the eleven tuning datasets, built from torch_em path lists +(`finetuning/v2/generalist/ais_decoder/ais_decoder_lib.py::build_datasets`); the trainer's validation set is +the last 5 % (at least 2 files) of every sorted train list, so the evaluation manifests (val splits) and the +test splits stay untouched. The file lists of every run are written to +`/checkpoints/ais_decoder_/data_manifest.json`. + +| dataset | train files (val tail) | samples per epoch | raw handling | +|---|---|---|---| +| livecell (8 cell types) | 3253 minus 30 images that also appear in the val split | 8 x 25 | grayscale, `MinInstanceSampler(6)` | +| tissuenet | 2580 | 200 | `raw/rgb` per-channel percentile normalisation, `labels/cell` | +| dynamicnuclearnet | 4950 | 200 | grayscale | +| deepbacs (mixed) | 125 | 120 | `_to_8bit` | +| dic_hepg2 | 302 | 120 | rgb png, channels kept distinct | +| neurips_cellseg (Training-labeled) | 1000 | 150 | mixed formats, to rgb | +| yeaz bf / phc 2d / phc stacks | 207 / 14 / 14 | 80 / 20 / 20 | stacks read as (1, 512, 512) patches | +| puma nuclei | 138 | 100 | rgb h5 | +| tnbc | 34 | 60 | rgb h5 (channel-first) | + +Excluded: **deepseas** (binary masks; connected components merge touching cells, which would teach "no +contact" exactly at contacts and give merged blobs one geodesic centre; its 40 manifest crops are also +train-split files) and **covid_if** (49 files without a split, 5 tuning crops, 44 production-scored). Both stay +evaluation datasets and are unseen for all four models. zarr/h5/stack datasets are wrapped in +`RandomSubsetDataset` because torch_em splits `n_samples` uniformly over the files (200 samples over 2451 +tissuenet files would only ever read the first 200 files). + +## 3. Training set-up + +`finetuning/v2/generalist/ais_decoder/train_ais_decoder.py`: `FrozenEncoderUniSAM2` (32 features wide, the +encoder in eval mode with `requires_grad=False`, so autograd stores no encoder activations), warm start from +the v4 `unetr_state` (a five-channel decoder keeps the four pretrained output rows and a fresh fifth row), +AdamW over the decoder parameters (lr 5e-5, weight decay 0.1), `ReduceLROnPlateau(0.9, patience 10)`, +bf16 autocast, `UniSAM2Trainer` (loss = metric = `DirectedDistanceLoss` of the variant), patch (512, 512), +percentile augmentation as in the generalist recipe. Checkpoints `/ais_decoder_training/checkpoints/ +ais_decoder_/{best,latest}.pt`; staging (`stage_ais_decoder_checkpoint.py`) writes the lean joint-format +file `/ais_decoder_training/staged/joint_sam2_hvit_t_multi_gpu/.pt` (v4 `model_state` + trained +`unetr_state`) for `MICRO_SAM2_JOINT_CHECKPOINT_ROOT=/ais_decoder_training/staged`, +`--joint-checkpoint `. + +Compute: `grete:shared`, one A100 (`-G A100:1`, 16 CPUs, 64 G, 12 h); the eight `3g.40gb` slices were held by +two-day jobs of one user (preemption off) and recent single-A100 jobs on `grete:shared` started within 1-1.5 h. + +### Smoke tests (session slice 1g.20gb, one CPU, 2026-09-07 01:20) + +`contact`, batch 4, one loader worker, 20 iterations: loaders built in ~1 min (1270 samples per epoch, +150 validation crops), `check_loader` 3 raw / 5 target channels, loader 3.9 samples/s with one worker +(1.02 s per batch of 4, so 12 workers deliver ~45 samples/s), GPU step 1.24 s per iteration at batch 4 +(7.2 GiB allocated, 11.0 GiB reserved), 20 iterations plus a 38-batch validation in 60 s, loss 1.25 -> 1.50 +validation metric (untrained fifth channel). Extrapolation to a full A100 (six to seven times the SMs of the +slice): ~0.4 s per iteration at batch 8, ~15 GiB allocated, so batch 8 fits a 40 GB node with margin and the +loader is not the bottleneck. + +### Budget and submission (2026-09-07 01:30) + +Queue at submission time: two of the eight `3g.40gb` slices free (the other six held by two-day jobs until +2026-09-08 evening), 14 `2g.20gb` slices free, grete:shared with 200 of 244 A100s allocated and about 16 free +on non-reserved nodes but other users' single-A100 arrays pending with reason "WaitingInQueue"; +`sbatch --test-only` estimated a 25 h start for a 16-CPU A100 job, which the recent starts (1-1.5 h) contradict. +Decision: identical training for all four (batch 8, `--epoch-scale 4` = 5080 samples / 635 iterations per +epoch, **48000 iterations**, lr 5e-5, 12 loader workers), spread over the two pools so that every model is +done within about 12 h: `baseline` and `contact` on the free `3g.40gb` slices (`grete:preemptible`, +`-t 14:00:00`; ~0.8 s per iteration expected, ~11 h), `fgcal` and `both` on `grete:shared` A100 (`-t 12:00:00`, +~0.4 s per iteration, ~6 h), `both` submitted last with `--dependency=after` on the other three. 48000 +iterations x 8 = 384k samples = 76 epochs of the 5080-sample epoch; the hardware only changes the wall time. + +First submission (jobs 15769310-13, 01:26): all four died with exit 139 within 3-11 min of training. Two +causes, both fixed before the resubmission: (1) Python 3.14 starts DataLoader workers through a fork server, +so every worker re-imported the environment (30-60 s each, serialised; the two 3g jobs spent nine minutes +before their first iteration, and the non-persistent validation workers would have paid it every epoch) - +`train_ais_decoder.py` now forces the fork start method; (2) a 512^2 zarr file with fewer than three objects +makes torch_em's `MinInstanceSampler` reject the same crop 500 times and raise, which ended the run +(`RandomSubsetDataset` now redraws another file, `FixedSubsetDataset` does the same for validation, and both +wrap every dataset; torch_em's image-collection datasets already rotate images after 50 failed crops). Measured +rejection rates of the container datasets: yeaz phase-contrast stack frames 9/290, yeaz bright field 1/40, +dynamicnuclearnet 0/120, tissuenet 0/60, puma 0/40, tnbc 0/34 - the sparse yeaz frames were the trigger. + +Second submission (02:47): `baseline` 15769606 and `contact` 15769607 on `3g.40gb` (started at once, both on +ggpu158), `fgcal` 15769609 on `grete:shared` A100 (started 02:52 on ggpu114 after five minutes in the queue), +`both` 15769611 after the three. Measured speed on a 3g.40gb slice: 1.08 iterations/s at batch 8, so 48000 +iterations take about 12.8 h (finish ~15:40). The evaluation is chained by SLURM: `ais_eval_` jobs +15769618-15769621 run `finetuning/v2/generalist/ais_decoder/evaluate_ais_decoder.sh best` after their +training succeeds (stage, predict v5 primary / training_extra / holdout and apg3d primary / holdout, then the +`current-defaults` runs plus `contact-ridge` / `contact-mask` for the five-channel decoders on the caches). + +## 4. Results + +Readout (`report_ais_decoders.py`, reference = the `baseline` decoder under the library defaults; the production +decoder `5a729846...` is the second reference, epoch A5 runs): + +- 4.1 Development set (primary + training_extra, eleven datasets, `--ndim 2`): balanced mSA, gate verdict, + merged + absorbed share, `fg_area_ratio`, `matched_iou` per variant and configuration (`current-defaults`; + `contact-ridge` and `contact-mask` for the five-channel decoders). +- 4.2 Holdout (five datasets). +- 4.3 3D crops (apg3d primary / holdout): family macros with `current-defaults` (and `contact-ridge`). +- 4.4 Field diagnostics (`diagnose_decoder_fields.py`): contact cosine at +-1 / +-3 px, contact Dice, magnitude + at contacts vs interior, foreground area ratio at threshold 0.5. +- 4.5 Training curves: validation loss per variant (TensorBoard under `/ais_decoder_training/logs/`). + +Training curves at 03:55 (validation loss per epoch of 635 iterations, loss = metric of each variant, so the +values are not comparable across variants): baseline 0.180, 0.174, 0.166, 0.166, 0.162, 0.157 (six epochs); +contact 0.933, 0.799, 0.762, 0.764, 0.733, 0.724 (the fresh contact head dominates the early loss); fgcal 0.493, +0.450, 0.434, 0.438, 0.428, 0.414, 0.410, 0.410, 0.405, 0.414, 0.402 (eleven epochs); both 1.149, 1.071, 1.017, +1.041, 1.008, 0.983, 0.979, 0.965, 0.950. All four decrease; none has plateaued yet. + +`fgcal` finished at 09:19 after 6.43 h (48000 iterations, peak 14.1 GiB allocated / 21.2 GiB reserved on an +A100-40GB; best epoch 64 of 76, validation loss 0.367); its evaluation chain (predict arrays 15772069 / 15772070, +screens 15772071 / 15772072) started at 09:21. + +### 4.0 Preliminary: `fgcal` against the production decoder (09:35, before the fine-tuned baseline exists) + +`current-defaults`, epoch A5, checkpoint `dd52aee4...` vs production `5a729846...`: + +| set | production | fgcal | up | worst | merged + absorbed (object-weighted) | unseeded | fg area ratio (mean over datasets) | +|---|---:|---:|---|---|---|---|---| +| dev (11) | 0.3437 | 0.4170 (+21.3 %) | 9 / 11 | covid_if -35.9 %, deepseas -25.8 % | 25.2 % -> 17.5 % | 20.7 % -> 14.1 % | see below | +| holdout (5) | 0.2437 | 0.3938 (+61.6 %) | 5 / 5 | tissuenet +19.0 % | 29.4 % -> 17.9 % | 20.4 % -> 14.5 % | | + +Per dataset (dev): livecell 0.277 -> 0.365, tissuenet 0.224 -> 0.263, dynamicnuclearnet 0.545 -> 0.831, deepbacs +0.181 -> 0.326, dic_hepg2 0.003 -> 0.174, neurips 0.226 -> 0.295, yeaz 0.616 -> 0.819, puma 0.469 -> 0.536, tnbc +0.368 -> 0.405, covid_if 0.740 -> 0.474, deepseas 0.134 -> 0.099. Merged + absorbed: livecell 36.9 -> 21.9 %, +tissuenet 25.6 -> 13.7 %, deepbacs 23.2 -> 9.8 %, neurips 28.0 -> 30.8 %. Foreground area ratio at 0.5: deepbacs +1.76 -> 1.25, livecell 1.11 -> 1.06, dynamicnuclearnet 0.92 -> 1.02, tissuenet 0.87 -> 0.75 (more under-coverage), +covid_if 1.05 -> 1.30, puma / tnbc / yeaz ~1.0 in both. + +Reading: the two datasets that lose are exactly the two the fine-tuned decoders never saw (covid_if, deepseas), +and dic_hepg2 / dynamicnuclearnet / yeaz (never in the joint training) gain the most, so this comparison mostly +measures "12 h of decoder fine-tuning on the tuning datasets' train splits", not the boundary-weighted loss. +The isolating comparison is against the fine-tuned `baseline` (same data, same budget), pending. + +Field diagnostics (`diagnose_decoder_fields.py`, dev caches, per-dataset medians; `ais/reports/decoder_fields_{fgcal,production}.csv`): + +| quantity | production | fgcal | +|---|---|---| +| distance magnitude in the background | 0.83-0.86 on every dataset (the label fill value) | 0.03-0.05 | +| magnitude at ground-truth contact pixels | 0.10-0.29 | 0.04-0.11 | +| flow cosine across a contact, +-1 px | 0.63 (tissuenet), 0.71 (dnn), 0.77 (yeaz), 0.88 (livecell) | 0.44, 0.47, 0.37, 0.88 | +| flow cosine across a contact, +-3 px | -0.58, -0.55, -0.52, -0.15 | -0.70, -0.72, -0.85, -0.34 | +| fg IoU at 0.5 (median) | livecell 0.84, dnn 0.78, deepbacs 0.66, neurips 0.65, tissuenet 0.76, covid_if 0.92 | 0.88, 0.93, 0.78, 0.74, 0.74, 0.77 | +| fg area ratio at 0.5 (median) | deepbacs 1.49, dic_hepg2 0.10, dnn 0.90, tissuenet 0.91, covid_if 1.05 | 1.13, 1.04, 1.01, 0.79, 1.28 | + +The background magnitude change matters for `boundary_magnitude_max`: the filter assumes a false region's +boundary runs through magnitude ~1; with ~0 in the background it no longer discriminates. Contact flips are +sharper but not negative at +-1 px. Attribution (fine-tuning vs the boundary loss) waits for the baseline. + +3D crops (apg3d primary + holdout, 75 crops, `current-defaults`), fgcal vs production: every LM family loses +(celegans_atlas 0.104 -> 0.011, embedseg_platy_ish 0.339 -> 0.135, embedseg_platy_nuclei 0.259 -> 0.086, +embedseg_skull 0.118 -> 0.079, gonuclear 0.256 -> 0.136, platynereis_nuclei 0.068 -> 0.006) and the EM CREMI +scores roughly double (cremi 0.99 -> 2.24, cremi_seen 0.59 -> 2.15, snemi 0.97 -> 1.95, humanneurons 1.35 -> +1.97); the volume foreground balloons (area ratio celegans 1.28 -> 2.38, gonuclear 2.02 -> 2.68) and merges rise +(celegans 33 -> 75 %). Expected for a 2D-only, LM-only decoder fine-tune (the 3D path of the decoder saw no data), +and the reason these decoders cannot replace the production one for volumes or EM; the 3D crops serve as the +regression instrument of the campaign only. + +`both` (fgcal + contact, checkpoint `25e2a32a...`, best epoch 73) against production, 2D (09:50): + +| set / configuration | production | fgcal defaults | both defaults | both contact-ridge | both contact-mask | +|---|---:|---:|---:|---:|---:| +| dev balanced (11) | 0.3437 | 0.4170 | 0.4090 | 0.4096 | 0.4093 | +| holdout balanced (5) | 0.2437 | 0.3938 | 0.3819 | 0.3819 | 0.3823 | +| dev merged + absorbed, object-weighted | 25.2 % | 17.5 % | 14.8 % | 12.6 % | 13.9 % | +| livecell mSA / merged + absorbed | 0.277 / 36.9 % | 0.365 / 21.9 % | 0.382 / 19.2 % | 0.385 / 15.7 % | 0.384 / 17.8 % | +| tissuenet mSA / merged + absorbed | 0.224 / 25.6 % | 0.263 / 13.7 % | 0.272 / 12.7 % | 0.268 / 11.9 % | 0.271 / 12.5 % | +| neurips mSA / merged + absorbed | 0.226 / 28.0 % | 0.295 / 30.8 % | 0.317 / 22.4 % | 0.314 / 19.7 % | 0.318 / 21.2 % | +| deepbacs mSA | 0.181 | 0.326 | 0.287 | 0.285 | 0.287 | +| deepseas / covid_if mSA (unseen) | 0.134 / 0.740 | 0.099 / 0.474 | 0.046 / 0.462 | 0.046 / 0.460 | 0.046 / 0.462 | + +The five-channel model wins on the three touching-cell datasets (livecell, tissuenet, neurips) and loses on +deepbacs and on the two unseen datasets, so its balanced score is 2 % below fgcal. The contact ridge at weight +1.0 removes another 2-4 points of merges on livecell / tissuenet / neurips for +0.1-0.7 % mSA; the mask mode at +0.5 changes little (the contact head is rarely above 0.5). Both post-processing settings are untuned. The +isolating pairs (contact vs baseline, both vs fgcal with the same data) complete when the 3g jobs finish. +Tables: `ais/reports/decoders_prelim_{primary_training_extra,holdout}*.csv`. + +`both` on the 3D crops scores exactly 0 on every LM family: the volumes get 180-350 seeds and a foreground +(fg IoU 0.34, area ratio 2.7) but every instance is removed by `boundary_magnitude_max=0.4`, because the +five-channel decoder's magnitude inside the true objects of a volume is 0.86 (median; fgcal 0.35, production +0.27), i.e. its 3D distance field has drifted towards the fill value. Without the filter one celegans crop gives +57 instances (ground truth 72; production 57). A 3D-only effect of the 2D fine-tune, recorded, not pursued. +Tables: `ais/reports/decoders_prelim_3d*.csv`. + +Contact head of `both` (`ais/reports/decoder_fields_both.csv`, medians, threshold 0.5): Dice against the true contact +lines livecell 0.57 (precision within 2 px 0.81, recall 0.55), yeaz 0.67 (0.88 / 0.65), dynamicnuclearnet 0.65 +(0.94 / 0.51), covid_if 0.45, tissuenet 0.26 (precision 0.93 but recall 0.16), neurips 0.19 (recall 0.02), +dic_hepg2 / deepbacs / tnbc / puma ~0 (dic_hepg2 has 2169 true contact pixels per crop and predicts none). The +head is precise but under-confident on the datasets with the largest merge losses, which is why the mask mode at +0.5 did nothing; a class-weighted or focal contact loss is the recipe change to consider for the big run. The +contact training also sharpened the flow: the +-1 px contact cosine drops from 0.47 / 0.44 / 0.37 (fgcal, dnn / +tissuenet / yeaz) to 0.11 / 0.32 / -0.03. + +### 4.6 Tuning launched in the meantime (10:53) + +Grid sweeps (`configs/ais_grid_lm_v4.json`, 1728 combinations) on the development caches of fgcal (jobs 15772848 / +15772849) and both (15772850 / 15772851), contact-configuration screens for both (15772852: ridge 0.5 / 2 / 4, mask +0.3 / 0.7, ridge 1 + mask 0.5, on dev and holdout), and a launcher (15772853, +`finetuning/v2/generalist/ais_decoder/launch_tuning_after_caches.sh`) that submits the same for baseline and contact +once their caches exist and then ranks every sweep into `ais/reports/dec__sweep_dev.csv` +(`report_ais_sweep.py`, reference = library defaults). Read the model comparison at tuned settings on the holdout, +not on the development set the sweep tuned on. + +fgcal sweep ranked (11:55, `ais/reports/dec_fgcal_sweep_dev.csv`, 1728 combinations, reference = library defaults on +the fgcal caches, dev balanced 0.4170): the best shared configuration reaches 0.4298 (+3.1 %, 6 / 11 up, worst +-6.4 %, mean ratio to the per-dataset optimum 0.935); nothing passes the gate. The top rows all use travel 800 +(n_iter 800, dt 0.5), density 50 (or 20), sigma 0.5, foreground weight 0.75, min_size 50, filter 0.4 or off - a +different regime from the production defaults (travel 25, density 10, sigma 1.0), consistent with a field that now +converges to sinks (magnitude ~0 in the background, sharper flips). Confirmation of the top-1 and the density-20 +variant on dev + holdout: `configs/ais_dec_fgcal_top{1,10}.json`, job dec_fgcal_top_screen. + +both sweep ranked (12:35, `ais/reports/dec_both_sweep_dev.csv`, contact terms not part of the cached sweep): best +shared configuration 0.4254 (+4.0 % over its defaults 0.4090, 8 / 11 up, worst -8.2 %), the same regime as fgcal +(travel 800, density 50, sigma 0.5, foreground weight 0.75, min_size 50). At the tuned shared setting fgcal stays +1 % ahead of both on the development set (0.4298 vs 0.4254). Screens of this shared top configuration alone and +with the contact terms (ridge 1; ridge 2 + mask 0.3) on dev + holdout for both: `configs/ais_dec_top1*.json`, +job dec_both_top_screen. + +Contact configurations on both (13:25, `ais/reports/dec_both_contact_{primary_training_extra,holdout}*.csv`, +reference = both under the library defaults, other parameters at the defaults): ridge weights 0.5 / 1 / 2 / 4 give ++0.2 / +0.2 / +0.3 / +0.3 % balanced on dev (4-5 of 11 up, worst -1.5 %) and +0.2 / 0.0 / +0.1 / -0.3 % on holdout; +mask thresholds 0.3 / 0.5 / 0.7 give 0.0 / +0.1 / 0.0 % (dev) and -0.0 / +0.1 / +0.1 % (holdout); ridge 1 + mask 0.5 ++0.2 / 0.0 %. The ridge does what it is meant to - seeded merges fall from 6.3 % to 4.1 % of the objects on dev (7.0 +to 4.7 % on holdout) with no change in unseeded objects - but the recovered objects hardly move mSA at IoU 0.5, so +with these decoders the contact channel is not where the remaining mSA is (merges are down from 13 % to 6 % of the +objects already by the fine-tuning). + +Tuned comparison of the two finished decoders (13:40, `ais/reports/dec_tuned_prelim_{primary_training_extra,holdout}*.csv`; +`dec-top1` = travel 800, density 50, sigma 0.5, fw 0.75, min_size 50, filter 0.4, the shared optimum of both sweeps): + +| configuration | dev balanced (11) | holdout balanced (5) | holdout seeded merges | +|---|---:|---:|---:| +| production, defaults | 0.3437 | 0.2437 | 16.7 % | +| fgcal, defaults | 0.4170 | 0.3938 | 8.7 % | +| fgcal, dec-top1 | 0.4298 | 0.4094 | 10.9 % | +| both, defaults | 0.4090 | 0.3819 | 7.0 % | +| both, dec-top1 | 0.4254 | 0.4098 | 8.9 % | +| both, dec-top1 + contact ridge 1 | 0.4271 | 0.4113 | 4.5 % | + +Holdout per dataset (fgcal top1 / both top1 + ridge): deepbacs 0.389 / 0.340, dic_hepg2 0.226 / 0.250, dynamicnuclearnet +0.822 / 0.825, livecell 0.355 / 0.382, tissuenet 0.255 / 0.261. Reading: at tuned settings the two decoders are within +0.5 % of each other on the holdout; the tuned regime (few, converged seeds) brings merges back for fgcal, which the +contact ridge removes for both without changing mSA; the shared configuration trades livecell / tissuenet for deepbacs +/ dic_hepg2 (the 6 / 11 "up" of the sweep). The isolating pairs against baseline and contact are pending. + +`baseline` finished at 15:42 after 12.87 h on a 3g.40gb slice (48000 iterations, best epoch 75 of 76, peak 14.1 GiB); +`contact` at 15:50 after 12.97 h. Their chains (staging, caches, default and contact screens), the finalisation job +and the launcher's sweeps follow automatically; the shared tuned configuration `dec-top1` is screened for baseline +(job dec_baseline_top_screen) and contact as well, so all four decoders can be read at the same tuned setting. + +### 4.1 The isolating comparison under the library defaults (16:05; contact pending) + +Reference = the fine-tuned `baseline` (same data, budget, initialisation; checkpoint under `staged/baseline.pt`): + +| decoder | dev balanced (11) | vs baseline | up / worst | holdout balanced (5) | vs baseline | up / worst | seeded merges dev / holdout | +|---|---:|---:|---|---:|---:|---|---| +| production | 0.3437 | -17.1 % | 2 / 11, dic_hepg2 -98 % | 0.2437 | -37.4 % | 0 / 5 | 13.3 % / 16.7 % | +| baseline | 0.4145 | - | - | 0.3894 | - | - | 8.1 % / 8.8 % | +| fgcal | 0.4170 | +0.6 % | 6 / 11, dic_hepg2 -8.8 % | 0.3938 | +1.1 % | 4 / 5, deepbacs -5.0 % | 8.3 % / 8.7 % | +| both | 0.4090 | -1.3 % | 7 / 11, deepseas -49 % | 0.3819 | -1.9 % | 3 / 5, dic_hepg2 -24 % | 6.3 % / 7.0 % | +| both + contact ridge 1 | 0.4096 | -1.2 % | 7 / 11 | 0.3819 | -1.9 % | 3 / 5 | 4.2 % / 4.8 % | + +Per dataset against baseline (dev): fgcal livecell 0.0 %, tissuenet +3.0 %, neurips +5.8 %, puma +4.4 %, yeaz +4.0 %, +dnn +0.7 %, deepbacs -5.0 %, dic_hepg2 -8.8 %, tnbc -2.3 %, covid_if -4.5 %, deepseas +11 %; both livecell +4.6 %, +tissuenet +6.6 %, neurips +13.7 %, yeaz +3.6 %, puma +2.8 %, tnbc +2.2 %, dnn +0.5 %, deepbacs -16.4 %, dic_hepg2 +-28.6 %, covid_if -6.9 %, deepseas -48.6 %. Holdout: fgcal deepbacs -5.0 %, dic_hepg2 +7.5 %, dnn +1.4 %, livecell ++0.4 %, tissuenet +5.0 %; both deepbacs -16.4 %, dic_hepg2 -23.7 %, dnn +1.7 %, livecell +5.3 %, tissuenet +10.6 %. + +Reading: +1. Almost the entire gain over the production decoder (+21 % dev, +60 % holdout) is the decoder fine-tune on the + tuning datasets' train splits, with the unchanged loss. The fine-tuned baseline already cuts merges from 13 % to + 8 % of the objects and moves the foreground area ratio to ~1 on most datasets (deepbacs 1.76 -> 1.19). +2. The boundary-weighted foreground loss (point 4.1) adds +0.6 % / +1.1 % balanced, on 6 / 11 and 4 / 5 datasets, + with a -5 to -9 % loss on deepbacs or dic_hepg2; the foreground area ratio and the merge share are unchanged + against baseline (tissuenet under-coverage 0.71 -> 0.75). It fails the generalization gate. +3. The contact channel (point 1.1, here on top of fgcal) is a strong, dataset-dependent lever: +5 to +14 % on the + touching-cell datasets (livecell, tissuenet, neurips) with the merge share down to 6 % (4 % with the ridge), but + -16 % on deepbacs and -24 to -29 % on dic_hepg2, so the balanced score is 1-2 % below baseline. The contact-vs- + baseline pair (pending) separates the channel from the fgcal loss it was stacked on. + +Mechanisms behind the both-vs-baseline differences (`ais/reports/decoders_isolating_dev_mechanisms.csv`, % of +objects): dic_hepg2 loses seeds (unseeded 43.9 -> 57.1 %, absorbed 34.1 -> 44.3 %; merges unchanged), deepbacs +splits its thin rods (1.7 -> 4.1 %) with a lower matched IoU (0.771 -> 0.744); livecell (merges 11.8 -> 9.2 %, +unseeded 15.2 -> 13.4 %), tissuenet (3.2 -> 2.5 %, 23.9 -> 22.6 %) and neurips (13.2 -> 10.0 %, 17.2 -> 14.9 %) gain on +both counts with higher matched IoU (0.777 -> 0.784, 0.737 -> 0.740, 0.772 -> 0.789). The contact head itself +never fires on dic_hepg2 or deepbacs, so their losses come from the shared features the extra task changed, not +from the ridge. + +### 4.2 All four decoders under the library defaults (16:25; `ais/reports/decoders_defaults_{primary_training_extra,holdout}*.csv`) + +| decoder (configuration) | dev balanced | vs baseline | up / 11 | worst | holdout balanced | vs baseline | up / 5 | worst | seeded merges dev | +|---|---:|---:|---|---|---:|---:|---|---|---:| +| production | 0.3437 | -17.1 % | 2 | dic_hepg2 -98 % | 0.2437 | -37.4 % | 0 | dic_hepg2 -99 % | 13.3 % | +| baseline | 0.4145 | - | - | - | 0.3894 | - | - | - | 8.1 % | +| fgcal | 0.4170 | +0.6 % | 6 | dic_hepg2 -8.8 % | 0.3938 | +1.1 % | 4 | deepbacs -5.0 % | 8.3 % | +| contact | 0.3952 | -4.7 % | 4 | dic_hepg2 -38 % | 0.3691 | -5.2 % | 2 | dic_hepg2 -40 % | 7.8 % | +| contact + ridge 1 | 0.4004 | -3.4 % | 4 | deepseas -26 % | 0.3770 | -3.2 % | 2 | dic_hepg2 -23 % | 4.5 % | +| both | 0.4090 | -1.3 % | 7 | deepseas -49 % | 0.3819 | -1.9 % | 3 | dic_hepg2 -24 % | 6.3 % | +| both + ridge 1 | 0.4096 | -1.2 % | 7 | deepseas -49 % | 0.3819 | -1.9 % | 3 | dic_hepg2 -21 % | 4.2 % | + +contact vs baseline per dataset (dev, defaults / ridge): tissuenet +8.1 / +6.5 %, neurips +5.6 / +7.9 %, livecell ++1.3 / +4.2 %, yeaz +2.3 / +2.1 %, puma -0.7 / -1.8 %, tnbc -0.2 / -1.2 %, dnn -3.1 / -3.5 %, deepbacs -12.0 / -11.5 %, +dic_hepg2 -38.1 / -8.5 %, covid_if -21.1 / -21.0 %, deepseas -25.6 / -25.9 %. Holdout: tissuenet +10.3 / +8.2 %, +livecell +1.5 / +4.8 %, dnn -2.6 / -2.6 %, deepbacs -12.0 / -11.5 %, dic_hepg2 -40.4 / -22.6 %. + +Reading (all four, same data, budget and initialisation): +- Point 4.1 (boundary-weighted foreground BCE): +0.6 % / +1.1 % balanced, 6 / 11 and 4 / 5 datasets up, a 5-9 % loss + on one dataset each time; foreground calibration and merge share unchanged against baseline. A marginal, non- + uniform effect; it does not pass the gate. +- Point 1.1 (contact channel): a strong dataset-dependent trade, not a general gain: +6 to +8 % on tissuenet and + neurips, +1 to +5 % on livecell and yeaz, against -12 % on deepbacs, -21 % on covid_if, -26 % on deepseas and + -38 % on dic_hepg2 (-8.5 % once the ridge recovers the absorbed objects). The losses come through the shared + features (fewer seeds on dic_hepg2 and covid_if, split rods on deepbacs), not through the ridge; the head itself + never fires on those datasets. Stacked on fgcal (`both`) the trade is milder (-1.3 % / -1.9 %) with the same sign + pattern. +- The dominant effect of the campaign is neither: the plain fine-tune on the tuning datasets' train splits lifts the + decoder from 0.344 to 0.415 (dev) and from 0.244 to 0.389 (holdout), removes 40 % of the merges and calibrates + the foreground area to ~1 on most datasets. covid_if and deepseas, the two datasets left out of training, lose + (-33 % and -33 % for baseline vs production), so part of this is in-domain specialisation. + +### 4.3 All four decoders at the shared tuned configuration (16:30; `ais/reports/decoders_tuned_{primary_training_extra,holdout}*.csv`) + +`dec-top1` (travel 800, density 50, sigma 0.5, foreground weight 0.75, min_size 50, filter 0.4) is the optimum of both +the fgcal and the both sweep; reference = baseline at dec-top1 (0.4200 dev, 0.4025 holdout; its own sweep pending). + +| decoder (configuration) | dev balanced | vs baseline | up / 11 | worst | holdout balanced | vs baseline | up / 5 | worst | seeded merges dev | +|---|---:|---:|---|---|---:|---:|---|---|---:| +| baseline (dec-top1) | 0.4200 | - | - | - | 0.4025 | - | - | - | 13.4 % | +| fgcal (dec-top1) | 0.4298 | +2.4 % | 9 | dic_hepg2 -5.1 % | 0.4094 | +1.7 % | 4 | deepbacs -4.5 % | 10.6 % | +| contact (dec-top1) | 0.4024 | -4.2 % | 5 | -26 % | 0.3872 | -3.8 % | 2 | deepbacs -18 % | 12.2 % | +| contact (dec-top1 + ridge 1) | 0.4155 | -1.1 % | 6 | -27 % | 0.4044 | +0.5 % | 3 | deepbacs -14 % | 4.1 % | +| both (dec-top1) | 0.4254 | +1.3 % | 8 | deepseas -47 % | 0.4098 | +1.8 % | 4 | deepbacs -16 % | 7.7 % | +| both (dec-top1 + ridge 1) | 0.4271 | +1.7 % | 8 | deepseas -48 % | 0.4113 | +2.2 % | 4 | deepbacs -16 % | 4.0 % | + +Holdout per dataset at dec-top1 (baseline / fgcal / contact + ridge / both + ridge): deepbacs 0.407 / 0.389 / 0.351 / +0.340, dic_hepg2 0.220 / 0.226 / 0.248 / 0.250, dynamicnuclearnet 0.804 / 0.822 / 0.784 / 0.825, livecell 0.341 / +0.355 / 0.381 / 0.382, tissuenet 0.240 / 0.255 / 0.258 / 0.261. The tuned regime (few converged seeds) raises the +merge share of the four-channel decoders from 8 % to 13 %; the contact ridge is the only thing that brings it to 4 %. + +### 4.4 Conclusions for the training recipe (2026-09-07, 16:35) + +1. In-domain data dominates. A decoder-only fine-tune with the unchanged loss on the tuning datasets' train + splits gains +21 % (dev) / +60 % (holdout) over the production decoder, halves the merge share and calibrates the + foreground area; every proposed loss change is a small correction on top of that. For the next big run the + composition of the training data (which of the evaluation datasets' train splits are included) matters far more + than the two loss changes. +2. Point 4.1 (boundary-weighted foreground BCE): consistently small and positive. +0.6 / +1.1 % at the defaults, + +2.4 / +1.7 % at the tuned setting, 9 of 11 dev datasets up at the tuned setting, but a 5-9 % loss on one dataset + (dic_hepg2 or deepbacs) each time, so it misses the gate's worst-loss bound. It does not change the foreground + area ratio or the merge share against the fine-tuned baseline. Cheap and safe to include, not decisive. +3. Point 1.1 (contact channel, plain Dice + BCE, ridge in the watershed): a dataset-dependent trade. +6 to +10 % on + tissuenet, +6 to +8 % on neurips, +1 to +5 % on livecell (the datasets whose merges motivated it), but -12 % on + deepbacs, -21 % on covid_if, -26 % on deepseas and -38 % on dic_hepg2 through the shared features (seeds lost, + rods split), with a head that never fires on those datasets. The ridge itself is effective and cheap (merges + 6-13 % -> 4 % at any setting) and recovers half of the dic_hepg2 loss; stacked on fgcal (`both`) the trade + narrows to -1.3 % / -1.9 % at the defaults and +1.7 % / +2.2 % at the tuned setting. Under the generalization rule + the channel as trained here is not a win; the levers to try before including it in a big run are a class- + weighted or focal contact loss (the head is precise but under-confident: recall 0.16 on tissuenet, 0.02 on + neurips at 0.5) and a lower contact loss weight so that the shared features do not lose seeds on large or thin + cells. The 3D path of the five-channel decoder also drifted (section 4.0), which a joint 2D + 3D run avoids. +4. Post-processing for fine-tuned decoders: their fields converge (magnitude ~0 in the background, sharper flips), + and the tuned optimum moves to long travel (800) with a high density threshold (50), sigma 0.5 and foreground + weight 0.75 (+1.3 to +3.1 % over the current defaults, 6-9 of 11 up, worst -5 to -8 %); the production defaults + are no longer the right regime for such decoders, and `boundary_magnitude_max` loses its premise. + +**Revised by 4.7 (18:15), once baseline had its own sweep:** point 2's "+2.4 % at the tuned setting" is measured +at `dec-top1`, which is fgcal's optimum and 1 % below baseline's own (baseline wants `foreground_threshold` 0.4, +fgcal and both 0.5). At each decoder's own optimum the boundary-weighted foreground loss is worth **+1.3 %**, and +its mechanism is the threshold calibration, not the merge share. Point 4's "the tuned optimum moves to density 50 +and sigma 0.5" holds for the loss-changed decoders only; baseline keeps the production density 10 / sigma 1.0 and +only lengthens the travel. + +(the 3D tables of all four and the unattended finalisation outputs are in 4.5, the sweep rankings in 4.7) + +### 4.5 Round-1 completions: the 3D table of all four, the field diagnostics, the mask mode (17:25) + +Written by `ais_decoder_finalize2` (15777315, four minutes once the screens were in; see 5.2 for why the first +attempt died): `ais/reports/decoders_final_{dev,holdout,3d}*.csv` and `decoder_fields_{baseline,contact}*.csv`. + +**All four on the 3D crops** (apg3d primary + holdout, 75 crops, `current-defaults`; the balanced score mixes LM +mSA with the negated CREMI error, so read the families, not the aggregate): + +| family | production | baseline | fgcal | contact | both | +|---|---:|---:|---:|---:|---:| +| celegans_atlas | 0.104 | 0.040 | 0.011 | 0.000 | 0.000 | +| embedseg_platy_ish | 0.339 | 0.156 | 0.135 | 0.000 | 0.000 | +| embedseg_platy_nuclei | 0.259 | 0.115 | 0.086 | 0.000 | 0.000 | +| embedseg_skull | 0.118 | 0.238 | 0.078 | 0.000 | 0.000 | +| gonuclear | 0.256 | 0.132 | 0.135 | 0.000 | 0.000 | +| platynereis_nuclei | 0.068 | 0.052 | 0.006 | 0.000 | 0.000 | +| cremi / cremi_seen (lower is better) | 0.99 / 0.59 | 1.87 / 1.23 | 2.24 / 2.15 | 2.19 / 2.05 | 2.09 / 1.89 | +| snemi / humanneurons (lower is better) | 0.97 / 1.35 | 1.69 / 1.97 | 1.95 / 1.98 | 2.16 / 2.37 | 1.92 / 2.04 | + +The regression is the 2D-only fine-tune itself, not the loss changes: the unchanged-loss `baseline` already loses +25-60 % of every LM family (embedseg_skull is the exception, 0.118 -> 0.238) and adds 0.6-1.0 to every CREMI +error, before any loss change. Both five-channel decoders are exactly 0 on all six LM families because +`boundary_magnitude_max=0.4` removes every instance of a field whose magnitude no longer dips at boundaries +(section 4.0). Read as: a 2D-only decoder fine-tune cannot replace the production decoder for volumes, and the +magnitude filter has to be re-decided for any fine-tuned decoder - not as a verdict on the two loss changes. + +**Field diagnostics of all four** (dev, per-dataset medians, `decoder_fields__summary.csv`): + +- Background distance magnitude 0.83-0.86 (production, the label fill value) -> 0.03-0.08 for *all four* + fine-tuned decoders. `boundary_magnitude_max` loses its premise for every one of them, not only for fgcal. +- dic_hepg2 is a production-decoder failure, not a loss effect: fg IoU 0.07 at an area ratio of 0.10 (it barely + predicts foreground there, hence mSA 0.003); every fine-tuned decoder reaches fg IoU 0.89-0.90 at ratio + 1.03-1.08. This single dataset carries most of the +21 % dev gain over production. +- The two datasets held out of training move the wrong way, which is where their losses come from: covid_if + fg IoU 0.92 -> 0.72-0.77 with the area ratio 1.05 -> 1.26-1.35 (over-coverage), deepseas fg IoU 0.46 -> + 0.16-0.38 with the ratio 1.90 -> 0.53-1.06 (`both` the worst at 0.16 / 0.53, and it is the variant with the + -49 % deepseas loss). +- The flow flip across a contact (cosine at +-1 px, lower is sharper) is sharpened by the fine-tune and again by + the contact channel: dynamicnuclearnet 0.71 -> 0.25 (baseline) -> 0.11 (contact), tissuenet 0.63 -> 0.39 -> + 0.28, yeaz 0.77 -> 0.40 -> -0.18. The channel does to the field exactly what it was meant to do; the mSA it + buys is the question, not the mechanism. +- fgcal against baseline moves the foreground in both directions rather than calibrating it: tissuenet + under-coverage 0.75 -> 0.79 (better), deepbacs over-coverage 1.03 -> 1.13 (worse), the rest within 0.02. + +**The mask mode**, added to the four-way defaults table: `contact` 0.3952 -> 0.3962 (dev) and 0.3691 -> 0.3706 +(holdout), `both` 0.4090 -> 0.4093 and 0.3819 -> 0.3823. Confirms 4.2 - the mask is inert because the head +rarely exceeds 0.5. + +**The fifth channel of the two round-1 decoders**, rescored with the mode-independent recalls (per-dataset +medians, threshold 0.5, `--contact-mode touching` = the target they were trained on): + +| dataset | target px | `contact` pred px / Dice / precision 2px / recall_touching / recall_bg | `both` pred px / Dice / precision / recall_touching / recall_bg | +|---|---:|---|---| +| yeaz | 7570 | 6967 / 0.69 / 0.84 / **0.76** / 0.15 | 5772 / 0.67 / 0.88 / 0.65 / 0.09 | +| dynamicnuclearnet | 170 | 264 / 0.65 / 0.79 / **0.72** / 0.01 | 126 / 0.65 / 0.94 / 0.51 / 0.00 | +| livecell | 11400 | 11954 / 0.60 / 0.77 / **0.63** / 0.07 | 8499 / 0.57 / 0.81 / 0.55 / 0.05 | +| covid_if | 935 | 2000 / 0.36 / 0.35 / 0.59 / 0.03 | 1497 / 0.45 / 0.46 / 0.58 / 0.02 | +| tissuenet | 7117 | 1281 / 0.30 / 0.93 / **0.19** / 0.00 | 950 / 0.26 / 0.93 / 0.16 / 0.00 | +| neurips_cellseg | 1007 | 296 / 0.21 / 0.46 / **0.13** / 0.00 | 14 / 0.20 / 0.10 / 0.02 / 0.00 | +| puma | 233 | 128 / 0.17 / 0.49 / 0.12 / 0.00 | 28 / 0.05 / 0.37 / 0.03 / 0.00 | +| tnbc | 180 | 19 / 0.08 / 0.31 / 0.03 / 0.00 | 0 / 0.01 / 0.00 / 0.00 / 0.00 | +| deepbacs | 132 | 30 / 0.02 / 0.07 / 0.02 / 0.00 | 4 / 0.00 / 0.00 / 0.00 / 0.00 | +| dic_hepg2 | 3790 | 40 / 0.01 / 0.16 / 0.001 / 0.00 | 0 / 0.00 / 0.00 / 0.00 / 0.00 | + +Three things this settles for round 2: + +1. `recall_bg_boundary` is 0.00-0.15 everywhere, so both heads did learn the *touching* target specifically and + ignore the background-facing rim - the target definition took, the confidence did not. +2. The head fires where merges are cheap (yeaz, dynamicnuclearnet, livecell: recall 0.63-0.76) and is nearly + silent exactly where the campaign lost mSA: tissuenet 0.19, neurips 0.13, deepbacs 0.02, dic_hepg2 0.001 + (3790 target pixels per crop, 40 predicted). Those losses therefore cannot come from the ridge - they come + from the shared features the extra task changed, as section 4.1 concluded. +3. The boundary-weighted foreground loss makes the head *less* confident, not more: every `both` recall is below + its `contact` counterpart (neurips 0.02 vs 0.13, puma 0.03 vs 0.12, tnbc and deepbacs to zero). The two loss + changes compete for the same decoder capacity. + +The round-2 target has a few percent of the pixels positive instead of under one, which is the structural +version of the "class-weighted or focal contact loss" lever of section 4.4 point 3: if under-confidence was +class imbalance, `boundary` fixes it, and its `recall_touching` on tissuenet / neurips / deepbacs / dic_hepg2 +is the number to look at. + +**Do not read the `fg_area_ratio` column of the summary CSVs**: it is a mean over datasets, and deepseas (12-91) +and neurips (2.3-12) dominate it because their crops carry few or tiny ground-truth objects. The per-dataset +column of `*_mechanisms.csv` is the readable one (baseline / fgcal / contact / both on deepbacs 1.19 / 1.25 / +1.40 / 1.25, tissuenet 0.71 / 0.75 / 0.77 / 0.74, dic_hepg2 1.04 / 1.11 / 1.16 / 1.12). + +### 4.7 Each decoder at its own sweep optimum, and the foreground threshold (19:00) + +The sweep rankings (`ais/reports/dec__sweep_dev.csv`, 1728 combinations, cached scorer, reference = +that decoder's library defaults) reproduce the screened full-pipeline runs to better than 0.05 %: baseline at +threshold 0.5 / density 50 / sigma 0.5 scores 0.4202 in the sweep against 0.4200 screened, fgcal 0.4298 against +0.4298, both 0.4254 against 0.4254. The sweep numbers below are therefore comparable to sections 4.2 / 4.3. + +| decoder | own optimum (dev balanced) | vs baseline's optimum | gain over its defaults | n_up | worst | fg threshold | density / sigma | +|---|---:|---:|---:|---|---|---:|---| +| baseline | 0.4244 | - | +2.4 % | 9 / 11 | -9.1 % | **0.4** | 10 / 1.0 | +| fgcal | 0.4298 | **+1.3 %** | +3.1 % | 6 / 11 | -6.4 % | 0.5 | 50 / 0.5 | +| both | 0.4254 | +0.2 % | +4.0 % | 8 / 11 | -8.2 % | 0.5 | 50 / 0.5 | +| contact | 0.4083 | **-3.8 %** | +3.3 % | 7 / 11 | -9.6 % | **0.6** | 10 / 1.0 | + +All four want the long travel (`n_iter` 800, `dt` 0.5), `foreground_weight` 0.75 and `min_size` 50; none passes +the gate; `boundary_magnitude_max` is irrelevant everywhere (0.4, 0.6 and off are within 0.001). + +Two things this changes: + +1. **The boundary-weighted foreground loss does calibrate the foreground, and the shared configuration hid it in + the opposite direction.** Every one of baseline's top 20 rows uses `foreground_threshold` 0.4; at 0.5 it only + reaches 0.4202 (+1.4 %). fgcal and both peak at 0.5. So the plain decoder needs its threshold lowered by a + tenth to reach its best, the boundary-calibrated ones are optimal at the natural 0.5 - which is exactly what + point 4.1 claims and what the area-ratio column was too coarse to show. Section 4.3 compared all four at + `dec-top1` (threshold 0.5), i.e. at fgcal's optimum and 1 % below baseline's, so the +2.4 % it reports for + fgcal is really **+1.3 %** (0.4298 against baseline's own 0.4244). Point 4.1 is a real but smaller effect, + and its mechanism is the threshold, not the merge share. +2. **The contact channel inflates the foreground, and the boundary-weighted BCE undoes it.** The optimal + threshold runs baseline 0.4 -> contact 0.6 -> fgcal / both 0.5. The field diagnostics say the same thing at a + fixed threshold: contact's `fg_area_ratio` at 0.5 is above baseline's on ten of eleven datasets (deepbacs + 1.21 vs 1.03, neurips 1.15 vs 1.03, tnbc 1.03 vs 0.88, puma 1.03 vs 0.91). The extra task pushes foreground + probability mass outward, and the calibrated loss pulls it back - which is why `both` sits between the two. +3. **The contact channel is a loss even at its own optimum.** Against baseline's own optimum, fgcal is +1.3 %, + both +0.2 % and contact **-3.8 %**. Section 4.2 measured -4.7 % at the shared defaults and 4.3 -4.2 % at + `dec-top1`; giving each decoder its best post-processing moves that by less than one point. The "it was only + mis-tuned" objection to section 4.4 point 3 is therefore closed: point 1.1 as implemented in round 1 loses. +4. **The "tuned regime moved" conclusion (4.4 point 4) is a property of the loss-changed decoders.** baseline + and contact keep the production density (10) and sigma (1.0) and only lengthen the travel; fgcal and both move + to density 50 / sigma 0.5. So the shift to "few, converged seeds" comes with the *foreground* loss change, + not with decoder fine-tuning as such. + +## 5. Round 2: the proper boundary channel (2026-09-07) + +Round 1 leaves point 1.1 undecided in the user's reading: the contact-only fifth channel (touching boundaries, +under 1 % of the pixels, ill-defined where three cells meet) is a dataset-dependent trade with an under-confident +head. Round 2 replaces it with the **classical boundary target**: the fifth channel holds the inner boundary of +every object, to neighbours and to background alike (`contact_mode="all"`, dilated by 1), which coincides with the +zero level set of the three geodesic distance channels the decoder already predicts, so the extra task no longer +asks for a quantity the other channels do not encode. + +| variant | fifth channel | foreground loss | +|---|---|---| +| `boundary` | inner boundary of every object, Dice + BCE | Dice (unchanged) | +| `boundary_fgcal` | same | Dice + boundary-weighted BCE (`boundary_weight=4`, radius 2) | + +`boundary` vs `baseline` isolates the channel, `boundary_fgcal` vs `fgcal` isolates it on top of the calibrated +foreground, and `boundary` vs `contact` isolates the target definition at a fixed loss. Everything downstream still +treats the channel as "contact" (sigmoid activation, `flow_instance_segmentation(contact=, contact_weight=, +contact_mask_threshold=)`, the `ais_contact_*.json` configs), so the round-1 readouts apply unchanged. + +### 5.1 Launch (17:16), and why the jobs first refused to start + +Submitted at 16:59 for `3g.40gb` slices with seven of the eight free, both jobs stayed `PENDING/WaitingInQueue` +for 17 minutes although slices, CPUs and memory were free on ggpu158 and ggpu192, and `sbatch --test-only` +claimed a start no earlier than 2026-09-08T03:13 for *every* pool (A100 on grete:shared, 3g, 2g, 1g on +preemptible) and independently of `--time`, `-c` and `--mem`. Cause: our own `dec_baseline_sweep_primary` array +sat at `TopOfQueue` on `grete:preemptible` with a marginally higher priority (103063 vs 103021), and an +unschedulable job at the head of the queue blocks the partition in the main scheduling loop for every +lower-priority job of the same user. `scontrol hold` on the four sweep arrays started both trainings within +seconds. The durable fix (the arrays only feed the sweep ranking, so they are the cheapest thing to delay): + +```bash +for j in 15776127 15776128 15776228 15776229; do scontrol update jobid=$j nice=100; done +``` + +which puts the sweeps below the rest of our chain (evaluation, finalisation, tuning launchers) while keeping them +ahead of the other user's queued preemptible job. Note for the next campaign: `--test-only` is worthless on +`grete:preemptible` because it ignores preemption - a sweep task started at 17:11 against a 03:14 estimate for +the same request. The only thing worth checking when a job does not start is whether one of our own arrays is at +the head of the queue. + +`boundary` = 15776831 (ggpu158), `boundary_fgcal` = 15776833 (ggpu192), both at 1.08 it/s for batch 8 (the +round-1 3g speed), so 48000 iterations plus the per-epoch validation land at 06:10-06:20 on 2026-09-08 inside +the 14 h limit (07:16). The first log lines confirm the target: `variant boundary: 5 output channels, loss +settings {'contact': True, 'contact_mode': 'all', 'boundary_weight': None}`. + +### 5.2 The round-1 finalisation died on an edited script + +`ais_decoder_finalize` (15772287) waited 7.4 h for the last screens, printed "all screens done" at 17:11 and then +aborted with `finalize_ais_decoder_reports.sh: line 44: syntax error near unexpected token 'done'`. The file is +syntactically fine; it had been edited at 16:53 while the job slept in the wait loop, and bash re-reads a running +script by byte offset, so the resumed parse landed mid-statement. None of the `decoders_final_*` tables or the +`baseline` / `contact` field diagnostics were written. Rerun as 15777315. **Rule from now on: submit a frozen +copy of every long-running driver**, `/jobs/frozen/_.sh`, never the repo path. + +The same 16:53 edit claimed a second job six hours later: `ais_decoder_tuning` (15772853, running since 12:28) +died at 19:03 with `break: only meaningful in a for, while or until loop` followed by +`syntax error near unexpected token 'done'` in `launch_tuning_after_caches.sh`, and its last log line is the +message of a branch it could not have reached - the signature of a shifted offset. Nothing was lost: it had +already submitted the `baseline` / `contact` sweeps at 16:03, and all four rankings exist (fgcal and both at +11:55 / 12:35, baseline and contact by hand at 18:12 / 18:57, section 4.7). Both files on disk pass `bash -n` +and the frozen copies under `/jobs/frozen/` are byte-identical to them, so `tuning2` and `finalize_r2` +are unaffected. **One edit to a driver can kill every job currently sleeping in it, hours apart.** + +### 5.3 The chain (nothing depends on the session) + +The session runs in a 12 h interactive job that ends at 05:06 on 2026-09-08, before the trainings do, so every +step is chained with SLURM dependencies. + +| job | what | starts | +|---|---|---| +| 15776831 / 15776833 | the two trainings | running since 17:16, done ~06:15 | +| 15776838 / 15776839 `ais_eval_` | `afterany` the training: stage, cache v5 primary / training_extra / holdout and apg3d primary / holdout, then the `current-defaults`, `contact-ridge` and `contact-mask` screens | ~06:15 | +| 15777359 `ais_decoder_tuning2` | `afterany` both evaluations (frozen `launch_tuning_after_caches.sh`): waits for the 2d caches, submits the two grid sweeps (1728 combinations) and the eight-configuration contact screen per variant, then ranks all six sweeps into `/ais/reports/dec__sweep_dev.csv` | ~06:20 | +| 15777357 `ais_decoder_finalize_r2` | `afterany` both evaluations (frozen `finalize_round2_reports.sh`): submits the `dec-top1` screens of the two new decoders `afterok` their prediction jobs, waits for every round-2 screen, then writes `decoders_all_defaults_{dev,holdout}`, `decoders_all_tuned_{dev,holdout}`, `decoders_all_3d`, `decoders__contact_dev` and the field diagnostics of both new decoders | ~06:20 | +| 15777315 `ais_decoder_finalize2` | the round-1 finalisation, rerun from a frozen copy | queued | + +`finalize_round2_reports.sh` is new (`finetuning/v2/generalist/ais_decoder/`); it replaces the manual "submit the +`dec-top1` screens once the caches exist, then run the section 5 commands" step of the hand-over, so the +successor only has to read the tables. + +### 5.4 What the two post-processing modes mean once the channel is a full boundary + +Both modes read the fifth channel unchanged (`micro_sam/v2/postprocessing.py`), but the target swap changes what +they do, which is worth stating before the numbers arrive: + +- `contact_weight` adds `w * contact` to the watershed height map. With the touching target the ridge sits only + between two objects; with the full boundary it also runs along every object's rim to the background. The + watershed is masked to the foreground, so a rim ridge mostly sits at the mask border and should be close to + inert - except that the target is dilated by one pixel, so the ridge reaches one pixel *inside* the object and + can shave structures only a few pixels wide (deepbacs rods, dic_hepg2 filaments). +- `contact_mask_threshold` excludes `contact > t` from the first seeded watershed and lets the instances claim + those pixels afterwards. With a full boundary this is no longer "keep the contact line free" but the classical + *erode, flood, dilate back* scheme: the first watershed runs on objects eroded by ~3 pixels. That should help + wherever objects touch, and it is the mode that was inert in round 1 only because the head rarely exceeded + 0.5 - a confident boundary head makes it active for the first time. The risk is the same one: an object thinner + than twice the band loses its interior entirely and can end up unseeded. + +So the expected signature of the boundary channel, if it works, is: mask mode finally moving the score, the +merge share falling on livecell / tissuenet / neurips, and a *new* kind of loss on the thin-object datasets - +which the mechanism columns separate (`seeded_split` and `gt_with_0_seeds` rather than `seeded_merged`). Both +modes are screened at 0.5 / 1 / 2 / 4 and 0.3 / 0.5 / 0.7 for each new decoder, so this is testable rather than +argued. + +### 5.5 `SBATCH_EXPORT=none` silently reverted the round-2 tuning to round 1 (2026-09-08, 06:15) + +Both trainings finished cleanly on the first attempt - `boundary` COMPLETED in 12:48:29 (48000 iterations, best +epoch 65 of 76, validation 0.786 at epoch 1 -> 0.576) and `boundary_fgcal` in 12:49 (best 0.804 from 1.076) - +and the evaluation chain staged both checkpoints with five output channels and submitted the caches, screens and +`dec-top1` screens as designed. + +`ais_decoder_tuning2` then did the wrong thing: at 06:13 it logged `caches of baseline ready, submitting sweeps` +and re-submitted the four **round-1** sweep arrays plus the 24-task `dec_contact_contact_screen`, and never +submitted anything for the two new decoders. Cause: **`SBATCH_EXPORT=none` is set in this environment** +(`echo $SBATCH_EXPORT`), so `sbatch` does not propagate the submitting environment and the `WAIT_VARIANTS` / +`VARIANTS` variables never reached the script, which fell back to its round-1 defaults. The note in the previous +hand-over - "`sbatch --export=ALL` is the default, so the two variables reach the script" - is wrong on this +system, so the original submission (15776840) carried the same latent bug; only the deliberate stop and restart +of the session caught it, because the failure is silent and produces plausible-looking work. + +Recovery (06:15-06:17), all of it visible in `/jobs/`: + +- cancelled the five redundant arrays (46 tasks) and `tuning2`, and moved their job directories to + `/jobs/_superseded/` so that `tasks_done` sees the completed round-1 directories as the newest again + (it reads `ls -td | head -1`, so an empty newer directory shadows a finished one); +- submitted by hand what the launcher should have: `dec_boundary_sweep_{primary,extra}` (15783735 / 15783736) + and `dec_boundary_contact_screen` (15783737) on the finished cache, and the same three for `boundary_fgcal` + (15783738 / 15783739 / 15783740) `afterok` its still-running `predict2d` array; +- `ais_rank_round2` (15783741) ranks both new sweeps `afterany` the four arrays. + +Fix in the repository: `launch_tuning_after_caches.sh` now takes the variants as arguments +(`--wait boundary boundary_fgcal --rank baseline contact ... boundary_fgcal`) and only falls back to the +environment when run directly in a shell. **Rule: never pass campaign parameters to a SLURM job through the +environment on this cluster** - put them in the command line or in the frozen script. + +### 5.6 Results of the boundary channel (2026-09-08, 06:30; `ais/reports/decoders_r2_early_*`) + +Both trainings ran the full budget on the first attempt: `boundary` COMPLETED in 12:48:29 (best epoch 65 of 76, +validation 0.786 at epoch 1 -> 0.576), `boundary_fgcal` in 12:49 (best 0.804 from 1.076). Checkpoints +`66368b4c` and `0753918a`, staged with five output channels. + +**1. The head is confident now - the class-imbalance diagnosis of 4.4 point 3 was right.** +`recall_touching` at threshold 0.5 (per-dataset medians, `decoder_fields_boundary_summary.csv`), round-1 +`contact` -> round-2 `boundary`: deepbacs 0.015 -> **0.505**, tnbc 0.032 -> **0.587**, puma 0.123 -> **0.566**, +neurips 0.128 -> **0.460**, tissuenet 0.192 -> 0.347, livecell 0.634 -> 0.724, dynamicnuclearnet 0.720 -> 0.899, +yeaz 0.757 -> 0.914, covid_if 0.591 -> 0.740. It learned the actual target rather than collapsing onto the +touching lines (`recall_bg_boundary` 0.43-0.90 against 0.00-0.15 for `contact`) and stayed precise (precision +within 2 px 0.68-0.98, Dice up to 0.87 on dynamicnuclearnet, 0.81 yeaz, 0.69 livecell). Exactly the four +datasets whose merges motivated the channel and whose head was silent in round 1 now fire. + +Two exceptions, and the first one taught us something about the instrument. On **`dic_hepg2`** the head has no +pixel above 0.5 on 33 of 50 crops (mean 13 predicted pixels against 8698 target pixels, per-crop maximum 0.41), +so every threshold-0.5 column calls it dead - but its *soft* probability is 0.133 on the true boundary against +0.0020 elsewhere, a 65-fold contrast, i.e. **well localised and merely under-confident** (deepbacs and livecell +run 0.65-0.69 against 0.0003). The consequence is visible in the scores: `contact_weight` adds +`w * contact` to the height map and therefore reads the soft map, so the ridge alone moves dic_hepg2 from +-13.8 % to +16.3 % against baseline at `dec-top1` - a 30-point swing out of a head that "predicts nothing". +`contact_mask_threshold` thresholds instead, and cannot use it. **Read the fifth channel's soft contrast, not +only its Dice and recall at 0.5**; the threshold columns understate a well-localised head. `deepseas` is +genuinely near-silent (Dice 0.056), as expected of binary masks with no true object boundaries. + +**2. Under the library defaults** (dev = 11 datasets, holdout = 5, reference = the fine-tuned `baseline`): + +| decoder (configuration) | dev balanced | vs baseline | up / 11 | holdout balanced | vs baseline | up / 5 | seeded merges dev | +|---|---:|---:|---|---:|---:|---|---:| +| `boundary` + ridge 1 | **0.4209** | **+1.5 %** | 7 | 0.3865 | -0.7 % | 3 | 4.5 % | +| `boundary` + mask 0.5 | 0.4201 | +1.3 % | 7 | 0.3855 | -1.0 % | 3 | 6.1 % | +| `boundary` defaults | 0.4192 | +1.1 % | 7 | 0.3838 | -1.4 % | 3 | 7.8 % | +| `fgcal` defaults | 0.4170 | +0.6 % | 6 | **0.3938** | **+1.1 %** | 4 | 8.3 % | +| `baseline` defaults | 0.4145 | - | - | 0.3894 | - | - | 8.1 % | +| `boundary_fgcal` defaults | 0.4124 | -0.5 % | 6 | 0.3746 | -3.8 % | 3 | 8.0 % | +| `both` defaults | 0.4090 | -1.3 % | 7 | 0.3819 | -1.9 % | 3 | 6.3 % | +| `contact` defaults | 0.3952 | -4.7 % | 4 | 0.3691 | -5.2 % | 2 | 7.8 % | + +**3. At the shared tuned configuration** (`dec-top1`, reference = `baseline` at `dec-top1` = 0.4200 dev / 0.4025 +holdout) the boundary channel gives **the best result of the whole campaign**: + +| decoder (configuration) | dev balanced | vs baseline | up / 11 | worst | holdout balanced | vs baseline | up / 5 | +|---|---:|---:|---|---|---:|---:|---| +| `boundary` + ridge 1 | **0.4340** | **+3.3 %** | **10** | deepbacs -14.1 % | 0.4085 | +1.5 % | 4 | +| `boundary` + ridge 2 + mask 0.3 | 0.4325 | +3.0 % | 10 | -15.0 % | 0.4069 | +1.1 % | 4 | +| `both` + ridge 1 | 0.4271 | +1.7 % | 8 | deepseas -48 % | **0.4113** | **+2.2 %** | 4 | +| `boundary` (no ridge) | 0.4254 | +1.3 % | 7 | -13.8 % | 0.4022 | -0.1 % | 3 | +| `contact` + ridge 1 | 0.4155 | -1.1 % | 6 | -26.5 % | 0.4044 | +0.5 % | 3 | + +**4. The round-1 collateral damage is repaired.** Per dataset at `dec-top1` + ridge 1 (dev), `contact` -> +`boundary`: covid_if -19.2 % -> **+0.2 %**, deepseas -26.5 % -> **+20.0 %**, dic_hepg2 +14.5 % -> +16.3 %, +dynamicnuclearnet -3.6 % -> +1.7 %, puma -1.3 % -> +2.1 %, tnbc +2.6 % -> +8.0 %, and the datasets the channel +was for stay up (livecell +10.7 %, tissuenet +9.1 %, yeaz +3.7 %, neurips +3.4 %). **Exactly one dataset is +down: deepbacs, -14.1 %** - and it is down by 13.7-16.4 % for `contact` and `both` too, so it is a property of +carrying a fifth channel at all, not of the target definition. + +**5. deepbacs is the predicted thin-object failure** (5.4), and its mechanism is visible: `boundary` + ridge 1 +against `baseline` at `dec-top1` splits more (`seeded_split` 2.8 % against 1.6 % of the objects), matches worse +(`matched_iou` 0.743 against 0.767) and above all over-covers (`fg_area_ratio` **1.36** against 1.19) - while +actually *improving* the two counts the channel targets (merges 2.8 % against 3.6 %, objects without a seed +5.2 % against 7.2 %). The rods are shaved and split, not merged. + +**6. Stacking the two loss changes still hurts**, as in round 1: `boundary_fgcal` is below `boundary` everywhere +(-0.5 % against +1.1 % dev, -3.8 % against -1.4 % holdout at the defaults) and its head is a few points less +confident (deepbacs 0.466 against 0.505, puma 0.486 against 0.566, tnbc 0.466 against 0.587). It does calibrate +the foreground it was meant to (`fg_area_ratio` deepbacs 1.06 against 1.17, neurips 1.03 against 1.07, deepseas +0.86 against 1.54) - but that did not buy mSA, and on deepbacs it made the score worse (-13.7 % against -7.7 % +at the defaults), so the over-coverage is not what costs deepbacs its score. + +**7. Gate verdict.** `boundary` + ridge 1 on dev: 10 of 11 datasets up, balanced +3.3 % (both bounds met), worst +-14.1 % against the -2 % bound - **it fails the gate on deepbacs alone**. On the holdout it is +1.5 % (4 of 5), +where `both` + ridge 1 reaches +2.2 %. So the proper boundary target turns point 1.1 from a broad +dataset-dependent trade (round 1: 4-6 datasets down, up to -38 %) into a broad gain with one identified, +channel-generic failure. That is a qualitatively different object from round 1 and the first version of the +fifth channel worth carrying further, but it is not yet a pass. + +Still running at the time of writing: the 24-configuration contact screens of both new decoders, the +`dec-top1` screens of `boundary_fgcal`, the four grid sweeps and their ranking (`ais_rank_round2`, 15783741), and +the 3d screens. The sweep will say which `foreground_threshold` the boundary decoders want, which is the test of +4.7 point 1 (`boundary`'s median `fg_area_ratio` is above `baseline`'s on nine of eleven datasets, so 0.5-0.6 is +the expectation). + +**8. The ridge and the mask separate cleanly** (`decoders_boundary_contact_{dev,holdout}*.csv`, `boundary` +against its own defaults 0.4192 dev / 0.3838 holdout, all other parameters at the library defaults): + +| configuration | dev | vs defaults | up / 11 | worst | holdout | vs defaults | seeded merges dev | seeded splits dev | +|---|---:|---:|---|---|---:|---:|---:|---:| +| defaults | 0.4192 | - | - | - | 0.3838 | - | 7.8 % | 1.80 % | +| ridge 0.5 | 0.4207 | +0.35 % | 4 | -1.0 % | **0.3874** | **+0.95 %** | 4.8 % | 2.03 % | +| ridge 1 | 0.4209 | +0.41 % | 4 | -2.1 % | 0.3865 | +0.71 % | 4.5 % | 2.08 % | +| ridge 2 | **0.4214** | **+0.51 %** | 4 | -3.1 % | 0.3857 | +0.51 % | 4.3 % | 2.17 % | +| ridge 4 | 0.4212 | +0.48 % | 4 | -3.3 % | 0.3853 | +0.39 % | 4.3 % | 2.21 % | +| mask 0.3 | 0.4201 | +0.21 % | **7** | **-0.5 %** | 0.3852 | +0.36 % | 5.5 % | 1.86 % | +| mask 0.5 | 0.4201 | +0.20 % | 5 | -0.3 % | 0.3855 | +0.46 % | 6.1 % | 1.79 % | +| mask 0.7 | 0.4198 | +0.13 % | 7 | -0.0 % | 0.3847 | +0.24 % | 7.0 % | 1.67 % | +| ridge 1 + mask 0.5 | 0.4209 | +0.41 % | 4 | -2.1 % | 0.3865 | +0.71 % | 4.5 % | 2.10 % | + +This revises the prediction of 5.4 in one respect and confirms it in another. The mask mode *does* move the +score now that the head is confident (+0.2 % dev, +0.36-0.46 % holdout, against 0.0-0.1 % in round 1), and it is +by far the more **uniform** lever: 7 of 11 datasets up with a worst case of -0.5 %, against the ridge's 4 of 11 +and -1.0 to -3.3 %. But the shaving of thin objects is a **ridge** effect, not a mask effect: `seeded_split` +rises monotonically with the ridge weight (1.80 % -> 2.21 %) and stays flat or falls under the mask +(1.67-1.86 %) - exactly as the erode-*and-dilate-back* structure of the mask mode implies, which is the half of +5.4's reasoning that was right. The ridge still wins on balanced mSA because it removes almost twice as many +merges (7.8 % -> 4.3 % against 5.5-7.0 %) and because it can exploit an under-confident head (point 1). + +**9. No post-processing setting can rescue deepbacs.** `boundary` is already -11.7 % there at `dec-top1` with +neither ridge nor mask, and its `fg_area_ratio` of 1.359 (baseline 1.185) is a property of the decoder, not of +the watershed. The ridge adds 2 points of loss on top (-14.1 %); the loss itself is in the field. So the gate +failure of point 7 is a training-recipe question (deepbacs' thin rods need the foreground calibrated, and +`boundary_fgcal` - which does calibrate it, 1.06 against 1.17 - scores *worse* there, -13.7 % against -7.7 % at +the defaults), not a tuning question. + +**10. The threshold test of 4.7 point 1: the full boundary does not inflate the foreground.** +`boundary`'s own sweep optimum (`dec_boundary_sweep_dev.csv`, 1728 combinations) is 0.4273 at +`foreground_threshold` **0.5** (0.4 gives 0.4252, 0.6 gives 0.4246, 0.7 gives 0.4158), so the optimal threshold +runs `baseline` 0.4 -> **`boundary` 0.5** -> `contact` 0.6. The touching target pushed foreground mass outward; +the full boundary does so far more mildly, and the section 5.6 expectation of "0.5-0.6" lands at the benign end. +The optimum also confirms 4.7 point 4: `boundary` keeps the *production* density (10) and sigma (1.0) and only +lengthens the travel to 800, exactly like `baseline` and `contact`, whereas `fgcal` and `both` - the two that +changed the *foreground* loss - move to density 50 / sigma 0.5. The regime shift belongs to the foreground loss, +not to the fifth channel. + +**11. Restating the headline honestly.** The +3.3 % of point 3 is measured against `baseline` at `dec-top1` +(0.4200), which is 1 % below baseline's own optimum (0.4244, 4.7) - the same overstatement that 4.7 caught for +`fgcal`. Against baseline's own optimum, `boundary` + ridge 1 at `dec-top1` is **+2.3 %**. The sweep cannot +settle this by itself because the cached scorer ignores the contact keywords, so `boundary`'s own optimum +(0.4273, +0.7 % over baseline's own optimum) is a *ridge-free* number and understates the decoder as much as +`dec-top1` overstates it. Screens of the missing cells were submitted at 06:42: `dec-base-top1` for `baseline` +(job dec_baseline_own_screen) and `dec-bnd-top1` / `dec-bnd-top1-ridge1` for `boundary` +(dec_boundary_own_screen), i.e. each decoder at its own sweep optimum, with and without the ridge. + +**12. The corrected comparison: every decoder against `baseline` at ITS OWN optimum** (screens +`dec_baseline_own_screen` / `dec_boundary_own_screen`, `ais/reports/decoders_own_optimum_*`). `baseline` at +`dec-base-top1` scores **0.4244 dev / 0.4052 holdout**, against 0.4200 / 0.4025 at `dec-top1`. Recomputing every +candidate's best configuration against that reference: + +| decoder (best configuration) | dev | vs baseline's own optimum | holdout | vs baseline's own optimum | +|---|---:|---:|---:|---:| +| `boundary` + ridge 1 @ `dec-top1` | 0.4340 | **+2.3 %** | 0.4085 | +0.8 % | +| `boundary_fgcal` + ridge 1 @ `dec-top1` | 0.4311 | +1.6 % | 0.4082 | +0.7 % | +| `fgcal` @ `dec-fgcal-top1` | 0.4298 | +1.3 % | 0.4094 | +1.0 % | +| `both` + ridge 1 @ `dec-top1` | 0.4271 | +0.6 % | **0.4113** | **+1.5 %** | +| `baseline` @ `dec-base-top1` | 0.4244 | - | 0.4052 | - | +| `contact` + ridge 1 @ `dec-top1` | 0.4155 | -2.1 % | 0.4044 | -0.2 % | + +**This retracts the "10 of 11 datasets up" of point 3.** Against `baseline` at its own optimum, `boundary` + +ridge 1 has **7 of 11 up on dev** and 3 of 5 on the holdout, with four datasets down: deepbacs -10.2 %, +neurips -6.1 %, tissuenet -5.0 %, puma -0.3 % (up: deepseas +28.1 %, dic_hepg2 +26.5 %, tnbc +6.5 %, covid_if ++4.2 %, livecell +4.2 %, yeaz +3.3 %, dynamicnuclearnet +1.0 %). tissuenet alone swings 14 points +(+9.1 % -> -5.0 %) purely from the reference, because `baseline` at threshold 0.4 / density 10 / sigma 1.0 is far +better there than at `dec-top1`. The lesson of 4.7 therefore applies to round 2 in full: **a shared tuned +configuration flatters whichever decoder it was tuned on**, and the only defensible reference is each decoder at +its own optimum. + +What survives the correction: the target change is worth **+4.4 points** over round 1 (`contact` -2.1 % -> +`boundary` +2.3 % on dev, both at their best configuration against the same reference), the head is confident +(point 1), and the collateral damage on the unseen datasets is repaired (covid_if +4.2 %, deepseas +28.1 %). +What does not: the dev gain is +2.3 % rather than +3.3 %, it does not confirm on the holdout (+0.8 %, where +round-1 `both` reaches +1.5 %), and four datasets are down rather than one. Under the user's rule - only +cross-dataset wins that hold on the holdout count - **the boundary channel is a real improvement over the +contact channel but still not a win over the plain fine-tune**, and no configuration of any of the six decoders +passes the gate. + +**13. `boundary`'s own sweep optimum is not its best configuration once the ridge exists.** +`dec-bnd-top1` (threshold 0.5, density 10, sigma 1.0) scores 0.4273 / 0.4000 and with ridge 1 0.4269 / 0.4007, +against 0.4340 / 0.4085 for `dec-top1` + ridge 1 (density 50, sigma 0.5). The ridge and the seed regime +interact: the ridge pays off in the few-converged-seeds regime, and the sweep - which cannot evaluate the contact +keywords at all - therefore optimises into the wrong basin. dic_hepg2 shows it starkly: -8.6 % at +`dec-bnd-top1-ridge1` against +26.5 % at `dec-top1-ridge1`. **A ridge-blind sweep cannot tune a five-channel +decoder**; the grid needs `contact_weight` as a dimension, which requires teaching the cached scorer the contact +keywords. + +**14. `boundary_fgcal`'s ridge and mask** (`decoders_boundary_fgcal_contact_*`): the ridge is worth at most ++0.13 % (dev, w0.5) and the mask +0.24 % (dev, t0.5) / +0.28 % (holdout) against its own defaults - an order of +magnitude less than for `boundary`, and higher ridge weights *hurt* (-0.35 % at w4). Its foreground is already +calibrated, so the seeds it would gain from a ridge are largely there; consistent with point 6. + +**15. All six sweep optima, and what each parameter tracks** (`dec__sweep_dev.csv`, 1728 combinations +each, cached scorer, reference = that decoder's own library defaults; `n_iter` 800, `dt` 0.5, +`foreground_weight` 0.75 and `min_size` 50 everywhere): + +| decoder | own optimum (dev) | `foreground_threshold` | density / sigma | fifth channel | foreground loss | +|---|---:|---:|---|---|---| +| `baseline` | 0.4244 | **0.4** | 10 / 1.0 | - | Dice | +| `contact` | 0.4083 | **0.6** | 10 / 1.0 | touching | Dice | +| `boundary` | 0.4273 | **0.5** | 10 / 1.0 | full boundary | Dice | +| `fgcal` | **0.4298** | 0.5 | **50 / 0.5** | - | Dice + boundary BCE | +| `both` | 0.4254 | 0.5 | **50 / 0.5** | touching | Dice + boundary BCE | +| `boundary_fgcal` | 0.4261 | 0.5 | **50 / 0.5** | full boundary | Dice + boundary BCE | + +The two parameters separate the two loss changes with no exceptions across six decoders: + +- **`density_threshold` / `sigma` track the foreground loss alone.** All three decoders trained with the + boundary-weighted foreground BCE want density 50 / sigma 0.5; all three without it want the production + density 10 / sigma 1.0. The fifth channel has no influence. This settles 4.7 point 4: the move to the + "few, converged seeds" regime is caused by the foreground loss, not by decoder fine-tuning and not by the + extra channel. +- **`foreground_threshold` tracks the fifth channel's target.** No channel 0.4, touching boundaries 0.6, full + boundaries 0.5 - i.e. the auxiliary task pushes foreground probability mass outward in proportion to how + ill-posed it is, and the calibrated foreground loss pins the threshold at 0.5 whatever the channel does + (`fgcal`, `both` and `boundary_fgcal` all 0.5). + +Ranking at each decoder's own **ridge-free** optimum: `fgcal` 0.4298 > `boundary` 0.4273 > `boundary_fgcal` +0.4261 > `both` 0.4254 > `baseline` 0.4244 > `contact` 0.4083. So without the contact ridge the +boundary-weighted foreground loss is the best single change, and the fifth channel only overtakes it once the +ridge is available (point 12) - which the sweep cannot see (point 13). + +**16. The 3d crops: the boundary channel does not repair the volume path, and its LM failure is the foreground** +(`ais/reports/decoders_all_3d*`, apg3d primary + holdout, 75 crops; regression instrument only - the decoders +were fine-tuned on 2d LM data and the 3d path saw none). Per family under `current-defaults`: + +| family | production | baseline | fgcal | boundary | boundary_fgcal | contact | both | +|---|---:|---:|---:|---:|---:|---:|---:| +| celegans_atlas | 0.104 | 0.040 | 0.011 | 0.000 | 0.000 | 0.000 | 0.000 | +| embedseg_platy_ish | 0.339 | 0.156 | 0.135 | 0.000 | 0.002 | 0.000 | 0.000 | +| embedseg_platy_nuclei | 0.259 | 0.115 | 0.086 | 0.000 | 0.000 | 0.000 | 0.000 | +| embedseg_skull | 0.118 | 0.238 | 0.078 | 0.000 | 0.009 | 0.000 | 0.000 | +| gonuclear | 0.256 | 0.132 | 0.135 | 0.000 | 0.004 | 0.000 | 0.000 | +| platynereis_nuclei | 0.068 | 0.052 | 0.006 | 0.000 | 0.000 | 0.000 | 0.000 | +| cremi / cremi_seen (lower better) | 0.99 / 0.59 | 1.87 / 1.23 | 2.24 / 2.15 | **1.86 / 1.42** | 1.97 / 2.04 | 2.19 / 2.05 | 2.09 / 1.89 | +| snemi / humanneurons (lower better) | 0.97 / 1.35 | 1.69 / 1.97 | 1.95 / 1.98 | **1.87** / 2.03 | 2.02 / 2.16 | 2.16 / 2.37 | 1.92 / 2.04 | + +Every five-channel decoder scores exactly 0 on all six LM families, the boundary target included, so the +better-posed channel does **not** repair the volume path. But the mechanism is not the one recorded for round 1 +in 4.0 (`boundary_magnitude_max` removing every instance): the mechanism columns show `boundary`'s 3d +**foreground ballooning** - `fg_area_ratio` 6.63 on celegans_atlas and **8.45** on gonuclear, against 2.04 / 3.59 +for `baseline` and 1.28 / 2.02 for production - with 2.7 to 9.1 background seeds per ground-truth object and +`matched_iou` undefined because nothing matches at IoU 0.5 at all. Objects are not missing for want of seeds +(`gt_with_0_seeds` 0.29-0.39, no worse than baseline); the volume is simply flooded. `boundary` is the *worst* +of the six on this measure, i.e. the extra 2d task makes the untrained 3d foreground worse the better it is +learned in 2d. + +Two things worth carrying to a joint 2d + 3d run: + +- **`fgcal` is the only variant that improves the 3d foreground** (gonuclear `fg_area_ratio` 2.68 against + baseline's 3.59, celegans 2.38 against 2.04 - and it is the only loss change that keeps an LM score at + baseline level, gonuclear 0.135 against 0.132). The boundary-weighted foreground BCE generalises to the + dimension it never saw; the fifth channel does the opposite. +- **On EM the boundary channel is harmless**: `boundary` matches `baseline` on cremi (1.86 against 1.87) and is + the best of the six on cremi_seen (1.42) and snemi (1.87), while `fgcal` is the worst on cremi (2.24). The + volume regression is specific to LM instance matching, not to volumes as such. + +## 6. Conclusive overview of the six decoders (2026-09-08, 07:20) + +Six decoders, identical data, budget (48000 iterations, batch 8) and initialisation (the v4 joint weights, the +image encoder frozen), differing only in the loss. Every figure below is **each decoder at its own best +configuration against `baseline` at its own best configuration** (0.4244 dev, 0.4052 holdout) - the reference +that sections 4.7 and 5.6 point 12 show to be the only defensible one. + +| decoder | fifth channel | foreground loss | dev | holdout | verdict | +|---|---|---|---:|---:|---| +| `boundary` | inner boundary of every object | Dice | **+2.3 %** | +0.8 % | best on dev, does not confirm | +| `boundary_fgcal` | same | Dice + boundary BCE | +1.6 % | +0.7 % | strictly below `boundary` | +| `fgcal` | - | Dice + boundary BCE | +1.3 % | +1.0 % | small, consistent, safest worst case | +| `both` | touching boundaries | Dice + boundary BCE | +0.6 % | **+1.5 %** | best on holdout, -48 % on deepseas | +| `baseline` | - | Dice | - | - | the reference | +| `contact` | touching boundaries | Dice | -2.1 % | -0.2 % | loses | +| production (v4) | - | - | -19.0 % | -39.9 % | not comparable (no fine-tune) | + +**1. In-domain data still dominates everything** (4.4 point 1, unrevised). The plain fine-tune with the +unchanged loss lifts the decoder from 0.3437 to 0.4244 on dev (+23 %) and from 0.2437 to 0.4052 on the holdout +(+66 %); the best loss change on top of that is worth +2.3 %, an order of magnitude less. For the next big run +the composition of the training data matters far more than either loss change. + +**2. Point 4.1 (boundary-weighted foreground BCE) - include it.** +1.3 % dev / +1.0 % holdout, the most +*consistent* of the changes (it is the only candidate whose worst dataset stays within -5.1 %, against -10 to +-48 % for every fifth-channel variant), and three independent mechanisms now explain it: it moves the optimal +`foreground_threshold` from 0.4 to the natural 0.5 (5.6 point 15), it is the sole cause of the tuned regime's +shift to density 50 / sigma 0.5 (5.6 point 15, six decoders with no exceptions), and it is **the only change +that generalises to the dimension it never saw** - the only variant that improves the 3d foreground and keeps +an LM volume score at baseline level (5.6 point 16). It still misses the gate's worst-loss bound, so it is not +a "win" under the strict rule, but it is cheap, safe and mechanistically understood. + +**3. Point 1.1 (the fifth channel) - the target definition was the whole question, and the answer is "better, +not yet good".** The touching target is unusable (-2.1 % dev, a head that never fires on four datasets, -21 to +-38 % collateral); the full inner boundary turns that into +2.3 % dev with a confident head (recall 0.35-0.91 +against 0.001-0.76) and repairs the collateral damage on the two unseen datasets (covid_if -19.2 % -> +4.2 %, +deepseas -26.5 % -> +28.1 %). What still blocks it: + +- it **does not confirm on the holdout** (+0.8 % against `both`'s +1.5 %) and has four datasets down on dev; +- **deepbacs -10.2 %**, a thin-object failure no post-processing can reach (5.6 point 9): the rods are shaved + and split and the 2d foreground over-covers (1.36 against baseline's 1.19). Every five-channel decoder loses + 10-16 % there; +- it **floods the 3d LM foreground** (`fg_area_ratio` up to 8.4 against baseline's 3.6) and is the worst of the + six on that measure, i.e. the better the channel is learned in 2d the worse the untrained 3d foreground gets; +- its gain **depends on the contact ridge**, which the tuning grid structurally cannot see (5.6 point 13), so + it cannot currently be tuned honestly alongside the other parameters. + +**Recommendation for the next big run**: include the boundary-weighted foreground loss; include the +full-boundary fifth channel **only** together with (a) joint 2d + 3d training, without which its volume +foreground is unusable, (b) `contact_weight` as a dimension of the tuning grid, and (c) a remedy for thin +objects - `boundary_fgcal` is not it (it calibrates the foreground but scores *worse* on deepbacs). Do not use +the touching-boundary target under any circumstances. + +**4. Method lessons that outlived the experiment.** + +- **Compare every decoder at its own optimum.** A shared tuned configuration flatters whichever decoder it was + tuned on: it inflated `fgcal` from +1.3 % to +2.4 % (4.7) and `boundary` from +2.3 % to +3.3 % with a + spurious "10 of 11 datasets up" (5.6 point 12), where tissuenet alone swung 14 points from the reference. +- **Read the auxiliary head's soft contrast, not its recall at 0.5.** On dic_hepg2 the boundary head has no + pixel above 0.5 on 33 of 50 crops yet separates boundary from background 65-fold in probability, and the + ridge - which reads the soft map - extracts a 30-point mSA swing from it (5.6 point 1). +- **A ridge-blind grid cannot tune a five-channel decoder** (5.6 point 13); `boundary`'s own sweep optimum is + the wrong basin once the ridge exists (dic_hepg2 -8.6 % against +26.5 %). +- **`SBATCH_EXPORT=none` on this cluster** silently reverts environment-passed campaign parameters to their + defaults (5.5), and **editing a driver script kills every job sleeping in it**, hours apart (5.2). + +## 7. The sweep results, cleanly (2026-09-08, 08:40; dic_hepg2 excluded) + +**The exclusion.** dic_hepg2 is dropped from every figure in this section on the user's instruction: its absolute +mSA is near the floor for every fine-tuned decoder (0.118-0.190 at the library defaults, against 0.25-0.84 for +eight of the other ten datasets) while its spread across the six decoders is 0.072 - so a 0.07 absolute wobble +becomes a +-40 % relative swing that dominates the balanced mean and the gate counts without representing +segmentation quality. The remaining ten development datasets are livecell, tissuenet, dynamicnuclearnet, +deepbacs, yeaz, neurips_cellseg, deepseas, puma, tnbc, covid_if; the holdout keeps four (livecell, tissuenet, +dynamicnuclearnet, deepbacs). Files: `ais/reports/dec__sweep_dev_no_dic.csv`, +`decoders_all_tuned_{dev,holdout}_no_dic*`, `decoders_own_optimum_dev_no_dic*`. +(deepseas has the same pathology - absolute mSA 0.046-0.112 with a 0.066 spread - and is kept here only because +it was not part of the instruction; a successor may want to drop it on the same grounds.) + +### 7.1 The optimum of each decoder (1728 combinations, ten datasets) + +Every optimum uses `n_iter` 800, `dt` 0.5 and `min_size` 50 (`boundary_fgcal`: 25), and `boundary_magnitude_max` +0.4 - which is irrelevant everywhere (0.4, 0.6 and off differ by less than 1e-3). + +| decoder | own optimum | vs `baseline`'s optimum | `foreground_threshold` | density / sigma | fg weight | up / 10 | worst | mean ratio to the per-dataset optimum | +|---|---:|---:|---:|---|---:|---|---|---:| +| `boundary` | **0.4502** | **+1.17 %** | 0.5 | 20 / 0.5 | 0.75 | 6 | -12.3 % | 0.944 | +| `fgcal` | 0.4472 | +0.50 % | 0.5 | 10 / 1.0 | 0.75 | 7 | -5.0 % | 0.942 | +| `boundary_fgcal` | 0.4454 | +0.09 % | 0.5 | 20 / 0.5 | 0.50 | 5 | -13.2 % | 0.931 | +| `baseline` | 0.4450 | - | **0.4** | 10 / 1.0 | 0.75 | 8 | -9.1 % | 0.950 | +| `both` | 0.4422 | -0.62 % | 0.5 | 20 / 0.5 | 0.50 | 6 | -9.8 % | 0.928 | +| `contact` | 0.4309 | -3.17 % | **0.6** | 10 / 1.0 | 0.50 | 7 | -11.3 % | 0.934 | + +No combination of any decoder passes the gate. Note the reordering against the eleven-dataset table of 5.6 +point 15: `boundary` now leads the ridge-free comparison (+1.17 %) instead of `fgcal`, and `both` drops below +`baseline`. + +### 7.2 What the sweep actually determines: read the plateau, not the top row + +Best balanced score per `foreground_threshold`, all other parameters free, as a loss in 1e-3 against each +decoder's own best threshold: + +| decoder | 0.4 | 0.5 | 0.6 | 0.7 | argmax | +|---|---:|---:|---:|---:|---:| +| `baseline` | **0** | -5.3 | -19.5 | -37.0 | **0.4** | +| `contact` | -13.2 | -3.7 | **0** | -5.0 | **0.6** | +| `fgcal` | -3.3 | **0** | -6.6 | -22.8 | 0.5 | +| `both` | -4.0 | **0** | -6.1 | -21.5 | 0.5 | +| `boundary` | -1.8 | **0** | -3.3 | -13.1 | 0.5 | +| `boundary_fgcal` | -0.1 | **0** | -7.4 | -28.4 | 0.5 | + +Same treatment for the seed regime, at each decoder's own best threshold: + +| decoder | best (d/sigma) | second | third | spread | +|---|---|---|---|---:| +| `baseline` | 10 / 1.0 | 50 / 0.5 (-2.4) | 20 / 0.5 (-2.8) | 2.8 | +| `contact` | 10 / 1.0 | 20 / 0.5 (-1.4) | 50 / 0.5 (-3.7) | 3.7 | +| `fgcal` | 10 / 1.0 | 20 / 0.5 (**-0.0**) | 50 / 0.5 (-1.4) | 1.4 | +| `both` | 20 / 0.5 | 50 / 0.5 (-1.0) | 10 / 1.0 (-1.2) | 1.2 | +| `boundary` | 20 / 0.5 | 10 / 1.0 (**-0.2**) | 10 / 0.5 (-1.2) | 1.2 | +| `boundary_fgcal` | 20 / 0.5 | 10 / 1.0 (**-0.0**) | 50 / 0.5 (-0.6) | 0.6 | + +**This retracts 5.6 point 15's second claim.** The "clean 3-3 separation of the seed regime by the foreground +loss, with no exceptions" was an artefact of dic_hepg2 plus reading a single top row: with dic_hepg2 removed the +two regimes are *identical to four decimals* for `fgcal` and `boundary_fgcal` and 0.2e-3 apart for `boundary`. +The seed regime is not determined by the loss - the sweep simply cannot distinguish density 10 / sigma 1.0 from +density 20 / sigma 0.5 for these decoders, and any claim built on which of the two the top row happened to pick +is noise. + +**The threshold claim survives, and it is the one real finding of the sweeps.** Its effects are 5 to 37e-3, an +order of magnitude above the regime differences, and the two informative contrasts are unambiguous: `baseline` +loses 19.5e-3 if forced to 0.6, and `contact` loses 13.2e-3 if forced to 0.4. So **the fifth channel shifts the +optimal foreground threshold, in proportion to how ill-posed its target is**: no channel 0.4, full inner +boundary 0.5, touching boundaries 0.6. (`boundary_fgcal` sits on a 0.4/0.5 plateau, the one soft case.) + +### 7.3 The screened comparison on the same ten datasets + +| comparison | dev | holdout | +|---|---:|---:| +| `baseline` at its own optimum | 0.4450 | - | +| `baseline` at `dec-top1` | 0.4382 (-1.5 %) | 0.4480 | +| `boundary` at its own optimum (`dec-bnd-top1`, ridge-free) | **0.4500 (+1.13 %)** | - | +| `boundary` at its own optimum + ridge 1 | 0.4495 (+1.02 %) | - | +| `boundary` at `dec-top1` + ridge 1 | 0.4497 (+1.05 %) | 0.4540 | +| `fgcal` at `dec-fgcal-top1` | 0.4458 (+0.18 %) | **0.4551** | + +**This retracts 5.6 point 13.** "A ridge-blind sweep cannot tune a five-channel decoder" rested entirely on +dic_hepg2: with it excluded, `boundary`'s own sweep optimum is its **best** configuration (0.4500), the ridge +adds nothing there (0.4495 with it), and at `dec-top1` the ridge is worth +0.5 % rather than the +2 % the +eleven-dataset table showed. The +26.5 % dic_hepg2 gain that made the ridge look essential was a swing on a +0.15-mSA dataset. The contact ridge is a small, real improvement in the tuned regime - not the thing that makes +the fifth channel work. + +**Net effect on the verdict of section 6.** On the ten datasets, at each decoder's own optimum against +`baseline` at its own optimum: `boundary` **+1.2 %** dev, `fgcal` +0.5 %, `boundary_fgcal` +0.1 %, `both` +-0.6 %, `contact` -3.2 %; on the four holdout datasets `fgcal` leads (+1.6 % against `boundary`'s +1.3 %, +referenced to `baseline` at `dec-top1`). The direction of section 6 is unchanged - the full-boundary target is +far better than the touching target, and no change passes the gate - but the boundary channel's dev advantage is +**+1.2 %, not +2.3 %**, and `fgcal` remains the change that holds up best on unseen data. + +## 8. The sweep optima on the nine informative datasets, and what they mean mechanically (08:45) + +`deepseas` is dropped as well as `dic_hepg2`, on the same grounds (absolute mSA 0.046-0.112, spread 0.066, so +its relative changes are floor noise). The nine remaining development datasets are livecell, tissuenet, +dynamicnuclearnet, deepbacs, yeaz, neurips_cellseg, puma, tnbc, covid_if. Files: +`ais/reports/dec__sweep_dev_core9.csv`. + +### 8.1 The optima + +Every optimum uses `n_iter` 800, `dt` 0.5, `boundary_magnitude_max` 0.4 and `seed_floor` "none". + +| decoder | balanced | vs `baseline` | `foreground_threshold` | `density_threshold` | `sigma` | `foreground_weight` | `min_size` | up / 9 | worst | +|---|---:|---:|---:|---:|---:|---:|---:|---|---| +| `boundary` | **0.4893** | +0.81 % | 0.5 | 20 | 0.5 | 0.75 | 50 | 6 | **-2.1 %** | +| `boundary_fgcal` | 0.4889 | +0.72 % | 0.5 | 20 | 0.5 | 0.50 | 25 | 5 | -6.6 % | +| `both` | 0.4868 | +0.28 % | 0.5 | 20 | 0.5 | 0.50 | 50 | 6 | -5.9 % | +| `fgcal` | 0.4865 | +0.23 % | 0.5 | 20 | 0.5 | 0.75 | 50 | 7 | -4.5 % | +| `baseline` | 0.4854 | - | **0.4** | 10 | 1.0 | 0.75 | 50 | 8 | -4.7 % | +| `contact` | 0.4722 | -2.72 % | **0.6** | 10 | 1.0 | 0.50 | 50 | 7 | -5.7 % | + +Removing the two floor-level datasets collapses the spread: the four "improved" decoders now sit **+0.2 % to ++0.8 %** above `baseline`, i.e. inside the band the sweep itself cannot resolve, and only `contact` is clearly +worse (-2.7 %). The one figure that improves markedly is `boundary`'s worst dataset, -12.3 % on ten datasets -> +**-2.1 %** on nine, because deepbacs is now its only real loss. Nothing passes the gate (`boundary` needs 7 of 9 +up and has 6, and +0.81 % against the +2 % bound). + +The two plateau checks of 7.2 are unchanged by the second exclusion: the `foreground_threshold` pattern holds +with the same large margins (`baseline` loses 19.2e-3 if forced to 0.6, `contact` 15.3e-3 if forced to 0.4, all +others prefer 0.5), and the seed regime remains unresolvable (the runner-up is within 0.1-1.3e-3 for the four +decoders that pick density 20 / sigma 0.5). + +### 8.2 How the two maps are computed (`micro_sam/v2/postprocessing.py::flow_instance_segmentation`) + +Both maps are built from the same two predicted quantities: the foreground probability `p` and the three +directed distance channels `d` (magnitude `|d|`, which dips to zero at object centres and at boundaries). + +**Seed map** - four steps, and every sweep parameter but two acts here: + +1. `fg = p > foreground_threshold` selects the pixels that take part. +2. Each `fg` pixel is advected along `-d` for `n_iter` steps of length `dt`; pixels of one object flow to its + centre. The per-pixel count of arrivals is the convergence density. +3. The density is Gaussian-smoothed with `sigma`. +4. `seeds = connected components of (density > density_threshold)`. + +So `n_iter x dt` is the travel budget (400 px for every optimum here, against 25 px in the library defaults), +`sigma` sets how far apart two convergence points may be and still merge into one seed, and +`density_threshold` sets how many arrivals a seed must collect - together they trade missed objects against +split ones. + +**Height map** - two parameters, one of them not in the grid: + +``` +h = foreground_weight * (1 - p) + (1 - foreground_weight) * (1 - |d| / max|d|) [+ contact_weight * contact] +``` + +i.e. a convex mix of the foreground's complement (a sharp edge signal) and the inverted, max-normalised distance +magnitude (a weak edge signal that also dips at object centres). `seed_floor` is "none" in every configuration +here, so the height under the seeds is left as it is. The watershed then floods `h` from `seeds` within `fg`; +`min_size` removes small instances and re-floods, and `boundary_magnitude_max` finally drops instances whose +median boundary magnitude does not dip. + +### 8.3 Each setting's best configuration, in words + +All six use the same 400-px travel budget and the same `boundary_magnitude_max` 0.4. + +- **`baseline`** (4 channels, Dice foreground). *Seeds*: the widest foreground, `p > 0.4`, advected and smoothed + with the broad `sigma` 1.0, seeds where at least 10 arrivals land. *Height*: 0.75 foreground + 0.25 inverted + magnitude. The only decoder that needs its foreground threshold lowered below 0.5, and the only one whose + seeding stays in the library's broad-smoothing regime. +- **`contact`** (5 channels, touching boundaries, Dice foreground). *Seeds*: the narrowest foreground, `p > 0.6` + - its foreground is inflated, so it must be cut back harder - otherwise identical to `baseline` (sigma 1.0, + density 10). *Height*: 0.5 foreground + 0.5 inverted magnitude, i.e. it leans more on the distance field than + `baseline` does. +- **`boundary`** (5 channels, full inner boundary, Dice foreground). *Seeds*: `p > 0.5`, sharp smoothing + (`sigma` 0.5) and a doubled density threshold (20), i.e. fewer, tighter, better-converged seeds. *Height*: + 0.75 foreground + 0.25 inverted magnitude, like `baseline`. The best of the six here. +- **`fgcal`** (4 channels, boundary-weighted foreground BCE). *Seeds*: `p > 0.5` - the calibrated foreground is + correct at the natural threshold - with `sigma` 0.5 and density 20. *Height*: identical to `baseline`, + 0.75 / 0.25. The most uniform of the six across datasets (7 of 9 up, worst -4.5 %). +- **`both`** (touching + calibrated foreground). *Seeds*: as `fgcal`. *Height*: 0.5 / 0.5, like `contact`. +- **`boundary_fgcal`** (full boundary + calibrated foreground). *Seeds*: as `fgcal`. *Height*: 0.5 / 0.5. The + only optimum that also halves `min_size` to 25. + +Two readings. The **height map** splits by the foreground loss and the channel type, not by score: every decoder +whose loss touches the foreground twice (`contact`, `both`, `boundary_fgcal`) falls back to the 0.5/0.5 mix, +while the three that predict a single clean foreground (`baseline`, `fgcal`, `boundary`) trust it at 0.75 - and +those three are the three best. The **seed map** splits by the foreground threshold exactly as 7.2 describes, +and otherwise only distinguishes `baseline`/`contact` (broad smoothing, density 10) from the four decoders that +prefer sharp smoothing with a doubled threshold - a difference the plateau check shows to be within noise. + +### 8.4 Exactly how the boundary channel enters the height map + +`contact_weight` is not a grid dimension, so none of the optima above uses it; screened separately (5.6 point 8) +it is worth +0.4 to +0.5 % for `boundary` and at most +0.13 % for `boundary_fgcal`. The construction is: + +```python +h = watershed_heightmap(p, d, foreground_weight) # convex mix of two [0,1] terms -> h in [0, 1] +if contact is not None and contact_weight: # None or 0 skips this entirely + h = h + contact_weight * np.clip(contact, 0, 1) # additive, and NOT renormalised +h = lower_height_under_seeds(h, seeds, seed_floor) # "none" in every configuration here +seg = watershed(h, markers=seeds, mask=fg_mask) +``` + +Three consequences that the bracketed `[+ contact_weight * contact]` of 8.2 hides: + +1. **The base map is bounded and the ridge is not.** `watershed_heightmap` normalises both of its terms to + [0, 1] and combines them with weights `w` and `1 - w`, so `h` lies in [0, 1]; the ridge is then added on top, + giving [0, 1 + `contact_weight`]. A weight of 1.0 is therefore as tall as the **entire dynamic range** of the + base map, not a small correction. +2. **It reads the raw sigmoid probability, never a threshold** (channel 4 is sigmoid-activated in + `micro_sam/v2/models/util.py`, so the `clip` is only a safety net). Every pixel contributes in proportion to + its confidence, which is why the ridge extracts signal from an under-confident head (5.6 point 1). +3. **Measured against the map it modifies** (livecell, `boundary`, `foreground_weight` 0.75, medians of 20 + crops): the base map sits at 0.236 on the touching lines against 0.091 in the interiors, a natural contrast + of **+0.145**, while the contact probability is 0.623 against 0.055. So the ridge multiplies the barrier the + watershed must climb between two touching objects by **3.0x at weight 0.5, 4.9x at weight 1, 8.8x at 2 and + 16.7x at 4**. + +That quantifies both measured behaviours at once: why the ridge removes merges so effectively (7.8 % -> 4.3 % of +the objects) and why `seeded_split` rises monotonically with the weight - beyond about weight 2 the ridge no +longer assists the boundary evidence in `h`, it overrides it, so any spurious boundary probability inside a thin +object cuts it in two. The screened optimum being weight 0.5-1 on the holdout is consistent with that. + +`contact_mask_threshold` never touches the height map: it thresholds the contact map to shrink the watershed +*mask*, floods the interiors first and then re-floods so the instances claim the excluded band - which is why it +does not shave thin objects the way the ridge does (5.6 point 8). diff --git a/finetuning/v2/evaluation/optimization/notes/AIS_DECODER_TRAINING_PROPOSAL.md b/finetuning/v2/evaluation/optimization/notes/AIS_DECODER_TRAINING_PROPOSAL.md new file mode 100644 index 000000000..d8b6952c6 --- /dev/null +++ b/finetuning/v2/evaluation/optimization/notes/AIS_DECODER_TRAINING_PROPOSAL.md @@ -0,0 +1,114 @@ +# Improving the UniSAM2 decoder for automatic instance segmentation (AIS) + +What the 2026-09 AIS post-processing campaign on the joint/v4 geodesic `hvit_t` model found about the +decoder's output, and how the decoder's training should change to remove the two losses that +post-processing cannot reach. Evidence and numbers: `finetuning/v2/evaluation/optimization/notes/AIS_V4_OPTIMIZATION.md`. + +## What the decoder predicts today + +The automatic branch regresses four channels: a foreground probability (Dice loss against the binary mask) +and three directed-distance channels (MSE, masked to the foreground) whose target is the *geodesic hybrid* +field (`micro_sam/v2/transforms/labels.py`, `GeodesicHybridDistanceTransform`): direction = gradient of the +geodesic distance from the object's centre, magnitude = per-object normalised distance to the object's own +boundary. Post-processing follows the negated field to sinks (seeds), then floods a height map +`0.5 (1 - fg) + 0.5 (1 - |d|)` inside `fg > 0.5`. + +Measured properties of the prediction (cached predictions of the 2026-09 development corpus): + +- The field is smooth where the target is discontinuous. At a contact line between two touching cells the + target's direction flips; the prediction turns over a 6-8 px band (cosine between the flow 1 px on either + side of a contact pixel: +0.90, inside an object +0.97; at ±3 px: ≈0 vs +0.70). The magnitude dips to 0.17 + (median) at contacts against 0.34 inside, instead of to zero, and with gaps along the line. +- At the sink (object centre) the magnitude is 0.04-0.2, not the target's 1: the network smears the + single-pixel zero of the source over the whole centre region. +- In the background the network emits the label transform's fill value (|d| ≈ 1.0-1.1) even though the + distance loss is masked there; the halo right around an object, however, carries a smooth continuation + of the object's field (only 0-13 % of halo pixels exceed 0.8). +- The thresholded foreground has a dataset-dependent bias: area ratio to the ground truth 1.13 (livecell), + 1.17 (tissuenet, yet 39 % of its object pixels fall below 0.5), 1.48 (dynamicnuclearnet), 3.5 (deepbacs, + thin rods). The matched-object IoU is 0.67-0.84, which caps mSA at the higher IoU thresholds. + +## Point 1: contact information (merges) + +Merges are the dominant loss on touching-cell data: livecell 25 % of the objects are seeded but end in +an instance that also covers a neighbour, and another 11 % are unseeded and absorbed; tissuenet 20 % + 5 %; +neurips_cellseg 17 % + 10 %. Oracles with the predicted seeds and foreground: a ground-truth ridge doubles +livecell (0.27 → 0.54) and deepbacs, +67 % on tissuenet. Every label-free rule tried on the predicted field +failed to separate "two seeds in one cell" from "two touching cells" (edge/interior magnitude ratio 0.58 vs +0.35 with heavy overlap). The information is not in the prediction. + +Proposed changes, in the order I would try them: + +1. **An explicit contact channel.** Add a fifth output channel trained on the *touching boundary* mask: + pixels of an object adjacent to another object (`find_boundaries(labels, mode="inner")` restricted to + pixels whose neighbourhood contains a second non-zero label, dilated by one pixel so the target is 2-3 px + wide and learnable). Loss: Dice or focal BCE (the class is rare). Post-processing then adds the channel + as a ridge term to the height map, or excludes it from the watershed mask and reassigns it afterwards + (the EM training already does the analogous thing implicitly: `expected_fg = fg & ~boundary`). This is + the cheapest change with the largest expected return, since the ridge oracle shows the ceiling. +2. **Boundary-excluded foreground for LM, as in the EM recipe.** Train the LM foreground target as + `fg & ~find_boundaries(labels, mode="inner")` (a one-pixel gap between touching objects, and a + one-pixel shrink at every boundary). The gap gives the watershed mask a separation it currently lacks + and the shrink counters the over-prediction of point 4 (see below). Risk: the shrink changes the + foreground calibration for every dataset by one pixel, which is a lot for 50-pixel objects; it has to be + paired with a dilation-by-one of every instance after the watershed, which the EM pipeline does not do + either. Cheaper than 1 (no new channel), less targeted. +3. **Sharpen the field at contacts through the loss.** The masked MSE weights every foreground pixel + equally; contact pixels are <2 % of them and the network averages the two objects' fields there. Weight + the distance loss by proximity to a contact (e.g. 5× within 3 px of a touching boundary), or add a + direction term (`1 - cos` between predicted and target unit vectors, weighted the same way) so that the + flip is penalised as a direction error and not only through the small magnitude residual. Expect a + sharper flip, not a sharp one: an L2 regressor will still average within its receptive field. +4. **Instance-affinity output for the merge decision.** A short-offset affinity channel (is the pixel 2 px + away the same instance?) decided per pixel is what the merge rule needed and could not compute from the + field. This is the most invasive option (a new head, a new loss, and the post-processing becomes a + mutex/affinity watershed, which `bioimage_cpp.segmentation.mutex_watershed` provides) and would replace + the seeded watershed rather than fix it. + +What would show that it worked: the merged + absorbed fraction on livecell / tissuenet in the D2 +decomposition of the benchmark (`benchmark_ais_optimization.py run`, columns `seeded_merged`, +`unseeded_absorbed`) falls from 36 % / 25 % towards the ridge oracle's level, and the contact-line +cosine at ±1 px turns negative. + +## Point 4: instance extent (foreground calibration) + +The foreground threshold is a compromise whose sign differs by dataset (tissuenet under-covers, deepbacs +over-covers by 3.5×), so no global rule generalizes, and the halo carries a smooth field, so the magnitude +cannot trim it. The extent is also where APG's advantage over AIS comes from: the same seeds with SAM2 +masks score 20 % higher in 2D and 80 % in 3D. + +Proposed changes: + +1. **Calibrate the foreground target to the boundary, not to the mask.** The Dice loss on the binary mask + rewards a soft, wide foreground (a boundary pixel predicted at 0.6 costs almost nothing). Options: + a per-pixel loss with boundary weighting (BCE with weights rising towards the boundary, or a boundary + Dice term on the ring of ±2 px), or a signed-distance regression for the object extent (predict the + signed distance to the object boundary, positive inside; the extent is the zero level set, which is + sub-pixel and calibrated by construction). The signed distance is the natural companion of the geodesic + channels and can replace the foreground channel entirely. +2. **Resolve the label-convention conflict explicitly.** The over-prediction on deepbacs (thin rods + annotated tighter than the visible cell) and the under-prediction on tissuenet are label conventions the + network averages over. Two remedies: (a) a per-dataset boundary offset during training (dilate or erode + the masks of the datasets whose annotation is systematically tight or loose, measured once against the + raw intensity edge), so that the network learns one convention; (b) a small conditioning input (the + dataset's convention as an offset in pixels) which is not available at inference for new data and so is + the weaker option. (a) is a data-preparation change and costs nothing at inference. +3. **Train the extent at the object's own scale.** The 1024-px resize of the encoder puts a 20-px nucleus + and a 200-px cell through the same boundary blur. Multi-scale sampling of the training patches (the + generalist loader already has patch shapes; add a scale augmentation targeting 30-80 px objects) or + an auxiliary loss on the boundary IoU at the native resolution would sharpen the small objects, where a + one-pixel error is 10-20 % of the IoU. + +What would show that it worked: the area ratio of `fg > 0.5` to the ground truth moves towards 1 on every +dataset at the same threshold, the matched-object IoU (`matched_iou` column) rises from 0.67-0.84, and +the per-dataset optimum of `foreground_threshold` in the sweep collapses to one value. + +## Order and cost + +Point 1.1 (contact channel) and 4.1 (boundary-calibrated foreground or signed distance) are one training +run each on the existing joint recipe (`finetuning/v2/generalist/train_joint.py`, `distance_type`), with a +new label transform in `micro_sam/v2/transforms/labels.py` and a loss term in +`micro_sam/v2/loss/directed_distance_based.py`. The benchmark and its caches evaluate a new checkpoint end +to end in under an hour (predict once, then the diagnostics), and the oracles give the ceiling for each +change before any post-processing is retuned. Everything the post-processing side can still do without a +better field is listed at the end of `AIS_V4_OPTIMIZATION.md`. diff --git a/finetuning/v2/evaluation/optimization/notes/AIS_DICE_REOPTIMIZATION.md b/finetuning/v2/evaluation/optimization/notes/AIS_DICE_REOPTIMIZATION.md new file mode 100644 index 000000000..16262a9e0 --- /dev/null +++ b/finetuning/v2/evaluation/optimization/notes/AIS_DICE_REOPTIMIZATION.md @@ -0,0 +1,190 @@ +# Dice-foreground AIS decoder re-optimization + +This is the bounded follow-up comparison of the already trained `baseline.pt` and `boundary.pt` checkpoints. +It does not train another decoder and does not include a foreground-BCE variant. Both checkpoints were trained +with the Dice foreground objective; the boundary model additionally predicts the full object-boundary channel. +Here “Dice” distinguishes these checkpoints from the foreground-calibration (`fgcal`) experiments: the already +trained auxiliary boundary head retains the Dice-plus-BCE loss recorded by its checkpoint's training code. + +## Fixed protocol + +- Checkpoints: `/baseline.pt` and `/boundary.pt`, where `` is + `/mnt/vast-nhr/projects/cidas/cca/experiments/micro_sam2/apg_optimization/ais_decoder_training/staged/joint_sam2_hvit_t_multi_gpu`. +- Development datasets (equal weight): LIVECell, TissueNet, DynamicNuclearNet, DeepBacs, YeaZ, + NeurIPS CellSeg, PUMA, TNBC and COVID-IF. These are `primary training_extra` with DIC-HepG2 and DeepSeas + excluded. No 3D data enters selection. +- Diagnostic holdout: LIVECell, TissueNet and DynamicNuclearNet only. Reused DeepBacs is excluded. This is a + robustness check, not another selection stage. +- Sealed OOD confirmation: Arvidsson 10/10, BitDepth NucSeg 70/70 (reported equally over four magnifications), + CellBinDB 48 (8 per six stain/acquisition types), microbeSEG 2/2 manual test images and VICAR 50 (10 per five + cell types). Inputs are centre-cropped to at most 512 x 512; smaller images remain at native size. +- The baseline and boundary checkpoint each get their own configuration. A shared post-processing setting is not + used for the checkpoint comparison. + +`benchmark_apg_optimization.py --prepare-only --subset ood_extended --ndim 2` creates and freezes the OOD +manifest as `subset_manifest_v5_ood_extended.json`. Selection uses the test loaders, validates both raw and label +files, and refuses changed counts or strata. `report_ais_checkpoint_comparison.py` independently checks its paths +against both decoder `data_manifest.json` files before it reports a result. + +## Search + +The corrected cached scorer in `parameter_search.py` now uses channel 4 exactly like production AIS: +`contact_weight` raises the watershed height at boundaries; `contact_mask_threshold` runs the open-mask watershed +and re-flood. These `contact_*` spellings are legacy post-processing API names; for `boundary.pt`, channel 4 is +the full object-boundary probability, not the earlier touching-contact target. The scorer rejects those +parameters for a four-channel checkpoint. JSON `null` consistently means “use the model default”; only the +explicit string `"off"` disables the default boundary-magnitude filter. + +Coarse candidate families: + +- baseline: `configs/ais_dice_reopt_base.json` (1,344 combinations); +- boundary: `configs/ais_dice_reopt_boundary.json` contains the no-auxiliary (1,344), ridge (4,032), mask + (4,032), and ridge-plus-mask (1,344) families in one 10,752-candidate sweep. The four component JSON files + retain the individual family specifications for inspection. + +The grids cover foreground threshold 0.30-0.60, density threshold 5-50, size 25/50, sigma 0.5/1.0, +400-1,600 flow iterations and foreground height weight 0.5/0.75/1.0. The boundary families jointly search these +with ridge weight or mask threshold. The second stage is deliberately local: take the top three rows of each +mechanism family and vary one coordinate. `prepare_ais_reoptimization_polish.py` creates this explicit candidate +grid and only extends flow to 2,400 iterations when the 1,600-iteration edge still gains at least 0.001 mSA (and +similarly tests 200 only when the 400 edge beats 800). + +Sweep sharding is cache-aware: every configuration with the same foreground threshold, smoothing, iteration +count and step size stays in one shard. Thus the expensive flow density is computed once per image and flow group, +not once per shard. The consolidated boundary grid also shares that computation across all four mechanism +families. Shards still form an exact disjoint partition of the requested combinations. The `cpu-test` submission +preset packs up to 48 four-thread shard commands into one exclusive 192-core test-node allocation. The coarse +layout uses 12 shards for each of four primary datasets (48 commands) and 9 for each of five training-extra +datasets (45 commands), so every submitted array occupies one node and uses most of its cores. + +Rank with `report_ais_sweep.py --no-reference`. It unions multiple mechanism grids and selects from the rows no +more than 0.001 mSA below the best. Within that plateau it favours the best worst-dataset relative optimum, then +fewer flow iterations and fewer boundary controls. The emitted JSON is directly accepted by `run`. + +## Execution recipe + +From `finetuning/v2/evaluation/optimization`, with the `new-stack` environment active: + +```bash +export MICRO_SAM2_JOINT_CHECKPOINT_ROOT=/mnt/vast-nhr/projects/cidas/cca/experiments/micro_sam2/apg_optimization/ais_decoder_training/staged +export MICRO_SAM2_JOINT_EXPORT_ROOT=/mnt/vast-nhr/projects/cidas/cca/experiments/micro_sam2/apg_optimization/model_exports +ROOT=/mnt/vast-nhr/projects/cidas/cca/experiments/micro_sam2/apg_optimization +REP=$ROOT/ais/reports/dice_reoptimization +PRIMARY_CORE="livecell tissuenet dynamicnuclearnet deepbacs" +EXTRA_CORE="yeaz neurips_cellseg puma tnbc covid_if" +CORE="$PRIMARY_CORE $EXTRA_CORE" +``` + +Prepare the manifest and cache development predictions. The four `--print-only` invocations print the task +graphs to inspect; omit that flag only after checking them. + +```bash +python benchmark_apg_optimization.py --prepare-only --subset ood_extended --ndim 2 +python ais_campaign_tasks.py predict --name dice_base_primary_predict --subsets primary --print-only \ + --extra "--joint-checkpoint baseline --ndim 2 --datasets $PRIMARY_CORE" +python ais_campaign_tasks.py predict --name dice_base_extra_predict --subsets training_extra --print-only \ + --extra "--joint-checkpoint baseline --ndim 2 --datasets $EXTRA_CORE" +python ais_campaign_tasks.py predict --name dice_boundary_primary_predict --subsets primary --print-only \ + --extra "--joint-checkpoint boundary --ndim 2 --datasets $PRIMARY_CORE" +python ais_campaign_tasks.py predict --name dice_boundary_extra_predict --subsets training_extra --print-only \ + --extra "--joint-checkpoint boundary --ndim 2 --datasets $EXTRA_CORE" +``` + +Run the coarse sweeps. Primary and training-extra are separate because their dataset sets do not overlap. Use +distinct `--name` values for each checkpoint and subset. The 12/9 shard counts pack each array into one full +test node and do not change the candidate set. + +```bash +python ais_campaign_tasks.py sweep --name dice_base_primary --preset cpu-test --subsets primary \ + --grid configs/ais_dice_reopt_base.json --datasets $PRIMARY_CORE --num-shards 12 --print-only \ + --extra "--joint-checkpoint baseline --ndim 2 --mode sparse" +python ais_campaign_tasks.py sweep --name dice_base_extra --preset cpu-test --subsets training_extra \ + --grid configs/ais_dice_reopt_base.json --datasets $EXTRA_CORE --num-shards 9 --print-only \ + --extra "--joint-checkpoint baseline --ndim 2 --mode sparse" +python ais_campaign_tasks.py sweep --name dice_boundary_primary --preset cpu-test --subsets primary \ + --grid configs/ais_dice_reopt_boundary.json --datasets $PRIMARY_CORE --num-shards 12 --print-only \ + --extra "--joint-checkpoint boundary --ndim 2 --mode sparse" +python ais_campaign_tasks.py sweep --name dice_boundary_extra --preset cpu-test --subsets training_extra \ + --grid configs/ais_dice_reopt_boundary.json --datasets $EXTRA_CORE --num-shards 9 --print-only \ + --extra "--joint-checkpoint boundary --ndim 2 --mode sparse" +``` + +After every shard succeeds, merge each grid with one `benchmark_ais_optimization.py sweep --merge` call using +`--num-shards 12` for primary and `--num-shards 9` for training_extra, then rank: + +```bash +mkdir -p "$REP" +python report_ais_sweep.py --grid configs/ais_dice_reopt_base.json --subset primary training_extra \ + --datasets $CORE --joint-checkpoint baseline --no-reference --output "$REP/baseline_coarse.csv" +python report_ais_sweep.py --grid configs/ais_dice_reopt_boundary.json \ + --subset primary training_extra --datasets $CORE \ + --joint-checkpoint boundary --no-reference --output "$REP/boundary_coarse.csv" +python prepare_ais_reoptimization_polish.py --ranking "$REP/baseline_coarse.csv" \ + --output configs/ais_dice_reopt_baseline_polish.json +python prepare_ais_reoptimization_polish.py --ranking "$REP/boundary_coarse.csv" \ + --output configs/ais_dice_reopt_boundary_polish.json +``` + +Each generator command prints the number of distinct flow-cache groups. Set `--num-shards` no higher than that +printed count. A polish grid is much smaller than the coarse grid: submit its individual shard commands with +`cpu-shared`, or combine the primary and training-extra task lists before using the packed `cpu-test` preset. +The sweep rejects a larger shard count instead of producing empty result files. Then merge with the same chosen +count, rank the union and write the two own-optimum configs: + +```bash +python report_ais_sweep.py --grid configs/ais_dice_reopt_base.json \ + configs/ais_dice_reopt_baseline_polish.json --subset primary training_extra --datasets $CORE \ + --joint-checkpoint baseline --no-reference --output "$REP/baseline_final.csv" \ + --select-config "$REP/baseline_optimum.json" --config-name baseline-dice-optimum +python report_ais_sweep.py --grid configs/ais_dice_reopt_boundary.json \ + configs/ais_dice_reopt_boundary_polish.json \ + --subset primary training_extra --datasets $CORE --joint-checkpoint boundary --no-reference \ + --output "$REP/boundary_final.csv" --select-config "$REP/boundary_optimum.json" \ + --config-name boundary-dice-optimum +``` + +Run the own-optimum configs on the three-dataset disjoint holdout for diagnosis. Only after configs are frozen, +cache and score `ood_extended` for both checkpoints. Keep diagnostics enabled. Generate these task graphs with: + +```bash +HOLDOUT="livecell tissuenet dynamicnuclearnet" +python ais_campaign_tasks.py screen --name dice_base_holdout --preset cpu-shared --subsets holdout --no-defaults \ + --configs "$REP/baseline_optimum.json" --print-only \ + --extra "--joint-checkpoint baseline --ndim 2 --datasets $HOLDOUT" +python ais_campaign_tasks.py screen --name dice_boundary_holdout --preset cpu-shared --subsets holdout --no-defaults \ + --configs "$REP/boundary_optimum.json" --print-only \ + --extra "--joint-checkpoint boundary --ndim 2 --datasets $HOLDOUT" +python ais_campaign_tasks.py predict --name dice_base_ood --subsets ood_extended --print-only \ + --extra "--joint-checkpoint baseline --ndim 2" +python ais_campaign_tasks.py predict --name dice_boundary_ood --subsets ood_extended --print-only \ + --extra "--joint-checkpoint boundary --ndim 2" +python ais_campaign_tasks.py screen --name dice_base_ood_score --preset cpu-shared --subsets ood_extended \ + --no-defaults --configs "$REP/baseline_optimum.json" --print-only \ + --extra "--joint-checkpoint baseline --ndim 2" +python ais_campaign_tasks.py screen --name dice_boundary_ood_score --preset cpu-shared --subsets ood_extended \ + --no-defaults --configs "$REP/boundary_optimum.json" --print-only \ + --extra "--joint-checkpoint boundary --ndim 2" +``` + +Submit the score tasks only after their prediction tasks finish. Finally pass the two OOD run directories and +the frozen manifest to: + +```bash +python report_ais_checkpoint_comparison.py --baseline-runs \ + --boundary-runs --manifest "$ROOT/subset_manifest_v5_ood_extended.json" \ + --output "$REP/ood_confirmation.json" +``` + +## Decision rule + +The JSON report supports a strong boundary-decoder improvement statement only when all four conditions hold: + +1. the 95% paired hierarchical-bootstrap CI for the macro mSA difference is above zero; +2. at least four of the five OOD domains improve; +3. no domain loses both more than 0.005 absolute mSA and more than 2% relative; +4. equal-domain macro mSA improves by at least 2% relative. + +The bootstrap resamples domains, then paired source images within every domain/acquisition stratum. The report +also writes per-domain and paired-sample CSVs, object-fate diagnostic deltas, generation time, both selected +parameter dictionaries, checkpoint/implementation checksums and the training-disjointness audit. microbeSEG is +always labelled as an `n=2` stress test, not treated as precise standalone evidence. diff --git a/finetuning/v2/evaluation/optimization/notes/AIS_V4_OPTIMIZATION.md b/finetuning/v2/evaluation/optimization/notes/AIS_V4_OPTIMIZATION.md new file mode 100644 index 000000000..bb520cbd0 --- /dev/null +++ b/finetuning/v2/evaluation/optimization/notes/AIS_V4_OPTIMIZATION.md @@ -0,0 +1,675 @@ +# AIS optimization for the joint/v4 geodesic `hvit_t` model + +Decision log of the AIS (decoder-based automatic instance segmentation) optimization campaign started +2026-09-06 on branch `ais-v4-optim` (forked from `apg-clean-up` at `4a3ef31`). Set-up, data, gates and +cluster mechanics: `EXPERIMENTAL_SETUP.md`; the plan: `~/.claude/plans/please-plan-a-campagin-cozy-willow.md`. +Paths are relative to `finetuning/v2/evaluation/`; `` is +`/mnt/vast-nhr/projects/cidas/cca/experiments/micro_sam2/apg_optimization`. + +## Why + +The v4 decoder predicts the geodesic hybrid field (`micro_sam/v2/transforms/labels.py`, +`GeodesicHybridDistanceTransform`): the direction of every pixel's vector is the gradient of the geodesic +distance from the object's centre, so `-d` converges to one sink per object; the magnitude is the +per-object normalised distance to the object's own boundary. The v2 decoder predicted the Euclidean +vector to the nearest boundary, whose negation converges onto the medial axis. The AIS post-processing +(`micro_sam/v2/postprocessing.py`, `flow_instance_segmentation`) and its `hvit_t` defaults (fg 0.5, +density 10, min_size 100, sigma 0.5, n_iter 50, dt 0.5, fg_weight 0.5) were derived for the v2 field: +a fixed 25 px travel, an absolute density threshold, a height map from the inverted magnitude. + +Scope decided with the user on 2026-09-06: sparse (flow) pipeline first, dense (multicut) afterwards; +the deliverable is new library logic and new `hvit_t` defaults in `postprocessing.py`; numpy prototypes +for primitives bioimage-cpp lacks, C++ port before the library switch if such a variant wins. `hvit_t` +only, no learned components, no per-dataset modes. + +## Harness (Phase 0, 2026-09-06) + +`optimization/benchmark_ais_optimization.py` predicts every manifest sample once and caches the +`(4, *spatial)` float32 prediction with its labels under +`/ais/predictions///.npz` (`predict`); every +configuration (`run`, `screen`), the parameter grid (`sweep`) and the diagnostics then run on the cache +on CPU. Run directories follow the APG layout, `/ais/hvit_t//--/` with `samples.csv`, `summary.csv`, `metadata.json`, so +`compare_apg_optimization.py` reads them. The implementation checksum covers the benchmark, `common.py`, +`parameter_search.py`, `micro_sam/v2/instance_segmentation.py` and `micro_sam/v2/postprocessing.py`. + +Per-sample columns beyond the metrics: `matched` / `unmatched` / `severed_objects` / `genuine_misses` +(the `benchmark_apg_3d.object_counts` definitions, computed from one contingency table), and the seed +diagnostics of a pipeline mirrored step by step (`sparse_pipeline`): `n_seeds`, `gt_with_0_seeds`, +`gt_with_1_seed`, `gt_with_2plus_seeds`, `background_seeds` (majority pixel in the background), +`seeded_unmatched` (seeded, lost in the watershed), `matched_before_min_size`, `fg_iou` and +`pipeline_mismatch` (mirrored segmentation differs from the library's; the bit-identity check of an epoch). +`report` joins the subsets of a screen and applies the generalization gate (up on all but two datasets, +no dataset below both −2 % and −0.005, balanced gain ≥ +2 %). Configuration files: +`configs/ais_*.json` (`{"name", "mode", "params_2d", "params_3d"}`; a flat dict is sparse overrides, +`{"sparse": ..., "dense": ...}` sets both). Task builder: `optimization/ais_campaign_tasks.py` +(`predict`, `screen`, `sweep`). Unit tests: `test/test_ais_optimization.py` (15 tests). + +Smoke test (deepbacs, 30 primary images, library defaults, session A100): balanced mSA 0.1604; of 892 +ground-truth objects 607 matched, 33 without a seed, **268 with two or more seeds**, 248 background seeds; +mirrored pipeline identical on all 30 images. The v4 field over-seeds the rods: a first sign that the +default travel (25 px) and the absolute density threshold do not fit a centre-directed field. + +## Log + +- 2026-09-06 19:30: harness written, unit tests green, smoke test passed. Prediction caching of the + v5 subsets (primary, training_extra, holdout) and the deep 3d crops (apg3d primary, holdout) started + on the session GPU. +- 2026-09-06 20:15: harness frozen (AIS epoch `f57b117edfda5420d9df761b1db4db2d`, commit of this state). The oracle markers were + changed from one pixel to the 3-neighbourhood inside the object after the first oracle run: the geodesic + magnitude is zero at an object's centre pixel (gradient of a field at its source), so the inverted + magnitude height map has a one-pixel spike there and the monotone flooding of `bioimage_cpp`'s watershed + floods a seed on a spike last (one pixel left to the object). Predicted seeds are multi-pixel blobs, so + the pipeline itself is unaffected, but any seed logic that places small seeds at the magnitude peak must + keep this in mind. Prediction caches: v5 primary 245, training_extra 157, holdout 238 samples (float32, + labels included); apg3d primary / holdout in progress. + +## Phase 1 (2026-09-06 evening): baseline, travel ladder (D1) and oracles (D3) + +All on the cached joint/v4 geodesic predictions (checksum `5a729846…`), library defaults, AIS epoch +`f57b117edfda5420d9df761b1db4db2d` (the epoch of the frozen Phase 0 harness; run directories under +`/ais/hvit_t/5a729846…/`). Per-dataset mSA of the defaults: + +| subset | balanced | livecell | tissuenet | dynamicnuclearnet | deepbacs | dic_hepg2 | volumes (n = 1 each) | +|---|---:|---:|---:|---:|---:|---:|---| +| primary (245) | 0.1841 | 0.2683 | 0.2102 | 0.5422 | 0.1604 | 0.0019 | celegans 0.131, embedseg 0.165, gonuclear 0.340, cremi CREMI 1.057, snemi CREMI 1.054 | +| holdout (238) | 0.1826 | 0.2726 | 0.2112 | 0.5223 | 0.1604 | 0.0021 | (same volumes) | +| training_extra (157) | 0.4183 | yeaz 0.6128, neurips_cellseg 0.2168, deepseas 0.1016, puma 0.4668, covid_if 0.7411, tnbc 0.3705 | | | | | | + +For comparison, APG defaults on the same manifests: primary 0.2955, holdout 0.2896, training_extra 0.4634 +(EXPERIMENTAL_SETUP.md §14.1). The dense multicut on the 12-slice cremi / snemi crops over-segments +massively (1831 and 3005 instances for 134 and 96 objects); Phase 5 material. + +Object fates of the defaults (primary + training_extra, 2d): livecell 17389 objects, 8323 matched, 2735 +without a seed, 2014 with two or more seeds, **6338 seeded but unmatched**; tissuenet 4011 / 2115 / 512 / +490 / 1384; deepbacs 892 / 607 / 33 / 268 / 252 (plus 248 background seeds); neurips_cellseg 5766 / 2141 / +1424 / 421 / 2204; deepseas 250 objects but 504 background seeds; yeaz 450 of 2689 objects split; +dic_hepg2 foreground IoU 0.09 (the foreground channel fails on DIC, nothing to post-process). The +"seeded but unmatched" category dominates everywhere; the refined decomposition (split / merged / +undersized / oversized, absorbed / missing) was added to the harness afterwards. + +### D1: travel ladder (`configs/ais_s0_travel_*.json`, cluster job 15766947, report `ais/reports/s0_travel_ladder_dev.csv`) + +Relative change of mSA against the defaults (travel 25 px) on the development manifests: + +| travel (px) | balanced (16 datasets) | livecell | tissuenet | dynamicnuclearnet | deepbacs | yeaz | neurips_cellseg | deepseas | puma | tnbc | 0-seed Δ | 2+-seed Δ | bg-seed Δ | +|---|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:| +| 12.5 | +8.5 % | −6.5 | −5.7 | +0.3 | +1.8 | −3.0 | −0.7 | −3.9 | +0.2 | −3.5 | +2705 | −2437 | −1799 | +| 50 | −2.5 % | +0.4 | −0.7 | 0.0 | −3.9 | −1.5 | +1.4 | −7.8 | −0.3 | −1.1 | −455 | +808 | +987 | +| 100 | −1.7 % | +0.2 | −0.9 | −0.2 | +5.5 | −0.9 | +1.8 | −6.5 | −0.5 | −2.2 | −131 | +438 | +1419 | +| 200 | +5.6 % | +0.3 | −0.9 | −0.2 | +10.6 | −0.6 | +3.1 | −4.6 | −0.2 | −1.7 | +202 | −28 | +1398 | +| 400 | +7.2 % | +0.3 | −0.9 | −0.2 | +14.6 | −0.7 | +3.1 | −1.6 | −0.2 | −1.7 | +311 | −192 | +1324 | + +The balanced figures are inflated by the single 12-slice embedseg volume (+81-130 %) and the near-zero +dic_hepg2; on the images the travel moves deepbacs (+14.6 % at 400 px) and neurips_cellseg (+3 %) and +costs everything else a little. Longer travel trades splits for background seeds (+1324 at 400 px) and +does not seed more objects. **Verdict: the travel is a secondary knob; run to convergence only together +with a seed rule that suppresses the background sinks.** No candidate passes the gate. + +### D3: oracles (`oracle`, primary / training_extra / holdout, defaults; `/ais/oracles/`) + +mSA when one part of the pipeline is replaced by the ground truth (predicted parts otherwise), primary ++ training_extra images: + +| dataset | baseline | GT seeds | GT height map | GT seeds + GT height map | GT foreground | GT seeds + GT foreground | +|---|---:|---:|---:|---:|---:|---:| +| livecell | 0.268 | 0.319 (+19 %) | 0.539 (+101 %) | 0.612 | 0.422 (+57 %) | 0.504 | +| tissuenet | 0.210 | 0.226 (+7 %) | 0.351 (+67 %) | 0.357 | 0.345 (+64 %) | 0.394 | +| dynamicnuclearnet | 0.542 | 0.568 (+5 %) | 0.556 (+3 %) | 0.575 | 0.972 (+79 %) | 0.988 | +| deepbacs | 0.160 | 0.266 (+66 %) | 0.407 (+154 %) | 0.488 | 0.621 (+287 %) | 0.921 | +| yeaz | 0.613 | 0.628 (+2 %) | 0.687 (+12 %) | 0.702 | 0.885 (+44 %) | 0.910 | +| neurips_cellseg | 0.217 | 0.332 (+53 %) | 0.317 (+46 %) | 0.431 | 0.602 (+178 %) | 0.738 | +| deepseas | 0.102 | 0.184 (+81 %) | 0.168 (+65 %) | 0.261 | 0.806 | 0.923 | +| puma | 0.467 | 0.507 (+8 %) | 0.500 (+7 %) | 0.520 | 0.883 (+89 %) | 0.928 | +| tnbc | 0.371 | 0.412 (+11 %) | 0.395 (+7 %) | 0.415 | 0.849 | 0.964 | +| covid_if | 0.741 | 0.755 (+2 %) | 0.783 (+6 %) | 0.791 | 0.864 | 0.885 | + +Reading: (1) on the touching-cell data (livecell, tissuenet, deepbacs) a perfect ridge map with the +predicted seeds and foreground doubles the score, so the assignment step (height map / watershed) is the +largest lever that post-processing controls; (2) perfect seeds add +5-20 % on most datasets and +50-80 % +where background seeds and misses are frequent (deepbacs, deepseas, neurips_cellseg); (3) the ground-truth +foreground ceiling is the largest everywhere, but it leaks the instance separation wherever objects do not +touch (dynamicnuclearnet, puma, tnbc, yeaz: nuclei), so it mixes foreground extent with separation. The +part of it that is extent (mSA on small nuclei swings on one boundary pixel) is reachable only through +the foreground threshold and the instance extent rule, which the sweep and the height-map work cover. +Holdout reproduces the primary picture (livecell 0.273 → 0.557 with GT ridges, tissuenet 0.211 → 0.351). + +Priorities for Phase 2/3 from D1-D3: height-map ridge terms (H1 divergence, H2 direction discontinuity) +and trajectory assignment (A1) first, seed rules that suppress background sinks and merge multi-sink +objects second (S1, S2, S5), travel to convergence as a parameter of both. The seed-variant prototype +(`scratchpad/proto_seeds.py`, cluster job) screens all of these on 8 images per dataset before any +library edit. + +### Probe of the predicted field (four primary images per dataset, 2026-09-06 19:25) + +| dataset | \|d\| background p50 / p90 | \|d\| foreground p10 / p50 | IoU(fg > 0.5, GT fg) | area(fg > 0.5) / area(GT) | \|d\| at the object centre / 5×5 ring | +|---|---|---|---:|---:|---:| +| livecell | 1.04 / 1.13 | 0.24 / 0.48 | 0.87 | 1.13 | 1.00 | +| tissuenet | 1.02 / 1.13 | 0.18 / 0.48 | 0.58 | 1.17 | 1.01 | +| dynamicnuclearnet | 1.01 / 1.11 | 0.17 / 0.57 | 0.67 | 1.48 | 0.93 | +| deepbacs | 1.02 / 1.11 | 0.22 / 0.53 | 0.48 | 3.49 | 1.00 | + +Three consequences. (1) The decoder predicts the label transform's fill value (magnitude ≈ 1) in the +background although the distance loss is masked there, so the magnitude cannot serve as a foreground cue, +and along a ray from an object's centre the magnitude runs 1 → 0 (boundary) → 1 (background): the +boundary is the magnitude *minimum*, which the inverted-magnitude height map already turns into a ridge. +(2) The thresholded foreground is systematically too large, by half on the nuclei and 3.5× on the thin +deepbacs rods; the seeded watershed then floods every instance out to the foreground edge, which is why +the ground-truth-foreground oracle is so far above everything else. The instance extent is therefore a +first-order problem: either a higher foreground threshold (the sweep must go beyond 0.7) or, scale-free, +an extent defined by the flow (pixels whose trajectory reaches the instance's sink, no refill; the halo +pixels carry the background direction and do not converge). (3) The one-pixel magnitude dip at the +centre of the training target is not reproduced by the network (ratio ≈ 1.0), so seeds placed at the +magnitude maximum are safe in practice; the oracle-marker precaution stays. + +### D2: fate of every ground-truth object under the defaults (primary + training_extra images, epoch `5700c6e0…`) + +Percent of ground-truth objects. "matched" at IoU 0.5; "seed0" = no seed component inside; "absorbed" / +"missing" = unseeded objects mostly covered by a neighbour's instance / by nothing; "seeded lost" = +seeded but unmatched, decomposed into "split" (two or more seeds), "merged" (the object's instance covers +at least half of another object too), "under" / "over" (extent errors); "bg seeds" = seed components +whose majority pixel is background, as percent of the object count; "iou" = mean IoU of the matched +objects; "pred/gt" = predicted over ground-truth instance count. + +| dataset | gt | matched | iou | seed0 | absorbed | missing | seeded lost | split | merged | under | over | bg seeds | fg IoU | pred/gt | +|---|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:| +| livecell | 17389 | 47.9 | 0.75 | 15.7 | 11.2 | 4.4 | 36.4 | 6.7 | **24.6** | 1.1 | 4.0 | 8.3 | 0.82 | 0.71 | +| tissuenet | 4011 | 52.7 | 0.71 | 12.8 | 4.7 | 8.1 | 34.5 | 5.3 | **20.4** | 3.0 | 5.8 | 2.3 | 0.71 | 0.67 | +| neurips_cellseg | 5766 | 37.1 | 0.74 | 24.7 | 10.2 | 14.4 | 38.2 | 3.4 | **16.8** | 9.7 | 8.4 | 24.3 | 0.59 | 0.74 | +| deepbacs | 892 | 68.0 | 0.67 | 3.7 | 3.5 | 0.2 | 28.3 | **14.5** | 9.8 | 0.3 | 3.7 | **27.8** | 0.64 | 1.06 | +| deepseas | 250 | 63.6 | 0.67 | 11.6 | 7.6 | 3.6 | 25.2 | 6.0 | 10.4 | 4.0 | 4.8 | **201.6** | 0.45 | 1.72 | +| dynamicnuclearnet | 2592 | 94.0 | 0.80 | 1.1 | 0.3 | 0.8 | 4.9 | 0.3 | 0.9 | 1.5 | 2.2 | 10.4 | 0.75 | 1.03 | +| puma | 3293 | 88.1 | 0.79 | 4.1 | 1.0 | 3.1 | 7.8 | 1.0 | 1.7 | 0.6 | 4.5 | 10.8 | 0.77 | 0.98 | +| tnbc | 399 | 83.7 | 0.77 | 3.5 | 0.8 | 2.8 | 12.8 | 1.8 | 2.3 | 1.3 | 7.5 | 20.8 | 0.66 | 1.00 | +| yeaz | 2689 | 86.9 | 0.84 | 6.6 | 2.3 | 4.3 | 6.5 | 2.9 | 1.9 | 0.8 | 0.9 | 3.4 | 0.87 | 0.92 | +| covid_if | 481 | 91.3 | 0.90 | 5.0 | 2.1 | 2.9 | 3.7 | 0.2 | 1.2 | 0.0 | 2.3 | 0.8 | 0.92 | 0.93 | +| dic_hepg2 | 490 | 1.8 | 0.63 | 32.9 | 0.6 | 32.2 | 65.3 | 47.1 | 0.0 | 13.7 | 4.5 | 105.3 | 0.09 | 1.57 | + +The size filter alone (`min_size` 100, refilled by the neighbours) costs tissuenet 302 of 2417 matches +(7.5 %), neurips_cellseg 110, livecell 94, puma 86 (`matched_before_min_size` column); the ground-truth +size floors are 10-50 px, so a shared default has to be lower. + +Mechanisms to address, in order of the objects they cost: + +- **M1 merges** (livecell 36 % merged + absorbed, tissuenet 25 %, neurips_cellseg 27 %): mostly two seeded + objects whose basins are not separated, i.e. the ridge of the height map is too weak or misplaced + (the D3 height-map oracle with the same predicted seeds doubles livecell). With `foreground_weight` 0.5 + the ridge between touching cells is only the magnitude term, 0.5·(1 − |d|), a quarter above the + interior, and any gap in the magnitude dip along the contact line lets one basin flood the other. The + direction of the field flips across the contact line whatever the magnitude does, so a direction + discontinuity ridge (H2) or the trajectory assignment (A1) attacks this directly. +- **M2 unseeded objects** (neurips_cellseg 25 %, livecell 16 %, tissuenet 13 %): the absolute density + threshold (10) is unreachable for small objects, whose converged particles number about their area + divided by the sink footprint. Relative or particle-count seeds (S1, S2) with a lower `min_size`. +- **M3 background seeds** (deepseas 2× the object count, deepbacs 28 %, neurips_cellseg 24 %, tnbc 21 %): + false-positive foreground regions converge into sinks. The predicted field there is the background + fill (|d| ≈ 1 throughout, no dip at the region's edge, no converging structure), so an instance-level + field-consistency filter (magnitude along the instance boundary, or the divergence at the sink) is a + cheap, label-free way to drop them. +- **M4 extent** (neurips_cellseg 18 % under/over, tnbc 9 %, matched IoU 0.67-0.75 on the cell datasets): + the foreground over-predicts (see the probe), so the instance boundary sits outside the object. A + higher foreground threshold or a flow-defined extent (A1 without refill). + +### Direction structure at the contact lines (8 images each, livecell / tissuenet) + +Merged pairs mostly hold **distinct** seeds (livecell 449 pairs with distinct seeds vs 154 sharing a +seed component; tissuenet 52 vs 17), so the merges are an assignment failure, as the D3 oracle said. +But the predicted direction field does not flip sharply at a contact line: the cosine between the flow +1 px on either side of a contact pixel is +0.90 (median) against +0.97 inside; only at ±3-4 px does it +reach 0 / −0.35 (interior +0.70 / +0.49), and it also reverses around every object centre. The +divergence separates contacts from interiors only weakly (+0.03 vs −0.02). The magnitude dip is the +sharper cue: |d| 0.17 at contacts against 0.34 inside (p50), i.e. a ridge of a quarter of the height +range with gaps. The network smooths the target's discontinuities over a 6-8 px band. + +### Prototype 1: seed and assignment variants (`proto_seeds.py`, cluster job 15767065, 8 / 6 images per dataset) + +Balanced mSA over the prototype images (primary: deepbacs, dic_hepg2, dynamicnuclearnet, livecell, +tissuenet; extra: covid_if, deepseas, neurips_cellseg, puma, tnbc, yeaz). Baseline (travel 25, defaults) +0.232 / 0.383. + +| variant | primary | extra | note | +|---|---:|---:|---| +| foreground threshold 0.7, travel 400 | 0.245 | 0.382 | deepbacs +51 %, tissuenet −6 %, dynamicnuclearnet −2.5 %, tnbc −15 % | +| foreground threshold 0.6, travel 400 | 0.243 | 0.388 | the best on both, small per-dataset losses (tissuenet −3 %) | +| density threshold 50 at travel 400 | 0.240 | 0.387 | dynamicnuclearnet +3 %, livecell −12 % | +| travel 400, defaults otherwise | 0.235 | 0.381 | | +| divergence / direction ridges (H1, H2) | 0.235 / 0.234 | 0.380 / 0.381 | no effect, as the direction analysis predicts | +| relative density seeds (S1, 0.25-0.5 of the local maximum) | 0.205 | 0.353 | livecell −52 % (large cells split) | +| particle-count sinks (S2) | 0.207 | 0.352 | livecell −50 %, deepbacs +17 % | +| trajectory assignment, refilled (A1) | 0.219 | 0.364 | worse everywhere; loose mask (0.3) without refill 0.186 / 0.323: the halo converges too | +| magnitude cores (S3), divergence sinks (S4) | ≤ 0.20 | ≤ 0.36 | worse | + +Reading: the structural seed rules over-segment the large livecell cells (their predicted field has +several weak sinks and a jittering centre) while they help the small-object data, so a scale-free seed +rule alone does not generalize; the assignment by trajectories inherits the blurred field and is worse +than the watershed; per-pixel direction ridges are empty. The gains that do generalize on this small +sample are the foreground threshold (0.6-0.7: the over-predicted foreground, mechanism M4) and running the +flow to convergence, both parameters. Next: the height-map prototype (`proto_h.py`, job 15767091: +sharpened magnitude dips, relative magnitude, multi-offset reversal ridges, background-instance filters) +and a case study of the merged pairs. + +### Prototype 2: height maps and instance filters (`proto_h.py`, job 15767091; seeds = converged density, threshold 10) + +Reference `lib_fw0.5` (the library height map, travel 400): balanced 0.235 primary / 0.381 extra. + +| variant | primary | extra | per dataset | +|---|---:|---:|---| +| **boundary-magnitude filter 0.4** (drop instances whose boundary median \|d\| > 0.4) | **0.247** | **0.398** | deepbacs +9 %, dynamicnuclearnet +8 %, deepseas ×4, neurips_cellseg +72 %, tnbc +1 %, nothing down | +| boundary filter 0.6 | 0.244 | 0.387 | same direction, smaller | +| mean-magnitude filter 0.7 | 0.237 | 0.383 | weaker | +| sharpened dips exp(−\|d\|/τ), powers, relative magnitude | 0.233-0.237 | 0.373-0.381 | ±1 %, relative magnitude −5 % on yeaz | +| foreground weight 0 / 0.25 / 0.75 | 0.233 / 0.236 / 0.226 | 0.377 / 0.380 / 0.379 | the current 0.5 is fine | +| multi-offset reversal ridges (k = 2-4) | 0.18-0.19 | 0.26-0.32 | catastrophic: the field also reverses around every centre | + +Reading: the shape of the height map is not the lever; the background-instance filter is the first +label-free rule that improves every dataset it touches (mechanism M3). It is scale-free (a real object +has a magnitude dip along its whole boundary, a false foreground region carries the background fill). + +### Why the merges happen (seed-quality probe, 6 images per dataset) + +At the density peak (the sink) the predicted magnitude is small: |d| 0.04-0.22 for the seeds of matched +objects. The network smears the target's zero at the centre over the whole centre region, so every +proper seed sits on a **peak** of the inverted-magnitude height map. With the monotone flooding of the +watershed a seed's front never drops below the seed's own height, so a seed whose centre dip is deeper +than the contact-line dip to its neighbour (contact |d| 0.17 median) loses the object to the neighbour: +in the merged pairs of livecell the losing seed's own instance is 8.5 px (median) before the size +filter, of tissuenet 1 px. The same property silently suppresses spurious seeds, which is why zeroing +the height under all seeds (halving livecell's merges, +7 % matched) still lowered mSA on every dataset: +the spurious seeds then flood too (deepbacs 0.178 → 0.125). Seed-quality cues that separate proper +seeds from background seeds: the foreground probability at the peak (proper p10 0.8-0.99, background +p50 0.5-0.7), the converged particle count (proper p10 ≥ 35-460, background p50 15-43; but the extra +seeds of split objects sit in between), and the inward flux ratio of the flow on a ring of radius 6 +(proper p50 0.8-0.97, background 0.1-0.3 on deepbacs, dynamicnuclearnet, deepseas, livecell; not on the +small tissuenet objects). neurips_cellseg's "background" seeds are confident cells the labels do not +contain (fg 0.94, flux 0.98), out of reach for post-processing. + +Prototype 3 (`proto_merge.py`): seed floors (none / zero / ring minimum) × size floor (100 / 25) × +boundary filter (off / 0.4) × decoder-consistency merge of adjacent instances without a dip on their +shared boundary (off / 0.7 / 0.85), at foreground 0.5 and 0.6. + +### Prototype 3: seed floors and the decoder-consistency merge (`proto_merge.py`, job 15767152) + +72 variants (foreground 0.5 / 0.6 × floor none / zero / ring × size floor 100 / 25 × boundary filter +off / 0.4 × merge off / 0.7 / 0.85), the same images as before. Top of the tables: primary +`fg0.6 · mono · min_size 25 · filter 0.4` 0.2515 and `fg0.6 · mono · 100 · filter 0.4` 0.2498 (baseline +0.2350); extra `fg0.5 · mono · 100 · filter 0.4` 0.398 (baseline 0.381). Every floor variant lands below +the monotone flooding with the same filter, on both subsets. Livecell (fg 0.5, size floor 100, filter +0.4): mono 0.2639 (599 matched, 880 predicted, 140 merged); ring floor 0.2486 (680 matched, **1259 +predicted**, 45 merged); zero floor 0.2593 (657 / 1097 / 73). The floors recover the merged objects and +release just as many extra instances: the seeds the monotone flooding silently suppressed were the extra +sinks inside the large cells. Only tissuenet gains from a floor (+5 % with size floor 25), because its +losses are small objects deleted by the size filter. The merge rule (edge / interior magnitude ratio +0.7 / 0.85) removes some of the extra instances but merges real neighbours as well (livecell mono +0.2639 → 0.2560 at 0.7). + +Pair probe (adjacent instances under the zero floor, 8 images): same-object pairs vs different-object +pairs on livecell (230 / 1735): edge-over-interior magnitude ratio p50 0.58 vs 0.35 (p25 0.37 vs p75 +0.48 overlap), mid-segment magnitude minimum 0.10 vs 0.07, foreground along the boundary 0.83 vs 0.87, +peak distance 21 vs 35 px, size ratio 0.30 vs 0.56. No cue separates the two populations; a rule that +merges most same-object pairs also merges about a fifth of the real pairs. **M1 is not fixable with +label-free rules on this field**: the decoder blurs the field over 6-8 px, so whether two sinks belong to +one object is not decidable from the prediction; the monotone flooding's implicit arbitration (the seed +with the lower height floods) is as good as any explicit rule tried. The remedy is training-side (a +sharper field, or a target with an explicit contact channel), out of scope here. + +### Decision (2026-09-06 20:15): what goes into epoch A1 + +- **Library (opt-in keyword)**: `boundary_magnitude_max` in `flow_instance_segmentation`, implemented by + `drop_instances_without_boundary_dip` (median |d| along an instance's boundary above the threshold → + the instance is a false foreground region). Default None (off) for every backbone; the default path is + bit-identical. The dense pipeline is untouched (Phase 5). +- **Parameters for the shared-default sweep**: travel to convergence (`n_iter` 800 at `dt` 0.5, the tracer + stops early), `foreground_threshold` (0.4-0.7; the datasets disagree in sign, so it is a compromise: + tissuenet under-covers, deepbacs over-covers), `min_size` (25-100), `density_threshold` (5-50), + `foreground_weight`, `boundary_magnitude_max` (off / 0.4 / 0.6). +- Not adopted: seed floors, decoder-consistency merge, relative or particle-count seeds, trajectory + assignment, height-map transforms, direction / divergence ridges (all recorded above with numbers). + +- 2026-09-06 20:30: **epoch A1 `a65e2eb08c23538f11544860736961a3`** (from `5700c6e0…`): `micro_sam/v2/postprocessing.py` gains + `drop_instances_without_boundary_dip` and the opt-in keyword `boundary_magnitude_max` (default None in + every backbone's table), mirrored in the harness (`sparse_pipeline`) and in the cached sweep scorer + (`parameter_search.score_image_sparse_cached`); tests in `test/test_v2_automatic_segmentation.py`. The + default path is unchanged; the baselines are rerun under this epoch and checked per sample. + +### Look ahead to Phase 5: the dense multicut on the deep EM crops (2026-09-06 21:00) + +AIS defaults on the apg3d manifests (epoch `5700c6e0…`): family macro **0.091** primary / 0.109 holdout +(APG: 0.327 / 0.342). The sparse LM families: celegans_atlas 0.10, gonuclear 0.20 (352 of 726 objects +split, **1127 background seeds**), embedseg_platy_ish 0.35, embedseg_platy_nuclei 0.23, embedseg_skull 0.10, +platynereis_nuclei 0.08 (720 background seeds for 127 objects). The dense EM families: cremi CREMI 0.94 +with 22004 instances for 840 objects, cremi_seen 0.40 (27698 / 7469), snemi 0.82 (14066 / 526), +humanneurons 1.30 (49939 / 1601). + +One cremi and one snemi crop by hand: the slice-wise oversegmentation already produces 15104 / 5566 +fragments for 295 / 93 objects (the EM foreground is predicted at 0.71 on average with 29 % of the +neuron voxels below 0.5, so the seeds shatter every cross-section and most fragment boundaries look like +membranes: median edge boundary value 0.74 / 0.81), and the multicut at `beta` 0.5 → 0.95 goes from 7400 +to 11456 instances (CREMI 1.02 → 2.08) — in elf's `compute_edge_costs` a **higher beta cuts more**, the +opposite of the `run_multicut` docstring ("higher values favour more merging"), and `EM_GRID` (0.5-0.8) +never enters the merging regime (< 0.5). Both the seeding granularity (fewer, larger fragments; the +boundary filter does not apply, the fragments are not instances) and the beta range are Phase 5 items. + +- 2026-09-06 21:05: epoch A1 baselines (`current-defaults`) on v5 primary / training_extra / holdout and + apg3d primary / holdout are identical per sample to the epoch `5700c6e0…` runs (755 samples: mSA and + instance counts equal, 0 pipeline mismatches); balanced 0.1841 / 0.4183 / 0.1826, apg3d dataset-balanced + 0.1091 / 0.1221. Filter screens `a1_filter_2d` (job 15767179, 27 tasks) and `a1_filter_3d` (15767180, + 18 tasks) and the sweeps `a1_sweep_primary` / `a1_sweep_extra` (15767181 / 15767182, grid + `configs/ais_grid_lm_v4.json`, 1728 combinations) submitted at 19:55. + +## Epoch A1 screen (2026-09-06 21:10, jobs 15767179 / 15767180, reports `ais/reports/a1_filter_*.csv`) + +Relative change of mSA against the defaults; "balanced" over the eleven development datasets (2D) or the +six sparse LM sources of the deep 3D crops (the dense EM sources are untouched by these parameters). + +| configuration | 2D dev balanced | up / 11 | worst | 2D holdout (5) | 3D LM primary (6) | 3D LM holdout (6) | +|---|---:|---:|---:|---:|---:|---:| +| filter 0.4, travel 25 (defaults otherwise) | **+1.4 %** | **9** | −0.1 % | +1.0 % | **+4.7 %, 5/6 up, passes** | **+9.7 %, 5/6 up, passes** | +| filter 0.6, travel 25 | +0.9 % | 7 | 0.0 % | +0.5 % | +0.8 % | +4.5 % | +| travel 400 alone | +0.3 % | 3 | −1.7 % (tnbc) | +2.5 % | +22 % (skull +281 %, platy_ish −6 %, platy_nuclei −6 %) | +19 % | +| filter 0.4, travel 400 | +1.7 % | 6 | −0.9 % (tissuenet) | **+3.5 %, 5/5 up, passes** | +27 % (2 sources down) | +28 % | +| filter 0.3, travel 400 | +1.8 % | 6 | −1.1 % | +3.5 % | +32 % | +35 % | +| filter 0.4, travel 400, min_size 25 | −0.3 % | 3 | −6.7 % (tnbc) | +4.5 % (tissuenet +12 %) | +2.6 % | +2.2 % | +| filter 0.4, travel 400, foreground 0.6 | −0.9 % | 5 | −6.5 % (dnn) | +1.0 % | +29 % | +27 % | + +Per dataset, filter 0.4 at travel 25 (2D): deepseas +23.6 %, dic_hepg2 +10.6 %, deepbacs +4.4 %, +neurips_cellseg +4.2 %, dynamicnuclearnet +1.4 %, tnbc +0.9 %, puma +0.2 %, covid_if / livecell / yeaz +0.0 %, tissuenet −0.1 %; 3D primary: embedseg_skull +25 %, platynereis_nuclei +16 %, gonuclear +5 %, +celegans_atlas +1.4 %, platy_nuclei +0.4 %, platy_ish 0.0 %; 3D holdout: platy_nuclei +35 %, platynereis ++29 %, skull +8 %, celegans +1 %. Travel 400 with the filter reaches +19 % on deepbacs and +18 % on +deepseas but costs covid_if, tissuenet, tnbc, yeaz 0.5-0.9 % each and, in 3D, the two large EmbedSeg +nuclei sources 5-7 % (the converged sinks split large nuclei). + +Reading: the boundary filter is a generalizing improvement (never below −0.1 % on any of the 22 dataset +× subset cells, up wherever background seeds exist) but alone it stays under the +2 % balanced bar in +2D; the travel is the second lever in 3D and on the small-object 2D data and needs a compensating change +where it splits large objects. The shared-default sweep (`configs/ais_grid_lm_v4.json`, 1728 +combinations over the eleven 2D datasets; a reduced grid over the six 3D LM sources) decides the +combination. + +## Phase 4: shared-default sweep, 2D (2026-09-06 21:40, jobs 15767181 / 15767182, ranking `ais/reports/a1_sweep_dev_ranking.csv`) + +Grid `configs/ais_grid_lm_v4.json` (1728 combinations: foreground 0.4-0.7, density 5-50, size floor +25 / 50 / 100, sigma 0.5 / 1.0, travel 25 / 400 px, foreground weight 0.25 / 0.5 / 0.75, filter off / 0.4 / +0.6) scored on every image of the eleven development datasets from the cache (`sweep`, 8-20 s per +image). `report_ais_sweep.py` ranks the combinations as shared defaults against the library defaults. +**7 of 1728 pass the gate**; all seven keep foreground 0.5, density 10, foreground weight 0.5 and use +sigma 1.0. + +| combination (changes to the defaults) | balanced | gain | up / 11 | worst | per dataset | +|---|---:|---:|---:|---:|---| +| defaults | 0.3357 | | | | | +| filter 0.4 | 0.3404 | +1.4 % | 9 | −0.1 % | deepseas +24, dic +11, deepbacs +4, neurips +4 | +| sigma 1.0 | 0.3420 | +1.9 % | 8 | −1.3 % (tnbc) | deepbacs +10, deepseas +30, livecell +2.9, neurips +6, tissuenet −1.2 | +| min_size 50 | 0.3333 | −0.7 % | 2 | −13.5 % | tissuenet +9.7, everything else down: the lower floor alone admits the small spurious seeds | +| sigma 1.0 + min_size 50 | 0.3430 | +2.2 % | 9 | −1.1 % | tissuenet +6.4 (the smoothing removes the spurious seeds the lower floor would keep) | +| **C1: sigma 1.0 + min_size 50 + filter 0.4** | **0.3437** | **+2.4 %** | **9** | **−0.8 % (tnbc)** | covid −0.2, deepbacs +12.7, deepseas +31.9, dic +59.6, dnn +0.4, livecell +3.1, neurips +4.4, puma +0.4, tissuenet +6.3, tnbc −0.8, yeaz +0.5 | +| C1 + travel 400 | 0.3437 | +2.4 % | 6 | −0.8 % | deepbacs +17.4, tissuenet +8.8, livecell +4.0, dnn +1.1; deepseas +20, five datasets −0.0 to −0.8 | +| sigma 1.0 + travel 400 + filter 0.4 (size floor 100) | 0.3431 | +2.2 % | 9 | −0.5 % | | + +The best by mean ratio to each dataset's optimum (0.878) is foreground 0.4 / filter 0.6, which fails the +gate (6 up); C1 is second (0.877). Reading: the sweep changes the *interpretation* of the seeding rather +than its logic: a wider smoothing of the convergence density (sigma 1.0) merges the jittering sinks of a +large cell into one seed and drops the isolated one-pixel seeds (tissuenet's were 1 px), which is what +the seed floors and the merge rule tried and failed to do structurally; with those gone the size floor can +follow the ground-truth floors (50), and the boundary filter removes the remaining false regions. Travel +to convergence is neutral in 2D (same balanced, more datasets marginally down); the 3D sweep decides it. + +Confirmation (job `a1_confirmation`, one task per manifest, trial `timing-1`, control and candidates on +the same node): `configs/ais_c1_sigma1_ms50_filter0p4.json`, `ais_c1_t400.json`, `ais_c3_sigma1_ms50.json` +on v5 primary / training_extra / holdout and apg3d primary / holdout. + +## Phase 4: shared-default sweep, 3D LM crops (2026-09-06 22:15, job 15767263, ranking `ais/reports/a1_sweep_3d_primary_ranking.csv`) + +Grid `configs/ais_grid_lm3d_v4.json` (576 combinations; size floor 50 / 100 / 200 voxels, foreground +weight fixed at 0.5) on the six sparse LM sources of the 57 primary deep crops. **22 of 576 pass the gate** +(all six sources up in the best of them). Balanced over the six sources, defaults 0.1714: + +| combination (changes to the defaults) | balanced | gain | up / 6 | worst | celegans | platy_ish | platy_nuclei | skull | gonuclear | platynereis | +|---|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:| +| C1 (sigma 1.0, min_size 50, filter 0.4) | 0.1804 | +5.3 % | 3 | −6.1 % | −6.1 | −0.3 | −0.7 | +27 | +8.7 | +43 | +| sigma 1.0, filter 0.4 (min_size 100) | 0.1910 | +11.4 % | 5 | −5.8 % | −5.8 | +2.5 | +0.7 | +56 | +14 | +66 | +| sigma 1.0, filter 0.4, min_size 200 | 0.2030 | +18.4 % | 5 | −4.6 % | −4.6 | +2.9 | +1.0 | +109 | +16 | +95 | +| **sigma 1.0, filter 0.4, min_size 200, foreground 0.6** | **0.2081** | **+21.4 %** | **6** | **+1.4 %** | +8.2 | +1.4 | +2.0 | +108 | +25 | +106 | +| same with foreground 0.7 | 0.2073 | +20.9 % | 5 | −4.3 % | +18 | −4.3 | +0.1 | +100 | +32 | +108 | +| sigma 1.0, filter 0.4, min_size 100, travel 400 | | +29.4 % | 3 | −3.5 % | | | | | | | +| top by balanced: density 20, foreground 0.6, min_size 200, sigma 1.0, filter 0.4 | 0.2644 | +54 % | 3 | −7.3 % | the two EmbedSeg platy sources lose | + +Joint view over the 384 combinations both grids share: 4 pass the 2D gate, 11 the 3D gate, **none both**. +The disagreement is the size floor (50 px is right for the 2D nuclei data, 200 voxels for the volumes: a +voxel floor of 50 keeps fragments that no 3D object is) and the foreground threshold (celegans_atlas turns +from −6 % to +8 % between 0.5 and 0.6, while 0.6 costs dynamicnuclearnet, tissuenet and tnbc in 2D). Sigma +1.0, the boundary filter at 0.4, density 10, foreground weight 0.5 and the default travel are shared by the +winners of both dimensions. Travel to convergence does not enter the 3D winners either (it splits the large +EmbedSeg nuclei), so the runtime stays as it is. + +Proposal: dimension-aware defaults, as the APG module has for volumes (`default_prompt_generation(..., +is_volume=True)`): images `{sigma 1.0, min_size 50, boundary_magnitude_max 0.4}` on top of the current +table; volumes additionally `{min_size 200, foreground_threshold 0.6}`. The volume part is confirmed on the +3D holdout before anything is promoted (`configs/ais_c1v_volume.json`). + +## Confirmation of C1 (2026-09-06 22:30, job 15767364, trial `timing-1`, control and candidates on one node per manifest) + +| configuration | 2D dev (11) | up | worst | 2D holdout (5) | up | worst | 3D LM primary (6) | 3D LM holdout (6) | +|---|---:|---:|---:|---:|---:|---:|---:|---:| +| **C1** sigma 1.0, min_size 50, filter 0.4 | **+2.4 %, passes** | 9 | −0.8 % | **+4.3 %, passes** | 5 | +0.9 % | +5.6 % (3 up, celegans −6.1 %) | +1.4 % (2 up, celegans −5.0 %) | +| C1 + travel 400 | +2.4 % | 6 | −0.8 % | +5.7 %, passes | 5 | +0.6 % | +18 % (2 up, platy_ish −7 %) | +21 % (2 up, celegans −7 %) | +| C3 sigma 1.0, min_size 50 (no filter) | +2.2 %, passes | 9 | −1.1 % | +4.1 % | 4 | −3.4 % (dic) | +2.1 % | −2.6 % | + +2D holdout per dataset, C1: deepbacs +12.7 %, dic_hepg2 +2.3 %, dynamicnuclearnet +0.9 %, livecell +2.8 %, +tissuenet +8.1 %. Object counts on the development images: +526 matched, −2589 objects with two or more +seeds, −2453 background seeds, −1364 splits (of 38 000 objects). C1 confirms in 2D; as a volume setting it +fails on celegans_atlas, which is what the volume overrides (foreground 0.6, size floor 200) address +(confirmation job `a1_confirmation_volume`, trial `timing-2`). + +Post-processing time on the same node: v5 primary 6.0 s (defaults) → 15.3 s (C1) for 240 images, i.e. +0.04 s +per image, all of it the boundary filter (sigma / size floor alone: 12.1 s → the density smoothing is not +the cost; C3 on holdout 5.2 s); apg3d primary 32.8 s → 159 s for 57 crops (+2.2 s per 32-slice crop). + +Runtime of the confirmation (2D images, prediction time from the A100 cache records, post-processing on +one cluster node): C1 with the first filter implementation +4 % to +31 % total per dataset (the filter's +`scipy.ndimage.median` per instance cost 0.06 s per image and 2.9 s per 32-slice crop); the same +configuration without the filter (C3) +0 % to +9 % (livecell +9.3 %, the wider density smoothing). The filter +was then rewritten (2026-09-06 22:45): the inner boundary from axis shifts and every instance's median from one +`lexsort` over the boundary pixels (mean of the two middle values for even counts, as `ndimage.median`) — +identical output on 60 images and 3 crops, 18× faster (3 ms per 512² image, 0.12 s per crop), so the +runtime overhead of C1 is that of C3. + +`default_postprocessing` gained an `ndim` argument and the table a `sparse_volume` sub-table of volume +overrides (empty until the volume confirmation), which `flow_instance_segmentation` resolves from the +foreground's dimensionality and the harness from `params_2d` / `params_3d`. + + +## Volume confirmation and promotion (2026-09-06 23:05, job 15767420, trial `timing-2`) + +`c1v-volume` (images: sigma 1.0, min_size 50, filter 0.4; volumes: the same plus min_size 200 and +foreground 0.6): 3D LM primary **+22.7 %, 6 / 6 up** (celegans +8.2, platy_ish +1.4, platy_nuclei +2.0, +skull +108, gonuclear +25, platynereis +90); 3D LM holdout **+22.0 %, 5 / 6 up, worst −0.2 %** (celegans ++20, platy_ish −0.2, platy_nuclei +38, skull +98, gonuclear +1.6, platynereis +131); the twelve-slice +volumes of the 2D manifests: celegans +14.5 %, embedseg +39.6 %, gonuclear +18 %; the images are C1 +(+2.4 %). Without the foreground change (`c1v-ms200`) holdout is +19.8 % with celegans −1.8 %. + +**Epoch A2 `576a85c8ffd4314627812fd30a3c1223`: promoted.** `DEFAULT_POSTPROCESSING["hvit_t"]["sparse"]` = foreground 0.5, +density 10, min_size 50, sigma 1.0, n_iter 50, dt 0.5, foreground weight 0.5, boundary_magnitude_max 0.4; +`["sparse_volume"]` = min_size 200, foreground 0.6. The other backbones keep their registry values and +an empty volume table; the dense pipeline is unchanged. The old values remain reachable as an explicit +configuration (`configs/ais_control_v4_old_defaults.json`, filter off via `Infinity`). + +## Phase 6: canonical runs, production and the 3D test manifest (submitted 2026-09-06 23:15) + +- Canonical A2 screens (job 15767503 `a2_canonical`, trial `a2-1`): `v4-old-defaults` (explicit old values) + against `current-defaults` (the promoted library defaults) on v5 primary / training_extra / holdout and + apg3d primary / holdout, plus `--ndim 2` runs of v5 primary and holdout for + `compare_apg_optimization.py --target quality`. +- Production (`submit_all_evaluations.py --segmentation_type automatic --segmentation_mode ais + --all_datasets --modality lm -m hvit_t --skip_tuning`, experiment folder + `experiments/v4_geodesic_ais_optimization`): jobs 15767555-57 with the new defaults (result tag + `a2-defaults`) and 15767589-91 with `--ais_params configs/ais_control_v4_old_defaults.json` (tag + `old-defaults`), 33 LM datasets each (23 2d + 10 3d LM; the dense EM pipeline is unchanged, its §14.3 + numbers stand). Reader: `report_ais_production.py -e --baseline default_old-defaults --candidate + default_a2-defaults`, which reports the twelve strictly unseen 2d datasets separately. +- 3D test manifest (`manifest_test_apg3d-v1.json`, 56 crops of the seven test-only LM datasets, opened + once): predictions cached on the session GPU, then `screen` old vs new defaults (trial `test-1`). + +### Canonical A2 runs (2026-09-06 23:50, job 15767503, trial `a2-1`; reports `ais/reports/a2_*`) + +New library defaults (`current-defaults`) against the explicit old values (`v4-old-defaults`), one node per +manifest: + +| instrument | old | new | change | up | worst | +|---|---:|---:|---:|---:|---:| +| 2D development, eleven datasets (generalization gate) | 0.3357 | 0.3437 | **+2.4 %, passes** | 9 / 11 | −0.8 % (tnbc) | +| 2D holdout, five datasets | 0.2337 | 0.2437 | **+4.3 %, passes** | 5 / 5 | +0.9 % | +| 3D LM deep crops, primary (6 sources) | 0.1765 | 0.2165 | **+22.7 %, passes** | 6 / 6 | +1.4 % | +| 3D LM deep crops, holdout (6 sources) | 0.1998 | 0.2438 | **+22.0 %, passes** | 5 / 6 | −0.2 % | + +`compare_apg_optimization.py --target quality` on the 2D-only runs of the five primary datasets: macro mSA +0.2366 → 0.2457 (**+3.9 %**) on primary and 0.2337 → 0.2437 (**+4.3 %**) on holdout, every dataset up +(deepbacs +12.7 %, tissuenet +6.3 / +8.1 %, livecell +3.1 / +2.8 %, dynamicnuclearnet +0.4 / +0.9 %, +dic_hepg2 +60 / +2 % on a near-zero base), every dataset's runtime within +2.9 % (total +1.3 % / +0.6 %, +post-processing on one node, prediction time from the cache records). The comparator's quality gate is +**not** met because its macro bar is +5 % on this five-dataset instrument; its runtime and per-dataset +checks pass, and the generalization gate of the campaign (the eleven-dataset rule) passes on every +instrument. Recorded as such: the promotion rests on the generalization rule, not on the +5 % quality +target that was set for the APG campaigns. + +## 3D test manifest, opened once (2026-09-07 00:00, screen `a2_apg3d_test`, trial `test-1`) + +56 crops, eight per test-only dataset, new defaults (with the volume overrides) against the old values: + +| dataset | old | new | change | matched old → new | +|---|---:|---:|---:|---| +| blastospim | 0.1209 | 0.1203 | −0.5 % | 61 → 61 | +| cartocell | 0.0142 | 0.0120 | −15 % (−0.002 absolute) | 24 → 22 | +| cellseg_3d | 0.000 | 0.000 | n/a (nothing matched either way) | 0 → 0 | +| mouse_embryo | 0.0663 | 0.0823 | **+24 %** | 217 → 218 | +| nis3d | 0.1079 | 0.1017 | −5.8 % (−0.006) | 762 → 622 | +| plantseg | 0.1573 | 0.2115 | **+35 %** | 553 → 497 | +| pnas_arabidopsis | 0.2889 | 0.2690 | −6.9 % (−0.020) | 1437 → 1362 | +| balanced (7) | 0.1083 | 0.1141 | +5.4 % | 3072 → 2784; background seeds 4065 → 783; unseeded objects 3065 → 3930 | + +Two of six scorable datasets up, two below the loss limit: **the volume overrides (foreground 0.6, size +floor 200 voxels) do not pass the out-of-domain check**, although they passed the tuning crops (+22 % on +primary and holdout). The pattern (background seeds −80 %, unseeded objects +28 %, matched −9 %) says the +filter does its job while the higher foreground threshold and the voxel floor remove real objects on the +unseen nuclei data (nis3d, pnas_arabidopsis). Decision rule, fixed before looking further: the volume +defaults fall back to the only volume candidate that passed the gate on both tuning instruments without a +dataset down, the old values plus the boundary filter (A1 screen: +4.7 % / +9.7 %, worst 0.0 %); it is +evaluated once on the test manifest (`configs/ais_v_filter_only.json`) and, if it fails too, volumes revert +to the old values. The image defaults are unaffected. The library change waits until the running production +jobs (which import the library per dataset) have finished, so that their `a2-defaults` results stay what +their tag says. + +- 2026-09-07 00:10: the 66 production jobs of 23:15 all failed at start-up (`torch.save ... Permission denied`): + `build_model(mode="ais")` exports the decoder into `MICRO_SAM2_JOINT_EXPORT_ROOT`, which was not set in the + submitting shell, so the library's default export root (not writable for this user) was used. Both + variables must be exported before submitting (EXPERIMENTAL_SETUP.md §2 pins both). Resubmitted with + `MICRO_SAM2_JOINT_EXPORT_ROOT=/model_exports` (the v4 export already exists there). + +### Volume fallback on the test manifest and epoch A3 (2026-09-07 00:30) + +`v-filter-only` (volumes: registry values + filter 0.4) on the 56 test crops: **+3.4 %, 6 / 6 scorable +datasets up** (blastospim +1.2, cartocell +1.9, mouse_embryo +6.8, nis3d +2.9, plantseg +9.3, +pnas_arabidopsis +0.7; cellseg_3d 0 either way), worst +0.7 %: passes. With the tuning crops (+4.7 % +primary, +9.7 % holdout, nothing down) this is the volume setting that generalizes. + +**Epoch A3 `e9d02380e340edfaccd30bf5cbf1bf03`: `DEFAULT_POSTPROCESSING["hvit_t"]["sparse_volume"]` = min_size 100, sigma 0.5** +(the registry values; foreground 0.5 and the filter 0.4 are inherited from the image table). Images +unchanged from A2. The 3D production jobs tagged `a2-defaults` were cancelled before they could import the +new table and are resubmitted as `a3-defaults`; the 2D jobs are unaffected (image table unchanged). + +### Canonical A3 runs on the 3D crops (2026-09-07 01:00, job 15767918, trial `a3-1`; reports `ais/reports/a3_apg3d_*`) + +Library defaults (`current-defaults`, epoch `e9d02380…`) against the old values, sparse LM sources: + +| instrument | old | new | change | up | worst | per source | +|---|---:|---:|---:|---:|---:|---| +| primary (57 crops) | 0.1765 | 0.1847 | **+4.7 %, passes** | 5 / 6 | 0.0 % | celegans +1.4, platy_ish 0.0, platy_nuclei +0.4, skull +25, gonuclear +5.2, platynereis +16 | +| holdout (18 crops) | 0.1998 | 0.2193 | **+9.7 %, passes** | 5 / 6 | 0.0 % | celegans +1.0, platy_ish 0.0, platy_nuclei +35, skull +8.4, gonuclear +0.2, platynereis +29 | +| test manifest (56 crops, opened once) | 0.1083 | 0.1120 | **+3.4 %, passes** | 6 / 6 scorable | +0.7 % | blastospim +1.2, cartocell +1.9, mouse_embryo +6.8, nis3d +2.9, plantseg +9.3, pnas_arabidopsis +0.7 | + +Identical to the `v-filter-only` configuration, as intended; the dense EM sources are unchanged. + +## Production, 2D test splits (2026-09-07 01:30; `experiments/v4_geodesic_ais_optimization/results/`, report `ais/reports/production_2d_old_vs_new.csv`) + +`evaluate_automatic_segmentation.py --skip_tuning` on the full test split of every 2D dataset, old defaults +(`--ais_params configs/ais_control_v4_old_defaults.json`, tag `old-defaults`) against the new library defaults +(tag `a2-defaults`; images are identical under A2 and A3). mSA: + +| dataset | old | new | change | | dataset | old | new | change | +|---|---:|---:|---:|---|---|---:|---:|---:| +| livecell | 0.2575 | 0.2660 | +3.3 % | | arvidsson* | 0.3581 | 0.3554 | −0.8 % | +| tissuenet | 0.2508 | 0.2583 | +3.0 % | | bitdepth_nucseg* | 0.2298 | 0.2340 | +1.8 % | +| dynamicnuclearnet | 0.5083 | 0.5509 | +8.4 % | | cellbindb* | 0.2787 | 0.2961 | +6.2 % | +| deepbacs | 0.2056 | 0.2319 | +12.8 % | | cellpose_data* | 0.1982 | 0.2063 | +4.1 % | +| dic_hepg2 | 0.0018 | 0.0028 | +55 % | | cvz_fluo* | 0.1404 | 0.1507 | +7.3 % | +| yeaz | 0.5964 | 0.6021 | +0.9 % | | dsb* | 0.4631 | 0.4862 | +5.0 % | +| neurips_cellseg | 0.2916 | 0.3138 | +7.6 % | | hpa* | 0.0003 | 0.0003 | +15 % | +| deepseas | 0.1048 | 0.1549 | +47.8 % | | **microbeseg*** | **0.1420** | **0.1258** | **−11.4 %** | +| puma | 0.4556 | 0.4613 | +1.3 % | | omnipose* | 0.2153 | 0.2537 | +17.8 % | +| covid_if | 0.7656 | 0.7686 | +0.4 % | | segpc* | 0.0066 | 0.0116 | +76 % | +| tnbc | 0.3277 | 0.3470 | +5.9 % | | usiigaci* | 0.0891 | 0.0995 | +11.7 % | +| | | | | | vicar* | 0.4032 | 0.4100 | +1.7 % | + +\* strictly unseen by any tuning. **All 23: 21 up, balanced 0.2735 → 0.2864 (+4.7 %)**; the eleven development +datasets 0.3423 → 0.3598 (+5.1 %, 11 up); the twelve unseen datasets 0.2104 → 0.2191 (+4.2 %, 10 up). +The gate's production variant (loss limit −5 % and 0.005) is violated by one dataset, microbeseg +(−0.016 absolute). Single-change ablations on microbeseg (tags `d-filter-only`, `d-sigma-only`, +`d-ms50-only`): filter alone 0.1420 (no change), size floor 50 alone 0.1580 (+11 %), **sigma 1.0 alone 0.1208 +(−15 %)**: the wider density smoothing merges the small, dense bacteria of microbeseg, the same mechanism +that lets it merge the jittering sinks of large cells everywhere else. This is the one known cost of the new +defaults; it is reported, not tuned away (the test split is not a tuning set). Against the §14.3 v4 numbers +the new AIS defaults now beat the v2 AIS defaults on every dataset that was compared there (deepbacs +0.2319 vs v2 0.2940 remains below v2). + +## Production, 3D LM test splits (2026-09-07 02:30; tags `old-defaults` vs `a3-defaults`, report `ais/reports/production_3d_old_vs_new.csv`) + +| dataset | old | new | change | | dataset | old | new | change | +|---|---:|---:|---:|---|---|---:|---:|---:| +| blastospim | 0.0644 | 0.0658 | +2.3 % | | mouse_embryo | 0.0341 | 0.0344 | +1.0 % | +| cartocell | 0.0088 | 0.0088 | +0.2 % | | nis3d | 0.1037 | 0.1039 | +0.2 % | +| celegans_atlas | 0.1118 | 0.1125 | +0.7 % | | plantseg | 0.1371 | 0.1469 | +7.2 % | +| cellseg_3d | 0.0000 | 0.0000 | 0 (nothing matched either way) | | pnas_arabidopsis | 0.3160 | 0.3162 | +0.1 % | +| embedseg | 0.4105 | 0.4310 | +5.0 % | | gonuclear | 0.2689 | 0.2873 | +6.8 % | + +**9 of 10 up, none down, balanced 0.1455 → 0.1507 (+3.6 %), gate passed.** The dense EM datasets are +unchanged (their §14.3 numbers stand). + +## Status (2026-09-07 02:30) + +Done: Phases 0-4 and 6 of the plan, Phase 5 for the sparse pipeline (3D crops, holdout, test manifest). +Promoted (epoch A3, commit 117f210 and the notes commits after it): `hvit_t` AIS defaults images +`sigma 1.0, min_size 50, boundary_magnitude_max 0.4`, volumes `min_size 100, sigma 0.5` plus the filter; +the new `drop_instances_without_boundary_dip` and the dimension-aware `default_postprocessing`. Every +instrument of the protocol improved (2D dev +2.4 %, 2D holdout +4.3 %, 3D primary +4.7 %, 3D holdout ++9.7 %, 3D test manifest +3.4 %, production 2D +4.7 % with 21 / 23 up, production 3D LM +3.6 % with 9 / 10 +up). Known cost: microbeseg −11.4 % on its test split (sigma 1.0; a decision for the user), arvidsson −0.8 %. +Open: Phase 5.3, the dense multicut (beta direction, oversegmentation granularity), and the C++ port of the +filter is not needed (the numpy version is 3 ms per image). Unfinished ideas that did not pass and should +not be retried without a sharper decoder field: seed floors, decoder-consistency merge, relative / +particle-count seeds, trajectory assignment, direction ridges (numbers above). + + +## Follow-up screen: seed floors on the promoted defaults (2026-09-07) + +Epoch A4 `184eba917bd0cff28b5719b5584f6967`: opt-in keyword `seed_floor` ('none' default, 'zero', 'ring') in +`flow_instance_segmentation`, implemented by `lower_height_under_seeds` (mirrored in the harness and the +sweep scorer, default path unchanged). Rationale: the monotone flooding lets a seed on a height peak lose its +object (the merge mechanism); the earlier floor test on the old defaults failed because it also released the +spurious seeds, which the promoted defaults (sigma 1.0, filter 0.4) now remove. Screen `f2_floor`: floors +zero / ring, each with the promoted size floor and with min_size 100, against the promoted defaults, on +v5 primary / training_extra / holdout and apg3d primary / holdout (trial `f2-1`, one node per manifest). + +### Result of the seed-floor screen (2026-09-06 23:45, job 15768572, trial `f2-1`, baseline = promoted A3 defaults) + +| configuration | 2D dev (11) | 2D holdout (5) | 3D LM holdout (6) | 3D LM primary (6) | +|---|---:|---:|---:|---:| +| zero floor | −0.6 % (4 up; livecell +2.5, tissuenet +5.5, deepbacs −6.7, tnbc −2.5, neurips −2.2) | −1.2 % (2 up) | **−51.5 %** (0 up, skull −84 %) | **−40.6 %** (0 up) | +| ring floor | −1.3 % (3 up; dic_hepg2 −24 %, deepbacs −8.9) | −1.2 % (2 up) | **−25.9 %** (1 up) | **−24.6 %** (2 up, skull −60 %) | +| zero / ring with min_size 100 | −1.1 % / −1.8 % | −3.1 % / −3.2 % | identical to the above (100 is the volume floor already) | | + +Object counts (2D dev, zero floor): merges −1895, matched +1476, but the released seeds add instances +faster than they add matches, so mSA falls on 7 of 11 datasets; in 3D every large nucleus carries several +weak sinks and all of them flood. **Not adopted in either dimension; the A3 defaults stay.** The seed floor +stays available as the opt-in keyword `seed_floor`. This closes the last post-processing lever the +diagnostics pointed at: the remaining merges need a sharper field from the decoder +(`AIS_DECODER_TRAINING_PROPOSAL.md` at the repository root). diff --git a/finetuning/v2/evaluation/optimization/notes/AIS_V4_OPTIMIZATION_SUMMARY.md b/finetuning/v2/evaluation/optimization/notes/AIS_V4_OPTIMIZATION_SUMMARY.md new file mode 100644 index 000000000..6ebca046e --- /dev/null +++ b/finetuning/v2/evaluation/optimization/notes/AIS_V4_OPTIMIZATION_SUMMARY.md @@ -0,0 +1,86 @@ +# AIS optimization for the joint/v4 geodesic `hvit_t` model: summary + +Concise findings of the 2026-09-06/07 campaign. The full decision log with every number, job id and run +directory is `AIS_V4_OPTIMIZATION.md`; the set-up is `EXPERIMENTAL_SETUP.md` (§13-14); the decoder-side +follow-up is `AIS_DECODER_TRAINING_PROPOSAL.md` at the repository root (uncommitted). + +## Outcome + +New `hvit_t` defaults for the flow (sparse) post-processing in `micro_sam/v2/postprocessing.py`: + +| | old (registry) | new, images | new, volumes | +|---|---|---|---| +| density smoothing `sigma` | 0.5 | **1.0** | 0.5 | +| size floor `min_size` | 100 | **50** | 100 | +| instance filter `boundary_magnitude_max` | off | **0.4** | **0.4** | +| seed floor `seed_floor` | (monotone flooding) | 'none' (floors screened, not adopted) | 'none' | +| foreground 0.5, density 10, travel 25 px, foreground weight 0.5 | unchanged | | | + +`drop_instances_without_boundary_dip` is new library logic: the geodesic decoder's distance magnitude falls +to zero along every real object boundary, so an instance whose boundary median exceeds the threshold is a +false foreground region (3 ms per image). `default_postprocessing` became dimension-aware (`sparse_volume` +overrides). The dense multicut is untouched. + +## Results (mSA, new defaults against the old ones) + +| instrument | old | new | change | datasets up | +|---|---:|---:|---:|---| +| 2D development corpus, 11 datasets | 0.3357 | 0.3437 | +2.4 % | 9 / 11 (worst −0.8 %) | +| 2D holdout, 5 datasets | 0.2337 | 0.2437 | +4.3 % | 5 / 5 | +| 3D LM deep crops, primary / holdout | 0.1765 / 0.1998 | 0.1847 / 0.2193 | +4.7 % / +9.7 % | 5 / 6, none down | +| 3D test manifest (7 test-only datasets, opened once) | 0.1083 | 0.1120 | +3.4 % | 6 / 6 scorable | +| production 2D test splits, 23 datasets | 0.2735 | 0.2864 | +4.7 % | 21 / 23 | +| of which the 12 never used for tuning | 0.2104 | 0.2191 | +4.2 % | 10 / 12 | +| production 3D LM test splits, 10 datasets | 0.1455 | 0.1507 | +3.6 % | 9 / 10, none down | + +Known costs: microbeseg −11.4 % (0.142 → 0.126, caused by the wider smoothing alone, which merges its +small dense bacteria) and arvidsson −0.8 %. The comparator's quality gate (+5 % macro on the five primary +datasets) reads +3.9 % / +4.3 % with every dataset up and runtime within +3 %; the campaign's generalization +gate passes on every instrument. APG remains ahead (2D primary 0.296 vs 0.246; 3D crops 0.33 vs 0.18). + +## What the diagnostics established + +- The v4 field is centre-directed and its magnitude dips at object centres; the old post-processing was + tuned to the v2 medial-axis field. Travel to convergence is *not* the fix (+7 % balanced but 5 / 16 up, + more background seeds). +- Loss decomposition of the old defaults: merges dominate the touching-cell data (livecell 25 % merged + + 11 % absorbed, tissuenet 20 % + 5 %), background seeds dominate deepseas (2× the object count), deepbacs and + neurips_cellseg; `min_size` 100 alone cost tissuenet 7.5 % of its matches. +- Oracles: a ground-truth ridge with the predicted seeds doubles livecell / deepbacs; ground-truth seeds add + +5-20 % (+50-80 % on small-object data); the ground-truth foreground ceiling is the largest but mixes + extent with separation. +- Mechanism of the merges: the watershed floods monotonically and every proper seed sits on a height peak + (the centre dip), so a seed with a deeper dip than the contact loses its object. The same property + suppresses spurious seeds, which is why a plain seed floor lost on the old defaults. +- What did not generalize (all recorded with numbers): relative / particle-count seeds (split large cells), + trajectory assignment (inherits the blurred field), direction and divergence ridges (the network smooths + the flip over 6-8 px), height-map transforms, a decoder-consistency merge (same-object vs different-object + seed pairs are not separable), stronger volume settings (won on the tuning crops, lost on the test manifest). +- What did: wider density smoothing (merges the jittering sinks of large cells, removes one-pixel seeds), + the ground-truth-like size floor once the spurious seeds are gone, and the boundary filter. + +## Follow-up screen: seed floors on the promoted defaults (not adopted) + +The one untested lever left by the diagnostics was to lower the height map under the seeds (so that the +monotone flooding cannot hold a seed's front at its own height), now that the promoted defaults remove the +spurious seeds that sank the same idea on the old defaults. Opt-in keyword `seed_floor` ('zero', 'ring'). +Against the promoted defaults: 2D development −0.6 % (zero) / −1.3 % (ring), 4 and 3 of 11 datasets up +(livecell +2.5 %, tissuenet +5.5 %, deepbacs −6.7 %); 2D holdout −1.2 %; 3D LM crops −40 to −52 % (zero) and +−26 % (ring) with no source up: the floors halve the merges but release more instances than they recover, +and in 3D every large nucleus carries several weak sinks that all flood. **The A3 defaults stay for images +and volumes**; the keyword remains available. The remaining merges need a sharper field from the decoder +(`AIS_DECODER_TRAINING_PROPOSAL.md`). + +## Where things are + +- Library: `micro_sam/v2/postprocessing.py` (defaults, `drop_instances_without_boundary_dip`, + `lower_height_under_seeds`, dimension-aware `default_postprocessing`); tests in + `test/test_v2_automatic_segmentation.py`. +- Harness: `finetuning/v2/evaluation/optimization/benchmark_ais_optimization.py` (predict / run / screen / + sweep / oracle / report on cached predictions), `ais_campaign_tasks.py`, `report_ais_sweep.py`, + `report_ais_production.py`, configurations `configs/ais_*.json`; unit tests `test/test_ais_optimization.py`. +- Data: caches, run directories, sweeps, oracles and reports under `/ais/`; production results under + `experiments/v4_geodesic_ais_optimization/results/` (tags `old-defaults`, `a2-defaults` for 2D, + `a3-defaults` for 3D). +- Open: the dense multicut (Phase 5.3: elf's `beta` cuts more when higher, the docstring says the opposite, + and the EM crops shatter into 15k fragments), and the decoder-side changes of the proposal. diff --git a/finetuning/v2/evaluation/optimization/notes/CAMPAIGN_OPERATIONS.md b/finetuning/v2/evaluation/optimization/notes/CAMPAIGN_OPERATIONS.md index dcfb6e9e6..ba789ab8c 100644 --- a/finetuning/v2/evaluation/optimization/notes/CAMPAIGN_OPERATIONS.md +++ b/finetuning/v2/evaluation/optimization/notes/CAMPAIGN_OPERATIONS.md @@ -17,7 +17,8 @@ Written for the campaigns started on 2026-09-02; the facts about the cluster wer ## Cluster and environment -- Environment: `micromamba activate super`. The evaluation and optimization launchers default to `super`. +- Environment: `micromamba activate new-stack`. The `super` environment that + `submit_all_evaluations.py` and `parameter_search.py` default to does not exist on this host. - Partition `grete:preemptible` (2-day limit): GRES `1g.10gb:1` (plentiful), `1g.20gb:1` (8 slices), `2g.20gb:1` (16 slices), `3g.40gb:1` (8). `grete:interactive` allows two jobs per user for 12 h. Every job needs `--constraint=inet`. Account `nim00007`; QOS `2h` and `normal` only. @@ -280,5 +281,4 @@ out to be one-pixel boundary conventions on small objects (see the closing secti of the closed campaigns; the output-root trees they name stay as data. - 2026-09-06: the production submitter defaults were fixed (`submit_all_evaluations.py`: environment `new-stack`, 3D jobs on `1g.20gb:1`; `parameter_search.py` array scripts activate `new-stack`), so the overrides this note - describes for the `super` environment were no longer needed at that point. -- The current evaluation and optimization launchers default to `super`, as requested by the user. + describes for the `super` environment are no longer needed. diff --git a/finetuning/v2/evaluation/optimization/notes/EXPERIMENTAL_SETUP.md b/finetuning/v2/evaluation/optimization/notes/EXPERIMENTAL_SETUP.md index 1970fcdac..a32f9af15 100644 --- a/finetuning/v2/evaluation/optimization/notes/EXPERIMENTAL_SETUP.md +++ b/finetuning/v2/evaluation/optimization/notes/EXPERIMENTAL_SETUP.md @@ -29,8 +29,9 @@ the refinement statistics columns, and the configuration files under `optimizati ## 2. Environment and cluster -- Environment: `micromamba activate super`. The launchers `submit_all_evaluations.py`, `parameter_search.py`, - and `submit_optimization_jobs.py` activate it by default. +- Environment: `micromamba activate new-stack`. Both submitters (`submit_all_evaluations.py`, + `parameter_search.py`) activate it by default since 2026-09-06; the earlier default `super` does not exist + on grete. - Partition `grete:preemptible` (2-day limit). GRES pools: `1g.10gb:1` (plentiful), `1g.20gb:1` (8 slices), `2g.20gb:1` (16 slices), `3g.40gb:1` (8). `grete:interactive` allows two jobs per user for 12 h. Every job needs `--constraint=inet`. Account `nim00007`; QOS `2h` and `normal` only. @@ -38,24 +39,32 @@ the refinement statistics columns, and the configuration files under `optimizati `set -u` fails on `/etc/bashrc`). - Presets (`submit_optimization_jobs.PRESETS`): - | preset | GRES | memory | time | QOS | CPUs | - |------------|-------------|--------|----------|------|------| - | `2d` | `1g.10gb:1` | 16G | 08:00:00 | | 4 | - | `2d-short` | `1g.10gb:1` | 16G | 02:00:00 | `2h` | 4 | - | `3d` | `2g.20gb:1` | 32G | 12:00:00 | | 4 | - | `3d-large` | `2g.20gb:1` | 64G | 12:00:00 | | 4 | - | `cpu` | `1g.10gb:1` | 64G | 04:00:00 | | 16 | + | preset | partition | GRES | memory | time | QOS | CPUs | + |------------|--------------------|-------------|--------|----------|------|------| + | `2d` | `grete:preemptible` | `1g.10gb:1` | 16G | 08:00:00 | | 4 | + | `2d-short` | `grete:preemptible` | `1g.10gb:1` | 16G | 02:00:00 | `2h` | 4 | + | `3d` | `grete:preemptible` | `2g.20gb:1` | 32G | 12:00:00 | | 4 | + | `3d-large` | `grete:preemptible` | `2g.20gb:1` | 64G | 12:00:00 | | 4 | + | `cpu-test` | `standard96s:test` | none | 500G | 00:59:00 | | 192 | + | `cpu-shared` | `standard96s:shared` | none | 16G | 01:00:00 | | 4 | + | `cpu` | `grete:preemptible` | `1g.10gb:1` | 64G | 04:00:00 | | 16 | + + `cpu-test` is the GPU-free, full-node preset for cached 2D sweeps. It packs 48 four-thread commands into + every 192-core allocation by default. `cpu-shared` is for individual cached screens. The legacy `cpu` preset + remains available for longer jobs that were already designed around the Grete partition. - The submitter writes `/jobs/_/` with `tasks.txt` (`tagcommand`), `job.sh`, `logs/`, `submit.json` (argv, resources, git revision, dirty flag) and `job_id.txt`, and submits `job.sh` as one array (`--array=0-N%throttle`, default throttle 8, `--requeue`, - `--open-mode=append`). Every task leaves `logs/.done` or `logs/.failed`; dependent + `--open-mode=append`). `--tasks-per-job` packs several task-file commands into one array element; it + defaults to 48 for `cpu-test` and one otherwise. Every task retains its own stdout, stderr and + `logs/.done` or `logs/.failed` marker; dependent stages wait on those markers or on `--dependency afterok:`, never on an output file. `status ` reports state, exit code, restarts and marker per task; `--resume-from ` re-submits the unfinished tasks; `--local` runs the same tasks sequentially on the session GPU. `MICRO_SAM2_JOINT_CHECKPOINT_ROOT` and `MICRO_SAM2_JOINT_EXPORT_ROOT` are pinned into `job.sh` (`PINNED_ENV_VARS`), so a job resolves the same checkpoints as the shell that submitted it. -- Production evaluations go through `submit_all_evaluations.py` (job arrays, one task per dataset and mode, 8 h, +- Production evaluations go through `submit_all_evaluations.py` (one job per dataset and mode, 8 h, `grete:preemptible`, `--constraint=inet`): 2D jobs `1g.10gb:1` / 16G, 3D jobs `1g.20gb:1` / 64G, both checkpoint variables pinned into the script; `--gpu`, `--memory`, `--env`, `--dry` override or inspect. - Always `--dry-run` first and read `job.sh`; `sbatch --test-only job.sh` checks the header. @@ -116,7 +125,7 @@ the refinement statistics columns, and the configuration files under `optimizati ## 5. 2D subsets (`optimization/benchmark_apg_optimization.py`) -Manifest schema version 5; files `/subset_manifest_v5{,_holdout,_training_extra,_deep3d}.json` +Manifest schema version 5; files `/subset_manifest_v5{,_holdout,_training_extra,_ood_extended,_deep3d}.json` (`_default_manifest_path`). Each manifest records its `manifest_checksum`, `selection_policy`, `schema_version` and `data_root`; `_validate_manifest` requires the exact schema version. @@ -125,6 +134,7 @@ Manifest schema version 5; files `/subset_manifest_v5{,_holdout,_training_ | primary | `SAMPLE_COUNTS_2D` | livecell 80 (10 per each of 8 `LIVECELL_TYPES`), tissuenet 40, dynamicnuclearnet 40, deepbacs 30, dic_hepg2 50 = 240 images, plus one 12-slice volume each of celegans_atlas, embedseg, gonuclear, cremi, snemi (245 samples) | `0f8fb67b3650a71f9f44b53037e89546` | | holdout | `SAMPLE_COUNTS_2D_HOLDOUT`, image-disjoint | 80 / 40 / 40 / 30 / 43 = 233 images plus the same 5 volumes (238 samples); deepbacs is reused verbatim (`HOLDOUT_REUSED_DATASETS`) because all 30 validation images are primary | `bf8f3c28befe1fb06d62309dc302d1c4` | | training_extra | `TRAINING_EXTRA_DATASETS`, `SAMPLE_COUNTS_2D_TRAINING_EXTRA` (caps) | yeaz 40, neurips_cellseg 40, deepseas 40, puma 26 (cap 40), covid_if 5, tnbc 6 (cap 20) = 157 images, no volumes | `cee6224d6a93cec5a54a5c522a0f7bf5` | +| ood_extended | sealed Dice-foreground decoder confirmation set; official test loaders, stratified where heterogeneous | Arvidsson 10, BitDepth NucSeg 70, CellBinDB 48, microbeSEG 2, VICAR 50 = 180 images, no volumes | `836f92a084b05f6fa5445f03355589d9` | | deep3d variant | `--crops-3d deep`, `CROP_SHAPE_3D_DEEP = (32, 512, 512)` | the 240 primary images with 32-slice volumes; SNEMI 30 slices overlap the production slab, so this is a regression instrument, not a tuning set | `f611a7125383e850798d0b5bf696f6f7` | - The eleven-dataset development corpus of the 2026-09 campaigns is primary + training_extra @@ -185,22 +195,7 @@ Constants: `DEEP_DEPTH = 32` (`MIN_REALIZED_DEPTH = 24` slices of annotation mak (HDF5 per crop) and `view_apg3d_cases.py` (napari). - Per sample: `parameter_search.compute_metrics` gives `msa` (`elf.evaluation.mean_segmentation_accuracy`) and, for `metric_mode="dense"`, `cremi`, `vi_split`, `vi_merge`, `adapted_rand`. 2D segmentations pass - through `drop_severed_objects` first, symmetric with the ground-truth filtering. Since 2026-09-14, - `compute_metrics` also logs `sbd` (symmetric best Dice) for `metric_mode="sparse"`. No ranking uses `sbd`. -- Production evaluation (`evaluate_automatic_segmentation.py`, `evaluate_automatic_baselines.py` and the - interactive scripts) scores through `common.run_dataset_evaluation`. Every metric is a mean over the samples. - - Instance segmentation: mSA, SA50, SA75, precision, recall and F1 (`micro_sam.v1.evaluation.run_evaluation`), - and `SBD`, the symmetric best Dice (`bioimage_py.evaluation.symmetric_best_dice_score`, background - ignored). SBD exists since 2026-09-14. Older result files have no `SBD` column. - - Dense EM: `cremi`, `vi_split`, `vi_merge` and `adapted_rand`, without SBD. - - Volumes also report the sums `unmatched` and `genuine_misses`. - - 3D test volumes are scored on the pinned crops of `eval_crops_3d.json` (`common.EVAL_CROPS_3D`), with one - sample per crop. The crops are 32 deep and at most 512 in plane. They tile the annotated bounding box - without overlap and stay clear of tuning data. Near-empty crops are left out. - - platynereis_nuclei is read inside the annotated block that training uses as its roi - (`common.PLATYNEREIS_NUCLEI_TEST_ROIS`). - - `submit_all_evaluations.py --per_sample` runs one array task per sample. A task that finds the rows of all - samples writes the dataset result (`common.evaluate_samples`). It refuses rows whose metric columns differ. + through `drop_severed_objects` first, symmetric with the ground-truth filtering. - 2D aggregation (`_summarize`): per-dataset mean and std, then the row `__dataset_balanced__` = the equal-weight mean of the per-dataset means. This is "balanced mSA". - 3D aggregation (`benchmark_apg_3d.summarize`): per-dataset mean with a 2000-sample bootstrap CI, @@ -288,14 +283,6 @@ Epochs of the 2026-09 campaigns: `aeb1aca09a5fff43d2b8bb8bacff2b06` (campaign st results unaffected). Historical run directories stay valid records under their own epochs; the 3D aggregate reads them through `sibling_run_dirs`. -Later epochs, recorded on 2026-09-15: - -- `6bfb3121744c127074739f6897a085c3` (commit `a5893c36`, 2026-09-13): job-array sweeps of a joint checkpoint. -- `86fd931ef4019032002d04cab12df118` (2026-09-15, uncommitted): pinned 3D evaluation crops in `eval_crops_3d.json`, - per-sample evaluation, the platynereis_nuclei test rois, the `covid_if_cells` channel layout, and APG overrides - from a JSON configuration. `compute_metrics` also logs `sbd` for sparse datasets. The ranking stays on mSA or - CREMI, and the tuning splits are unchanged. - ## 12. Output root layout ``` @@ -339,6 +326,20 @@ Historical trees written only by code that lives on `apg-optim-fable` (data, rea confirmation on holdout, one production run on the 23 (2D) or the test manifest (3D) at the very end, with the twelve strictly unseen 2D datasets as the out-of-domain check. +Status (2026-09-06): implemented as `optimization/benchmark_ais_optimization.py` (`predict` caches the +decoder predictions per manifest sample under `/ais/predictions/`, `run` / `screen` / `sweep` / +`oracle` / `report` work on the cache), task builder `optimization/ais_campaign_tasks.py`, configurations +`optimization/configs/ais_*.json`, decision log `notes/AIS_V4_OPTIMIZATION.md`. The AIS implementation +checksum covers five files (the benchmark, `common.py`, `parameter_search.py`, +`micro_sam/v2/{instance_segmentation, postprocessing}.py`). AIS epochs: `f57b117edfda5420d9df761b1db4db2d` +(frozen Phase 0 harness) → `5700c6e0f471b360013551a442b1e53d` (harness only: refined loss decomposition +columns) → `a65e2eb08c23538f11544860736961a3` (epoch A1, 2026-09-06: opt-in `boundary_magnitude_max` +filter in `micro_sam/v2/postprocessing.py`, default off; the cached sweep scorer applies it) → +`576a85c8ffd4314627812fd30a3c1223` (epoch A2, 2026-09-06: the optimized `hvit_t` defaults with volume overrides, the fast filter +and dimension-aware `default_postprocessing`) → `e9d02380e340edfaccd30bf5cbf1bf03` (epoch A3, 2026-09-07: volume defaults +reverted to the registry smoothing and size floor plus the filter after the 3D test manifest). Decision log +and results: `notes/AIS_V4_OPTIMIZATION.md`. + ## 14. Baseline results of the cleaned harness (2026-09-06) Reruns of the default settings with the joint/v4 hvit_t geodesic checkpoint (checksum `5a729846…`) on the @@ -415,3 +416,26 @@ AIS v4 is mixed (−2.0 % on average: dsb +9.0 %, gonuclear +9.4 %, livecell +1. deepbacs −30.1 %) and worsens humanneurons. APG beats AIS on every dataset except dynamicnuclearnet, as under v2. Note that deepbacs APG gained +28.5 % on its validation subset (section 14.1 vs the v2 control) but is flat on the test split. + +### 14.4 AIS with the optimized `hvit_t` defaults (2026-09-07, epoch A3 `e9d02380e340edfaccd30bf5cbf1bf03`) + +The AIS optimization campaign (`notes/AIS_V4_OPTIMIZATION.md`) promoted new `hvit_t` post-processing +defaults into `micro_sam/v2/postprocessing.py`: images `sigma 1.0, min_size 50, boundary_magnitude_max 0.4` +(the new instance filter that drops instances without a distance-magnitude dip along their boundary), volumes +`min_size 100, sigma 0.5` (the registry values) with the same filter; everything else unchanged, the dense +multicut untouched. The old values remain reachable as `optimization/configs/ais_control_v4_old_defaults.json`. + +Development / confirmation (cached predictions, `/ais/`): 2D eleven-dataset development corpus balanced +mSA 0.3357 → 0.3437 (+2.4 %, 9 up, worst −0.8 %), 2D holdout 0.2337 → 0.2437 (+4.3 %, 5 / 5 up), 3D LM crops +primary 0.1765 → 0.1847 (+4.7 %), holdout 0.1998 → 0.2193 (+9.7 %), 3D test manifest (opened once) 0.1083 → +0.1120 (+3.4 %, 6 / 6 up). `compare_apg_optimization.py --target quality` on the five primary datasets: macro ++3.9 % (primary) / +4.3 % (holdout), every dataset up, runtime within +2.9 % per dataset; the +5 % macro bar of +that gate is not reached, the generalization gate of section 9 is. + +Production 2D test splits (`experiments/v4_geodesic_ais_optimization/results/`, tags `old-defaults` vs +`a2-defaults`, `report_ais_production.py`): 21 of 23 datasets up, balanced 0.2735 → 0.2864 (+4.7 %); the twelve +strictly unseen datasets 0.2104 → 0.2191 (+4.2 %, 10 up). Regressions: microbeseg 0.1420 → 0.1258 (−11.4 %, +attributed by ablation to `sigma 1.0` alone) and arvidsson −0.8 %. Reference rows for the datasets of 14.3: +livecell 0.2660, deepbacs 0.2319, dsb 0.4862, dynamicnuclearnet 0.5509 (AIS old: 0.2575 / 0.2056 / 0.4631 / +0.5083). 3D LM production (tag `a3-defaults`): 9 of 10 datasets up, none down, balanced 0.1455 → 0.1507 +(+3.6 %; embedseg 0.4310, gonuclear 0.2873, plantseg 0.1469); the dense EM rows of 14.3 are unchanged. diff --git a/finetuning/v2/evaluation/optimization/package_apg3d_cases.py b/finetuning/v2/evaluation/optimization/package_apg3d_cases.py index 480dcc994..c3390784b 100644 --- a/finetuning/v2/evaluation/optimization/package_apg3d_cases.py +++ b/finetuning/v2/evaluation/optimization/package_apg3d_cases.py @@ -53,8 +53,6 @@ def load_run(checkpoint: str, config: str, subset: str) -> tuple: other implementation checksums, the current implementation winning when a crop was run under both. """ from benchmark_apg_3d import load_volume_config, run_dir, sibling_run_dirs - from common import checkpoint_checksum, get_joint_checkpoint - from optimization.apg3d_manifest import load_manifest campaign_root, checkpoint_root = CHECKPOINTS[checkpoint] if checkpoint_root is not None: @@ -62,11 +60,7 @@ def load_run(checkpoint: str, config: str, subset: str) -> tuple: else: os.environ.pop("MICRO_SAM2_JOINT_CHECKPOINT_ROOT", None) config_name, params_3d = load_volume_config(CONFIGS[config]) - manifest = load_manifest(subset, campaign_root) - checkpoint_id = checkpoint_checksum(get_joint_checkpoint("hvit_t", "best")) - path = run_dir( - campaign_root, subset, config_name, params_3d, checkpoint_id, manifest["manifest_checksum"], "trial-1", - ) + path = run_dir(campaign_root, subset, config_name, params_3d) rows: Dict[str, dict] = {} for sibling in sibling_run_dirs(path): for crop in sorted((sibling / "crops").glob("*.json")): diff --git a/finetuning/v2/evaluation/optimization/prepare_ais_reoptimization_polish.py b/finetuning/v2/evaluation/optimization/prepare_ais_reoptimization_polish.py new file mode 100644 index 000000000..d8eb1f02d --- /dev/null +++ b/finetuning/v2/evaluation/optimization/prepare_ais_reoptimization_polish.py @@ -0,0 +1,132 @@ +"""Generate the bounded second-stage AIS grid from a finished coarse Dice-foreground ranking. + +The coarse search captures interactions between seed, watershed and boundary-use parameters. This +script takes the three best rows of every mechanism family and varies one coordinate at a time. Its +output uses the explicit-candidate grid format understood by ``benchmark_ais_optimization.py sweep``. +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Dict, Iterable, List, Optional, Sequence + +import numpy as np +import pandas as pd + +import benchmark_ais_optimization as ais + + +def _parameter_row(row: pd.Series) -> Dict[str, object]: + params = {} + for key in ais.SPARSE_KEYS: + if key not in row: + continue + value = row[key] + if pd.isna(value) or (isinstance(value, str) and value.lower() == "none"): + continue + # Mixed optional columns (for example ridge and mask controls in the same ranking) are + # read by pandas as strings because they also contain the sentinel "none". Convert their + # numeric entries back here so the generated JSON cannot pass string thresholds to numpy. + if isinstance(value, str): + if key == "boundary_magnitude_max" and value.lower() == ais.EXPLICIT_OFF: + value = float("inf") + else: + try: + value = float(value) + except ValueError as error: + raise ValueError(f"Invalid value {value!r} for AIS parameter '{key}'.") from error + if isinstance(value, np.generic): + value = value.item() + if key in ("n_iter", "min_size"): + value = int(value) + elif key == "boundary_magnitude_max" and np.isinf(value): + value = ais.EXPLICIT_OFF + params[key] = value + return params + + +def _with_values(base: Dict[str, object], key: str, values: Iterable[object]) -> Iterable[Dict[str, object]]: + for value in values: + candidate = dict(base) + candidate[key] = value + yield candidate + + +def _edge_iterations(family: pd.DataFrame, margin: float) -> List[int]: + by_iteration = family.groupby("n_iter")["balanced"].max() + extra = [] + if 1600 in by_iteration and 1200 in by_iteration and by_iteration[1600] - by_iteration[1200] >= margin: + extra.append(2400) + if 400 in by_iteration and 800 in by_iteration and by_iteration[400] - by_iteration[800] >= margin: + extra.append(200) + return extra + + +def polish_combinations(ranking: pd.DataFrame, top_per_family: int = 3, edge_margin: float = 0.001) -> List[Dict]: + """Return unique one-coordinate refinements of the top coarse rows.""" + if "balanced" not in ranking: + raise ValueError("The ranking needs a 'balanced' column.") + if "mechanism_family" not in ranking: + ranking = ranking.copy() + ranking["mechanism_family"] = "single-grid" + + candidates = {} + for _, family in ranking.groupby("mechanism_family", sort=True, dropna=False): + family = family.sort_values("balanced", ascending=False) + extra_iterations = _edge_iterations(family, edge_margin) + for _, row in family.head(top_per_family).iterrows(): + base = _parameter_row(row) + variants = [base] + threshold = float(base["foreground_threshold"]) + variants.extend(_with_values( + base, "foreground_threshold", + sorted({round(max(0.25, min(0.65, threshold + offset)), 3) for offset in (-0.025, 0, 0.025)}), + )) + variants.extend(_with_values(base, "foreground_weight", (0.5, 0.625, 0.75, 0.875, 1.0))) + variants.extend(_with_values(base, "min_size", (0, 10, 25, 50, 75, 100))) + variants.extend(_with_values( + base, "boundary_magnitude_max", (ais.EXPLICIT_OFF, 0.25, 0.3, 0.35, 0.4, 0.5, 0.6), + )) + if extra_iterations: + variants.extend(_with_values(base, "n_iter", extra_iterations)) + if "contact_weight" in base: + variants.extend(_with_values(base, "contact_weight", (0.25, 0.5, 0.75, 1.0, 1.5, 2.0, 3.0))) + if "contact_mask_threshold" in base: + variants.extend(_with_values( + base, "contact_mask_threshold", (0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8), + )) + for candidate in variants: + identity = json.dumps(candidate, sort_keys=True, separators=(",", ":")) + candidates[identity] = candidate + return list(candidates.values()) + + +def main(argv: Optional[Sequence[str]] = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--ranking", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--top-per-family", type=int, default=3) + parser.add_argument("--edge-margin", type=float, default=0.001) + args = parser.parse_args(argv) + if args.top_per_family < 1: + parser.error("--top-per-family must be positive.") + ranking = pd.read_csv(args.ranking) + combinations = polish_combinations(ranking, args.top_per_family, args.edge_margin) + resolved = [ais.resolve_postprocessing({"sparse": combo}, "hvit_t")["sparse"] for combo in combinations] + cache_keys = ais.SWEEP_CACHE_KEYS["sparse"] + n_cache_groups = len({tuple(combo[key] for key in cache_keys) for combo in resolved}) + args.output.parent.mkdir(parents=True, exist_ok=True) + with open(args.output, "w") as f: + json.dump({"combinations": combinations}, f, indent=2, sort_keys=True) + f.write("\n") + print( + f"Wrote {len(combinations)} polish candidates in {n_cache_groups} flow-cache groups to {args.output}. " + f"Use no more than {n_cache_groups} sweep shards." + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/finetuning/v2/evaluation/optimization/report_ais_checkpoint_comparison.py b/finetuning/v2/evaluation/optimization/report_ais_checkpoint_comparison.py new file mode 100644 index 000000000..8f08a62d2 --- /dev/null +++ b/finetuning/v2/evaluation/optimization/report_ais_checkpoint_comparison.py @@ -0,0 +1,468 @@ +"""Paired, hierarchical comparison of Dice-foreground baseline and boundary AIS checkpoints. + +Each checkpoint must be evaluated with its own development-selected post-processing configuration on +the exact same manifest. The report balances acquisition strata within a dataset, balances datasets +in the macro score, and resamples datasets and paired images for its confidence interval. It also +audits the sealed OOD paths against the two decoder-training manifests. +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +from pathlib import Path +from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple + +import numpy as np +import pandas as pd + +OPTIMIZATION_ROOT = Path(__file__).resolve().parent +sys.path.insert(0, str(OPTIMIZATION_ROOT)) + +import benchmark_ais_optimization as ais # noqa +import benchmark_apg_optimization as apg # noqa + + +DEFAULT_TRAINING_ROOT = ( + Path("/mnt/vast-nhr/projects/cidas/cca/experiments/micro_sam2/apg_optimization") + / "ais_decoder_training/checkpoints" +) +DEFAULT_TRAINING_MANIFESTS = ( + DEFAULT_TRAINING_ROOT / "ais_decoder_baseline/data_manifest.json", + DEFAULT_TRAINING_ROOT / "ais_decoder_boundary/data_manifest.json", +) +FATE_COLUMNS = ( + "matched", "unmatched", "genuine_misses", "gt_with_0_seeds", "gt_with_1_seed", + "gt_with_2plus_seeds", "seeded_unmatched", "seeded_split", "seeded_merged", + "seeded_undersized", "seeded_oversized", "unseeded_absorbed", "unseeded_missing", +) +OTHER_COUNT_COLUMNS = ("predicted_objects", "n_seeds", "background_seeds", "pipeline_mismatch") +EXTENT_COLUMNS = ("fg_iou", "fg_area_ratio", "matched_iou") +TIME_COLUMNS = ("initialization_seconds", "generation_seconds", "total_seconds") + + +def _json_default(value: Any) -> Any: + if isinstance(value, Path): + return str(value) + if isinstance(value, np.generic): + return value.item() + raise TypeError(f"Cannot serialize {type(value).__name__} to JSON.") + + +def _atomic_json(path: Path, value: Dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_name(f"{path.name}.tmp.{os.getpid()}") + with open(temporary, "w") as f: + json.dump(value, f, indent=2, sort_keys=True, default=_json_default) + f.write("\n") + os.replace(temporary, path) + + +def _atomic_csv(path: Path, value: pd.DataFrame) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_name(f"{path.name}.tmp.{os.getpid()}") + value.to_csv(temporary, index=False) + os.replace(temporary, path) + + +def _single(values: Iterable[Any], description: str) -> Any: + unique = {json.dumps(value, sort_keys=True, default=_json_default): value for value in values} + if len(unique) != 1: + raise ValueError(f"Expected one {description}, found {len(unique)} distinct values.") + return next(iter(unique.values())) + + +def load_side(run_dirs: Sequence[Path], side: str) -> Tuple[Dict[str, Any], pd.DataFrame]: + """Load compatible complete runs belonging to one checkpoint/configuration.""" + if not run_dirs: + raise ValueError(f"No {side} run directories were supplied.") + loaded = [(path.resolve(strict=True), *ais.load_run(path.resolve(strict=True))) for path in run_dirs] + metadata = [entry[1] for entry in loaded] + for field in ( + "config_name", "config_checksum", "checkpoint_checksum", "implementation_checksum", "model_type", + "params_2d", "dimensions", + ): + _single((item.get(field) for item in metadata), f"{side} {field}") + if any(item.get("mode") not in ("sparse", "auto") for item in metadata): + raise ValueError(f"Every {side} run must use sparse AIS post-processing.") + if any(item.get("dimensions") != [2] for item in metadata): + raise ValueError(f"Every {side} run must be 2d-only (dimensions=[2]).") + if len({item.get("manifest_checksum") for item in metadata}) != len(metadata): + raise ValueError(f"The {side} inputs contain duplicate runs for a manifest.") + + samples = pd.concat([entry[2] for entry in loaded], ignore_index=True) + if samples["sample_id"].duplicated().any(): + duplicated = samples.loc[samples["sample_id"].duplicated(), "sample_id"].tolist() + raise ValueError(f"The {side} runs contain duplicate sample ids: {duplicated[:5]}.") + if not (samples["ndim"] == 2).all(): + raise ValueError(f"The {side} sample table contains non-2d rows.") + if "stratum" not in samples: + samples["stratum"] = "" + samples["stratum"] = samples["stratum"].fillna("").astype(str) + summary = { + "side": side, + "run_dirs": [str(entry[0]) for entry in loaded], + "config_name": _single((item["config_name"] for item in metadata), f"{side} config name"), + "config_checksum": _single((item["config_checksum"] for item in metadata), f"{side} config checksum"), + "checkpoint_checksum": _single( + (item["checkpoint_checksum"] for item in metadata), f"{side} checkpoint checksum", + ), + "checkpoint_name": _single((item.get("checkpoint_name") for item in metadata), f"{side} checkpoint name"), + "implementation_checksum": _single( + (item["implementation_checksum"] for item in metadata), f"{side} implementation checksum", + ), + "model_type": _single((item["model_type"] for item in metadata), f"{side} model type"), + "params_2d": _single((item["params_2d"] for item in metadata), f"{side} 2d parameters"), + "manifest_checksums": sorted(item["manifest_checksum"] for item in metadata), + "subsets": sorted(str(item.get("subset")) for item in metadata), + "hardware": _single((item.get("hardware", {}) for item in metadata), f"{side} prediction hardware"), + "postprocessing_hardware": _single( + (item.get("postprocessing_hardware", {}) for item in metadata), + f"{side} post-processing hardware", + ), + } + return summary, samples + + +def pair_samples(baseline: pd.DataFrame, boundary: pd.DataFrame) -> pd.DataFrame: + """Pair the exact same source samples, retaining metrics and diagnostics from both sides.""" + identity = ["sample_id", "dataset", "stratum"] + base_ids = set(map(tuple, baseline[identity].itertuples(index=False, name=None))) + boundary_ids = set(map(tuple, boundary[identity].itertuples(index=False, name=None))) + if base_ids != boundary_ids: + raise ValueError( + "Baseline and boundary runs do not contain the same samples: " + f"{len(base_ids - boundary_ids)} baseline-only, {len(boundary_ids - base_ids)} boundary-only." + ) + required = ["msa", *EXTENT_COLUMNS, *TIME_COLUMNS, "gt_objects", *FATE_COLUMNS, *OTHER_COUNT_COLUMNS] + missing = [column for column in required if column not in baseline or column not in boundary] + if missing: + raise ValueError( + f"Both runs must include full AIS diagnostics; missing shared columns: {missing}. " + "Do not use --no-diagnostics for the confirmation runs." + ) + paired = baseline[identity + required].merge( + boundary[identity + required], on=identity, how="inner", validate="one_to_one", + suffixes=("_baseline", "_boundary"), + ) + if paired[["msa_baseline", "msa_boundary"]].isna().any().any(): + raise ValueError("Paired mSA values must all be finite.") + if not np.array_equal(paired["gt_objects_baseline"], paired["gt_objects_boundary"]): + raise ValueError("The two runs disagree on ground-truth object counts.") + return paired.sort_values(identity).reset_index(drop=True) + + +def validate_manifest_coverage(paired: pd.DataFrame, manifest: Dict[str, Any]) -> None: + """Require the paired table to cover every sealed manifest sample, with matching domain metadata.""" + identity = ("sample_id", "dataset", "stratum") + expected = { + (sample["sample_id"], sample["dataset"], str(sample.get("stratum", ""))) + for sample in manifest["samples"] + } + actual = set(map(tuple, paired[list(identity)].itertuples(index=False, name=None))) + if actual != expected: + raise ValueError( + "The paired runs do not exactly cover the sealed manifest: " + f"{len(expected - actual)} missing and {len(actual - expected)} unexpected sample identities." + ) + + +def _strata(group: pd.DataFrame) -> List[pd.DataFrame]: + if (group["stratum"] != "").any(): + if (group["stratum"] == "").any(): + raise ValueError(f"Dataset '{group['dataset'].iloc[0]}' mixes declared and missing strata.") + return [part for _, part in group.groupby("stratum", sort=True)] + return [group] + + +def balanced_scores(group: pd.DataFrame) -> Tuple[float, float]: + """Return baseline/boundary mSA, giving declared strata equal weight.""" + scores = np.asarray([ + [part["msa_baseline"].mean(), part["msa_boundary"].mean()] for part in _strata(group) + ]) + return float(scores[:, 0].mean()), float(scores[:, 1].mean()) + + +def _balanced_column(group: pd.DataFrame, column: str) -> float: + values = np.asarray([part[column].dropna().mean() for part in _strata(group)], dtype="float64") + return float(values[np.isfinite(values)].mean()) if np.isfinite(values).any() else np.nan + + +def _stratified_bootstrap(group: pd.DataFrame, n_bootstrap: int, rng: np.random.Generator) -> np.ndarray: + scores = np.zeros((n_bootstrap, 2), dtype="float64") + strata = _strata(group) + for part in strata: + values = part[["msa_baseline", "msa_boundary"]].to_numpy(dtype="float64") + indices = rng.integers(0, len(values), size=(n_bootstrap, len(values))) + scores += values[indices].mean(axis=1) + return scores / len(strata) + + +def bootstrap( + paired: pd.DataFrame, n_bootstrap: int = 20_000, seed: int = 0, +) -> Tuple[Dict[str, float], Dict[str, Dict[str, float]]]: + """Paired hierarchical bootstrap over domains and images within each domain/stratum.""" + if n_bootstrap < 100: + raise ValueError("Use at least 100 bootstrap replicates.") + rng = np.random.default_rng(seed) + datasets = sorted(paired["dataset"].unique()) + domain_samples = np.stack([ + _stratified_bootstrap(paired[paired["dataset"] == dataset], n_bootstrap, rng) + for dataset in datasets + ], axis=1) + domain_indices = rng.integers(0, len(datasets), size=(n_bootstrap, len(datasets))) + rows = np.arange(n_bootstrap)[:, None] + macro = domain_samples[rows, domain_indices].mean(axis=1) + + def interval(values: np.ndarray) -> Tuple[float, float]: + values = values[np.isfinite(values)] + if not len(values): + return np.nan, np.nan + low, high = np.quantile(values, (0.025, 0.975)) + return float(low), float(high) + + delta = macro[:, 1] - macro[:, 0] + with np.errstate(divide="ignore", invalid="ignore"): + relative = macro[:, 1] / macro[:, 0] - 1.0 + delta_low, delta_high = interval(delta) + relative_low, relative_high = interval(relative) + overall = { + "absolute_ci_low": delta_low, + "absolute_ci_high": delta_high, + "relative_ci_low": relative_low, + "relative_ci_high": relative_high, + "probability_boundary_better": float((delta > 0).mean()), + } + by_dataset = {} + for index, dataset in enumerate(datasets): + sample = domain_samples[:, index] + domain_delta = sample[:, 1] - sample[:, 0] + with np.errstate(divide="ignore", invalid="ignore"): + domain_relative = sample[:, 1] / sample[:, 0] - 1.0 + low, high = interval(domain_delta) + rel_low, rel_high = interval(domain_relative) + by_dataset[dataset] = { + "absolute_ci_low": low, "absolute_ci_high": high, + "relative_ci_low": rel_low, "relative_ci_high": rel_high, + "probability_boundary_better": float((domain_delta > 0).mean()), + } + return overall, by_dataset + + +def dataset_table(paired: pd.DataFrame, intervals: Dict[str, Dict[str, float]]) -> pd.DataFrame: + rows = [] + count_columns = [column for column in (*FATE_COLUMNS, *OTHER_COUNT_COLUMNS) if f"{column}_baseline" in paired] + for dataset, group in paired.groupby("dataset", sort=True): + baseline, boundary = balanced_scores(group) + relative = boundary / baseline - 1.0 if baseline else np.nan + row: Dict[str, Any] = { + "dataset": dataset, + "n_samples": int(len(group)), + "n_strata": len(_strata(group)), + "baseline_msa": baseline, + "boundary_msa": boundary, + "absolute_delta": boundary - baseline, + "relative_gain": relative, + "improved": bool(boundary > baseline), + "material_loss": bool( + boundary - baseline < ais.GATE["max_absolute_loss"] + and relative < ais.GATE["max_relative_loss"] + ), + **intervals[dataset], + } + for column in EXTENT_COLUMNS: + base = _balanced_column(group, f"{column}_baseline") + candidate = _balanced_column(group, f"{column}_boundary") + row[f"baseline_{column}"] = base + row[f"boundary_{column}"] = candidate + row[f"delta_{column}"] = candidate - base + for column in TIME_COLUMNS: + row[f"baseline_{column}"] = float(group[f"{column}_baseline"].sum()) + row[f"boundary_{column}"] = float(group[f"{column}_boundary"].sum()) + gt_total = float(group.get("gt_objects_baseline", pd.Series(dtype=float)).sum()) + for column in count_columns: + base = float(group[f"{column}_baseline"].sum()) + candidate = float(group[f"{column}_boundary"].sum()) + row[f"baseline_{column}"] = base + row[f"boundary_{column}"] = candidate + row[f"delta_{column}"] = candidate - base + if column in FATE_COLUMNS and gt_total: + row[f"delta_{column}_per_gt"] = (candidate - base) / gt_total + rows.append(row) + return pd.DataFrame(rows) + + +def audit_training_disjointness( + manifest: Dict[str, Any], data_root: Path, training_manifest_paths: Sequence[Path], +) -> Dict[str, Any]: + """Audit both dataset identities and resolved source paths against decoder training manifests.""" + ood_datasets = {sample["dataset"] for sample in manifest["samples"]} + ood_paths = { + (data_root / sample[key]).resolve() + for sample in manifest["samples"] + for key in ("raw_path", "label_path") + } + training_datasets, training_paths = set(), set() + manifests = [] + for path in training_manifest_paths: + path = path.resolve(strict=True) + with open(path) as f: + record = json.load(f) + datasets = record.get("datasets") + if not isinstance(datasets, dict): + raise ValueError(f"Training manifest '{path}' has no dataset mapping.") + training_datasets.update(datasets) + for splits in datasets.values(): + if not isinstance(splits, dict): + continue + for paths in splits.values(): + if isinstance(paths, list): + training_paths.update( + Path(value).expanduser().resolve() for value in paths if isinstance(value, str) + ) + manifests.append({"path": str(path), "variant": record.get("variant"), "n_datasets": len(datasets)}) + variants = {record["variant"] for record in manifests} + if variants != {"baseline", "boundary"}: + raise ValueError( + "The disjointness audit needs the Dice-foreground baseline and boundary training manifests; " + f"found variants {sorted(variants)}." + ) + dataset_overlap = sorted( + dataset for dataset in ood_datasets + if any(name == dataset or name.startswith(f"{dataset}_") for name in training_datasets) + ) + path_overlap = sorted(map(str, ood_paths & training_paths)) + if dataset_overlap or path_overlap: + raise RuntimeError( + f"The OOD set overlaps decoder training: datasets={dataset_overlap}, paths={path_overlap[:5]}." + ) + return { + "passed": True, + "ood_datasets": sorted(ood_datasets), + "training_manifests": manifests, + "dataset_overlap": dataset_overlap, + "path_overlap": path_overlap, + } + + +def compare( + baseline_runs: Sequence[Path], boundary_runs: Sequence[Path], manifest_path: Path, + training_manifest_paths: Sequence[Path], n_bootstrap: int = 20_000, seed: int = 0, + expected_subset: str = "ood_extended", expected_baseline_config: str = "baseline-dice-optimum", + expected_boundary_config: str = "boundary-dice-optimum", +) -> Tuple[Dict[str, Any], pd.DataFrame, pd.DataFrame]: + baseline_meta, baseline = load_side(baseline_runs, "baseline") + boundary_meta, boundary = load_side(boundary_runs, "boundary") + if baseline_meta["checkpoint_checksum"] == boundary_meta["checkpoint_checksum"]: + raise ValueError("Baseline and boundary runs unexpectedly use the same checkpoint.") + for field in ("implementation_checksum", "model_type", "manifest_checksums", "subsets"): + if baseline_meta[field] != boundary_meta[field]: + raise ValueError(f"Baseline and boundary {field} differ.") + if baseline_meta["subsets"] != [expected_subset]: + raise ValueError( + f"Expected only subset '{expected_subset}', found {baseline_meta['subsets']}." + ) + expected_configs = {"baseline": expected_baseline_config, "boundary": expected_boundary_config} + actual_configs = { + "baseline": baseline_meta["config_name"], "boundary": boundary_meta["config_name"], + } + if actual_configs != expected_configs: + raise ValueError( + "The confirmation must compare each checkpoint at its own frozen optimum: " + f"expected {expected_configs}, found {actual_configs}." + ) + + with open(manifest_path.resolve(strict=True)) as f: + manifest = json.load(f) + data_root = Path(manifest["data_root"]).resolve(strict=True) + apg._validate_manifest(manifest, data_root, "standard", expected_subset) + if baseline_meta["manifest_checksums"] != [manifest["manifest_checksum"]]: + raise ValueError("Run metadata does not match the supplied manifest checksum.") + + disjointness = audit_training_disjointness(manifest, data_root, training_manifest_paths) + paired = pair_samples(baseline, boundary) + validate_manifest_coverage(paired, manifest) + overall_interval, dataset_intervals = bootstrap(paired, n_bootstrap, seed) + domains = dataset_table(paired, dataset_intervals) + baseline_macro = float(domains["baseline_msa"].mean()) + boundary_macro = float(domains["boundary_msa"].mean()) + relative_gain = boundary_macro / baseline_macro - 1.0 if baseline_macro else np.nan + checks = { + "ci_excludes_zero": bool(overall_interval["absolute_ci_low"] > 0), + "at_least_four_of_five_domains_improve": bool( + len(domains) == 5 and int(domains["improved"].sum()) >= 4 + ), + "no_material_domain_loss": bool(not domains["material_loss"].any()), + "relative_macro_gain_at_least_two_percent": bool(relative_gain >= ais.GATE["min_balanced_gain"]), + } + timing_comparable = ( + baseline_meta["hardware"] == boundary_meta["hardware"] + and baseline_meta["postprocessing_hardware"] == boundary_meta["postprocessing_hardware"] + ) + report = { + "comparison": ( + "Dice-foreground boundary checkpoint at its own optimum vs Dice-foreground baseline at its own optimum" + ), + "baseline": baseline_meta, + "boundary": boundary_meta, + "manifest": str(manifest_path.resolve()), + "manifest_checksum": manifest["manifest_checksum"], + "n_samples": int(len(paired)), + "n_domains": int(len(domains)), + "domain_weighting": "equal; acquisition strata are equal-weight within each declared domain", + "baseline_macro_msa": baseline_macro, + "boundary_macro_msa": boundary_macro, + "absolute_delta": boundary_macro - baseline_macro, + "relative_gain": relative_gain, + "bootstrap": {"replicates": n_bootstrap, "seed": seed, **overall_interval}, + "claim_checks": checks, + "strong_improvement_claim_supported": bool(all(checks.values())), + "timing_comparable": timing_comparable, + "timing_seconds": { + side: {column: float(paired[f"{column}_{side}"].sum()) for column in TIME_COLUMNS} + for side in ("baseline", "boundary") + }, + "disjointness_audit": disjointness, + "caveats": ["microbeSEG has only two official manual test images and is a stress-test domain."], + } + return report, domains, paired + + +def main(argv: Optional[Sequence[str]] = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--baseline-runs", type=Path, nargs="+", required=True) + parser.add_argument("--boundary-runs", type=Path, nargs="+", required=True) + parser.add_argument("--manifest", type=Path, required=True) + parser.add_argument("--training-manifests", type=Path, nargs="+", default=list(DEFAULT_TRAINING_MANIFESTS)) + parser.add_argument("--expected-subset", default="ood_extended") + parser.add_argument("--baseline-config-name", default="baseline-dice-optimum") + parser.add_argument("--boundary-config-name", default="boundary-dice-optimum") + parser.add_argument("--bootstrap", type=int, default=20_000) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--output", type=Path, required=True, help="JSON path; CSV tables are written beside it.") + args = parser.parse_args(argv) + report, domains, paired = compare( + args.baseline_runs, args.boundary_runs, args.manifest, args.training_manifests, + args.bootstrap, args.seed, args.expected_subset, args.baseline_config_name, args.boundary_config_name, + ) + _atomic_json(args.output, report) + _atomic_csv(args.output.with_name(f"{args.output.stem}_domains.csv"), domains) + _atomic_csv(args.output.with_name(f"{args.output.stem}_paired_samples.csv"), paired) + print(domains[[ + "dataset", "n_samples", "n_strata", "baseline_msa", "boundary_msa", "absolute_delta", + "relative_gain", "absolute_ci_low", "absolute_ci_high", "material_loss", + ]].to_string(index=False, float_format=lambda value: f"{value:.4f}")) + print( + f"\nMacro mSA: {report['baseline_macro_msa']:.4f} -> {report['boundary_macro_msa']:.4f} " + f"({100 * report['relative_gain']:+.2f}%); 95% paired hierarchical CI " + f"[{report['bootstrap']['absolute_ci_low']:+.4f}, {report['bootstrap']['absolute_ci_high']:+.4f}]." + ) + print(f"Strong improvement claim supported: {report['strong_improvement_claim_supported']}") + print(f"Report: {args.output}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/finetuning/v2/evaluation/optimization/report_ais_decoders.py b/finetuning/v2/evaluation/optimization/report_ais_decoders.py new file mode 100644 index 000000000..aad7735dd --- /dev/null +++ b/finetuning/v2/evaluation/optimization/report_ais_decoders.py @@ -0,0 +1,209 @@ +"""Compare the AIS decoder variants of the 2026-09 training campaign on the benchmark's cached runs. + +Reads the run directories of the staged checkpoints (`/ais/hvit_t//`), joins the +requested subsets, and reports per (variant, configuration): balanced mSA and the generalization gate against +the baseline model under the library defaults, the mechanism columns of the D2 decomposition as a share of +the ground-truth objects (merged, absorbed, unseeded, background seeds) and the extent figures +(matched IoU, foreground IoU, foreground area ratio). Reader only: not part of the implementation checksum. + + python report_ais_decoders.py --subsets primary training_extra --ndim 2 --output /ais/reports/decoders_dev + python report_ais_decoders.py --kind apg3d --subsets primary holdout --ndim 3 --configs current-defaults +""" + +import argparse +import json +import os +import sys +from pathlib import Path +from typing import Dict, List, Optional, Sequence, Tuple + +import numpy as np +import pandas as pd + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) +import benchmark_ais_optimization as ais # noqa: E402 +from common import checkpoint_checksum # noqa: E402 + +DEFAULT_ROOT = Path("/mnt/vast-nhr/projects/cidas/cca/experiments/micro_sam2/apg_optimization") +DEFAULT_STAGED = DEFAULT_ROOT / "ais_decoder_training" / "staged" / "joint_sam2_hvit_t_multi_gpu" +VARIANTS = ("baseline", "contact", "fgcal", "both") +MECHANISMS = ("seeded_merged", "unseeded_absorbed", "gt_with_0_seeds", "seeded_split", "background_seeds") +EXTENT = ("matched_iou", "fg_iou", "fg_area_ratio") + + +def find_runs( + root: Path, model_type: str, checkpoint_id: str, manifest_checksums: Sequence[str], config_names: Sequence[str], + dimensions: List[int], epoch: Optional[str], +) -> Dict[str, List[Tuple[Path, Dict]]]: + """Complete run directories of one checkpoint, keyed by configuration name; the newest epoch if not given.""" + runs: Dict[str, List[Tuple[Path, Dict]]] = {} + for metadata_path in sorted((root / ais.CAMPAIGN / model_type / checkpoint_id).glob("*/metadata.json")): + with open(metadata_path) as f: + metadata = json.load(f) + if metadata.get("status") != "complete" or metadata.get("manifest_checksum") not in manifest_checksums: + continue + if metadata.get("config_name") not in config_names: + continue + if not set(dimensions) <= set(metadata.get("dimensions", [])): + continue + if epoch is not None and metadata.get("implementation_checksum") != epoch: + continue + runs.setdefault(metadata["config_name"], []).append((metadata_path.parent, metadata)) + # One run per (config, manifest): keep the newest epoch / trial when several exist. + for name, entries in runs.items(): + by_manifest: Dict[str, Tuple[Path, Dict]] = {} + for run_dir, metadata in sorted(entries, key=lambda e: e[0].stat().st_mtime): + by_manifest[metadata["manifest_checksum"]] = (run_dir, metadata) + runs[name] = list(by_manifest.values()) + return runs + + +def load_samples(entries: Sequence[Tuple[Path, Dict]], ndim: int, datasets: Optional[Sequence[str]]) -> pd.DataFrame: + samples = pd.concat([ais.load_run(run_dir)[1] for run_dir, _ in entries], ignore_index=True) + samples = samples[samples["ndim"] == ndim] + if datasets: + samples = samples[samples["dataset"].isin(datasets)] + return samples.reset_index(drop=True) + + +def mechanism_table(samples: pd.DataFrame) -> pd.DataFrame: + """Per dataset: the mechanism counts as a share of the ground-truth objects and the extent means.""" + rows = [] + for dataset, group in samples.groupby("dataset", sort=True): + gt = float(group["gt_objects"].sum()) + row = {"dataset": dataset, "gt_objects": int(gt), "msa": float(group["msa"].mean())} + for column in MECHANISMS: + row[column] = float(group[column].sum() / gt) if column in group and gt else float("nan") + for column in EXTENT: + row[column] = float(group[column].mean()) if column in group else float("nan") + rows.append(row) + return pd.DataFrame(rows) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--output-root", type=Path, default=DEFAULT_ROOT) + parser.add_argument("--staged-dir", type=Path, default=DEFAULT_STAGED, + help="Directory of the staged .pt files.") + parser.add_argument("--variants", nargs="+", default=list(VARIANTS)) + parser.add_argument("--baseline-variant", default="baseline") + parser.add_argument("--production-checkpoint", default=None, + help="A joint checkpoint (e.g. the v4 production one) to include as the variant 'production'.") + parser.add_argument("--configs", nargs="+", default=["current-defaults", "contact-ridge", "contact-mask"]) + parser.add_argument("--baseline-config", default="current-defaults") + parser.add_argument("--kind", default="v5", choices=ais.KINDS) + parser.add_argument("--subsets", nargs="+", default=["primary", "training_extra"]) + parser.add_argument("--ndim", type=int, default=2) + parser.add_argument("--datasets", nargs="*", default=None) + parser.add_argument("--epoch", default=None, help="Implementation checksum to select; default: any (newest).") + parser.add_argument("--model-type", default="hvit_t") + parser.add_argument("--data-root", type=Path, default=Path("/mnt/vast-nhr/projects/cidas/cca/data")) + parser.add_argument("--campaign-root", type=Path, default=DEFAULT_ROOT / "3d_v2") + parser.add_argument("--output", default=None, help="Prefix of the CSVs to write.") + args = parser.parse_args() + + manifests = [ + ais.load_campaign_manifest(args.kind, subset, args.output_root, args.data_root, args.campaign_root) + for subset in args.subsets + ] + checksums = [m["manifest_checksum"] for m in manifests] + dimensions = [args.ndim] + + checkpoints = {} + for variant in args.variants: + path = args.staged_dir / f"{variant}.pt" + if path.exists(): + checkpoints[variant] = checkpoint_checksum(str(path)) + else: + print(f"[skip] no staged checkpoint for '{variant}' at {path}") + if args.production_checkpoint: + checkpoints["production"] = checkpoint_checksum(args.production_checkpoint) + + tables: Dict[Tuple[str, str], pd.DataFrame] = {} + for variant, checkpoint_id in checkpoints.items(): + runs = find_runs( + args.output_root, args.model_type, checkpoint_id, checksums, args.configs, dimensions, args.epoch, + ) + for config_name, entries in runs.items(): + found = {m["manifest_checksum"] for _, m in entries} + if found != set(checksums): + print(f"[skip] {variant}/{config_name}: runs for {len(found)}/{len(checksums)} subsets only") + continue + tables[(variant, config_name)] = load_samples(entries, args.ndim, args.datasets) + print(f"[ok] {variant:10s} {config_name:18s} checkpoint {checkpoint_id[:8]} " + f"epoch {entries[0][1]['implementation_checksum'][:8]} n={len(tables[(variant, config_name)])}") + + reference = (args.baseline_variant, args.baseline_config) + if reference not in tables: + raise SystemExit(f"The reference {reference} has no complete runs; nothing to compare against.") + baseline_scores = ais.dataset_scores(tables[reference]) + baseline_mechanisms = mechanism_table(tables[reference]).set_index("dataset") + + summary_rows, detail_rows, mechanism_rows = [], [], [] + for (variant, config_name), samples in tables.items(): + scores = ais.dataset_scores(samples) + verdict = ais.gate_table(baseline_scores, scores) + mechanisms = mechanism_table(samples).set_index("dataset") + weights = mechanisms["gt_objects"] + row = { + "variant": variant, "config": config_name, "n_samples": int(len(samples)), + "balanced": verdict["balanced_candidate"], "balanced_gain": verdict["balanced_gain"], + "n_up": verdict["n_up"], "n_datasets": verdict["n_datasets"], "worst_relative": verdict["worst_relative"], + "passed": verdict["passed"], + } + for column in MECHANISMS: + # Object-weighted share over the datasets, and its change against the reference in percentage points. + share = float(np.nansum(mechanisms[column] * weights) / weights.sum()) + reference_share = float(np.nansum(baseline_mechanisms[column] * weights) / weights.sum()) + row[column] = share + row[f"{column}_delta"] = share - reference_share + for column in EXTENT: + row[column] = float(mechanisms[column].mean()) + summary_rows.append(row) + for dataset in verdict["datasets"]: + detail_rows.append({ + "variant": variant, "config": config_name, "dataset": dataset, + "baseline": float(baseline_scores[dataset]), "candidate": float(scores[dataset]), + "relative": verdict["relative"][dataset], + }) + for dataset, values in mechanisms.iterrows(): + mechanism_rows.append({"variant": variant, "config": config_name, "dataset": dataset, **values.to_dict()}) + + summary = pd.DataFrame(summary_rows).sort_values("balanced", ascending=False).reset_index(drop=True) + details = pd.DataFrame(detail_rows) + mechanisms_all = pd.DataFrame(mechanism_rows) + + pd.set_option("display.width", 250) + shown = summary.copy() + for column in ("balanced_gain", "worst_relative"): + shown[column] = shown[column].map(lambda v: "n/a" if v is None or not np.isfinite(v) else f"{100 * v:+.1f}%") + for column in MECHANISMS: + shown[column] = shown[column].map(lambda v: f"{100 * v:.1f}%") + shown[f"{column}_delta"] = shown[f"{column}_delta"].map(lambda v: f"{100 * v:+.1f}") + print("\nSummary (mechanisms as % of ground-truth objects, deltas in percentage points vs the reference):") + print(shown.to_string(index=False, float_format=lambda v: f"{v:.4f}")) + pivot = details.pivot_table(index=["variant", "config"], columns="dataset", values="relative") + print("\nRelative mSA change per dataset vs the reference:") + print(pivot.to_string(float_format=lambda v: f"{100 * v:+.1f}%")) + absolute = details.pivot_table(index=["variant", "config"], columns="dataset", values="candidate") + print("\nmSA per dataset:") + print(absolute.to_string(float_format=lambda v: f"{v:.4f}")) + extent = mechanisms_all.pivot_table(index=["variant", "config"], columns="dataset", values="fg_area_ratio") + print("\nForeground area ratio (fg > threshold / ground truth) per dataset:") + print(extent.to_string(float_format=lambda v: f"{v:.2f}")) + lost = mechanisms_all.assign(lost=mechanisms_all["seeded_merged"] + mechanisms_all["unseeded_absorbed"]) + merged = lost.pivot_table(index=["variant", "config"], columns="dataset", values="lost") + print("\nMerged + absorbed objects (% of ground truth) per dataset:") + print(merged.to_string(float_format=lambda v: f"{100 * v:.1f}%")) + + if args.output: + os.makedirs(os.path.dirname(os.path.abspath(args.output)), exist_ok=True) + summary.to_csv(f"{args.output}.csv", index=False) + details.to_csv(f"{args.output}_datasets.csv", index=False) + mechanisms_all.to_csv(f"{args.output}_mechanisms.csv", index=False) + print(f"\nwritten {args.output}.csv, _datasets.csv, _mechanisms.csv") + + +if __name__ == "__main__": + main() diff --git a/finetuning/v2/evaluation/optimization/report_ais_production.py b/finetuning/v2/evaluation/optimization/report_ais_production.py new file mode 100644 index 000000000..090468eb1 --- /dev/null +++ b/finetuning/v2/evaluation/optimization/report_ais_production.py @@ -0,0 +1,140 @@ +"""Compare two tagged AIS production evaluations dataset by dataset. + +Reads the result files `evaluate_automatic_segmentation.py` writes +(`/results/_micro_sam2__ais__ckpt-.csv`) for a baseline tag +and a candidate tag, and reports the metric per dataset (mSA, or the CREMI score for the dense EM +datasets), the relative change, the balanced means over the 2d datasets, over the twelve 2d datasets no +tuning ever saw (EXPERIMENTAL_SETUP.md, section 3) and over the 3d datasets, and the generalization gate. + +Usage: + python report_ais_production.py -e --baseline default_old-defaults \\ + --candidate default_a2-defaults +""" + +from __future__ import annotations + +import argparse +import re +import sys +from pathlib import Path +from typing import Dict, Optional, Sequence + +import numpy as np +import pandas as pd + +EVALUATION_ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(EVALUATION_ROOT)) + +from common import DATASETS_2D, DATASETS_3D_EM, DATASETS_3D_LM, DATASETS_DENSE, VAL_SPLITS # noqa +from optimization.benchmark_ais_optimization import GATE # noqa + +# The 2d development corpus of the 2026-09 campaigns; every other 2d dataset is strictly unseen by any tuning. +DEVELOPMENT_2D = ( + "livecell", "tissuenet", "dynamicnuclearnet", "deepbacs", "dic_hepg2", + "yeaz", "neurips_cellseg", "deepseas", "puma", "covid_if", "tnbc", +) +UNSEEN_2D = tuple(d for d in DATASETS_2D if d not in DEVELOPMENT_2D) + + +def read_results(experiment: Path, model_type: str, tag: str, checksum: Optional[str]) -> Dict[str, pd.Series]: + """The result row of every dataset with the given tag, keyed by dataset.""" + results = {} + pattern = re.compile( + rf"^(?P.+)_micro_sam2_{re.escape(model_type)}_ais_{re.escape(tag)}_ckpt-(?P[0-9a-f]+)\.csv$" + ) + for path in sorted((experiment / "results").glob("*.csv")): + match = pattern.match(path.name) + if match is None or (checksum is not None and not match.group("ck").startswith(checksum)): + continue + table = pd.read_csv(path) + if len(table) != 1: + raise ValueError(f"Expected one row in '{path}', got {len(table)}.") + results[match.group("dataset")] = table.iloc[0] + return results + + +def score(row: pd.Series, dataset: str) -> float: + """mSA, or the negated CREMI score on the dense EM datasets (higher is better either way).""" + if dataset in DATASETS_DENSE: + return -float(row["cremi"]) if "cremi" in row else float("nan") + return float(row["mSA"]) + + +def compare(baseline: Dict[str, pd.Series], candidate: Dict[str, pd.Series]) -> pd.DataFrame: + rows = [] + for dataset in sorted(set(baseline) & set(candidate)): + base, cand = score(baseline[dataset], dataset), score(candidate[dataset], dataset) + rows.append({ + "dataset": dataset, + "group": "2d" if dataset in DATASETS_2D else ("3d_lm" if dataset in DATASETS_3D_LM else "3d_em"), + "unseen_2d": dataset in UNSEEN_2D, + "has_val_split": dataset in VAL_SPLITS, + "metric": "-cremi" if dataset in DATASETS_DENSE else "msa", + "baseline": base, "candidate": cand, + "relative": cand / base - 1.0 if base else np.nan, "absolute": cand - base, + }) + return pd.DataFrame(rows) + + +def gate(table: pd.DataFrame) -> Dict[str, object]: + if table.empty: + return {"n": 0} + relative, absolute = table["relative"].to_numpy(), table["absolute"].to_numpy() + up = int((absolute > 0).sum()) + violates = (relative < GATE["max_relative_loss"]) & (absolute < GATE["max_absolute_loss"]) + balanced_gain = float(table["candidate"].mean() / table["baseline"].mean() - 1.0) + return { + "n": int(len(table)), "n_up": up, "balanced_baseline": float(table["baseline"].mean()), + "balanced_candidate": float(table["candidate"].mean()), "balanced_gain": balanced_gain, + "worst_relative": float(np.nanmin(relative)), "loss_limit_ok": bool(not violates.any()), + "passed": bool( + up >= len(table) - GATE["max_down"] and not violates.any() and balanced_gain >= GATE["min_balanced_gain"] + ), + } + + +def main(argv: Optional[Sequence[str]] = None) -> int: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("-e", "--experiment_folder", type=Path, required=True) + parser.add_argument("-m", "--model_type", default="hvit_t") + parser.add_argument("--baseline", required=True, help="Result tag of the baseline, e.g. default_old-defaults.") + parser.add_argument("--candidate", required=True, help="Result tag of the candidate, e.g. default_a2-defaults.") + parser.add_argument("--checksum", default=None, help="Checkpoint checksum prefix the result files must carry.") + parser.add_argument("--output", type=Path, default=None) + args = parser.parse_args(argv) + + baseline = read_results(args.experiment_folder, args.model_type, args.baseline, args.checksum) + candidate = read_results(args.experiment_folder, args.model_type, args.candidate, args.checksum) + table = compare(baseline, candidate) + missing = sorted((set(baseline) ^ set(candidate))) + if table.empty: + raise SystemExit(f"No dataset has both tags (baseline {len(baseline)}, candidate {len(candidate)} results).") + pd.set_option("display.width", 200) + shown = table.copy() + shown["relative"] = shown["relative"].map(lambda v: f"{100 * v:+.1f}%") + print(shown[["dataset", "group", "unseen_2d", "metric", "baseline", "candidate", "relative"]].to_string( + index=False, float_format=lambda v: f"{v:.4f}")) + for name, mask in ( + ("all 2d", table["group"] == "2d"), + ("2d strictly unseen (out of domain)", (table["group"] == "2d") & table["unseen_2d"]), + ("2d development", (table["group"] == "2d") & ~table["unseen_2d"]), + ("3d LM", table["group"] == "3d_lm"), + ("3d EM (dense, -CREMI)", table["group"] == "3d_em"), + ): + verdict = gate(table[mask]) + if verdict["n"]: + print(f"\n{name}: n {verdict['n']}, up {verdict['n_up']}, balanced {verdict['balanced_baseline']:.4f} -> " + f"{verdict['balanced_candidate']:.4f} ({100 * verdict['balanced_gain']:+.1f} %), worst " + f"{100 * verdict['worst_relative']:+.1f} %, loss limit ok {verdict['loss_limit_ok']}, " + f"gate {verdict['passed']}") + if missing: + print(f"\nDatasets with only one of the two tags so far: {missing}") + if args.output is not None: + args.output.parent.mkdir(parents=True, exist_ok=True) + table.to_csv(args.output, index=False) + print(f"\nTable: {args.output}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/finetuning/v2/evaluation/optimization/report_ais_sweep.py b/finetuning/v2/evaluation/optimization/report_ais_sweep.py new file mode 100644 index 000000000..a4334c547 --- /dev/null +++ b/finetuning/v2/evaluation/optimization/report_ais_sweep.py @@ -0,0 +1,252 @@ +"""Rank the combinations of an AIS parameter sweep as shared defaults under the generalization gate. + +Reads the per-dataset CSVs that `benchmark_ais_optimization.py sweep` wrote for one grid on one or more +manifests (e.g. primary and training_extra), joins them over the datasets, and reports for every +combination the balanced mSA, the per-dataset change against a reference combination (the current +library defaults by default), the generalization gate verdict and the mean ratio to each dataset's own +optimum. Not part of the implementation checksum: it only reads results. + +Usage: + python report_ais_sweep.py --grid configs/ais_grid_lm_v4.json --subset primary training_extra \\ + --output /ais/reports/a1_sweep_dev.csv --top 25 +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path +from typing import Dict, List, Optional, Sequence + +import numpy as np +import pandas as pd + +OPTIMIZATION_ROOT = Path(__file__).resolve().parent +sys.path.insert(0, str(OPTIMIZATION_ROOT.parent)) + +from optimization import benchmark_ais_optimization as ais # noqa + + +def load_sweep_tables( + grid_path: Path, subsets: Sequence[str], output_root: Path, data_root: Path, campaign_root: Path, + model_type: str, joint_checkpoint: str, datasets: Optional[Sequence[str]] = None, kind: str = "v5", +) -> Dict[str, pd.DataFrame]: + """The per-dataset sweep tables of one grid over the given manifests, keyed by dataset.""" + with open(grid_path) as f: + grid = json.load(f) + checkpoint_id = ais._checkpoint_identity(model_type, joint_checkpoint) + tables: Dict[str, pd.DataFrame] = {} + for subset in subsets: + manifest = ais.load_campaign_manifest(kind, subset, output_root, data_root, campaign_root) + sweep_dir = ais.sweep_dir(output_root, checkpoint_id, manifest["manifest_checksum"], grid_path.stem, grid) + for path in sorted(sweep_dir.glob("*.csv")): + if ".shard" in path.name or path.stem in ("shared_config",): + continue + if datasets and path.stem not in datasets: + continue + if path.stem in tables: + raise ValueError( + f"Dataset '{path.stem}' appears in more than one requested subset for grid '{grid_path.stem}'." + ) + tables[path.stem] = pd.read_csv(path) + if not tables: + sweeps = output_root / ais.CAMPAIGN / "sweeps" + raise FileNotFoundError(f"No sweep tables for grid '{grid_path.stem}' under {sweeps}.") + if datasets: + missing = sorted(set(datasets) - set(tables)) + if missing: + raise FileNotFoundError(f"Grid '{grid_path.stem}' is missing requested datasets: {missing}.") + return tables + + +def _parameter_columns(table: pd.DataFrame) -> List[str]: + return [c for c in table.columns if not c.endswith(("_mean", "_std")) and c != "n_images"] + + +def load_sweep_tables_many( + grid_paths: Sequence[Path], subsets: Sequence[str], output_root: Path, data_root: Path, campaign_root: Path, + model_type: str, joint_checkpoint: str, datasets: Optional[Sequence[str]] = None, kind: str = "v5", +) -> Dict[str, pd.DataFrame]: + """Load and union compatible sweep families, preserving the grid name as a categorical parameter.""" + collected: Dict[str, List[pd.DataFrame]] = {} + parameter_columns = set() + for grid_path in grid_paths: + tables = load_sweep_tables( + grid_path, subsets, output_root, data_root, campaign_root, model_type, joint_checkpoint, datasets, kind, + ) + for dataset, table in tables.items(): + table = table.copy() + if "mechanism_family" not in table: + table["mechanism_family"] = grid_path.stem + parameter_columns.update(_parameter_columns(table)) + collected.setdefault(dataset, []).append(table) + if not collected: + raise FileNotFoundError("No sweep tables were found for the requested grids.") + + combined = {} + for dataset, parts in collected.items(): + normalized = [] + for table in parts: + table = table.copy() + for column in parameter_columns: + if column not in table: + table[column] = "none" + normalized.append(table) + combined[dataset] = pd.concat( + normalized, ignore_index=True, sort=False, + ).drop_duplicates().reset_index(drop=True) + return combined + + +def rank_shared(tables: Dict[str, pd.DataFrame], reference: Optional[Dict[str, object]] = None) -> pd.DataFrame: + """Join the datasets on the parameter columns and score every combination as a shared default.""" + datasets = sorted(tables) + keys = _parameter_columns(tables[datasets[0]]) + merged = None + for dataset in datasets: + table = tables[dataset][keys + ["msa_mean"]].rename(columns={"msa_mean": dataset}).copy() + # NaN-safe join key for the optional parameters. + for key in keys: + table[key] = table[key].astype(object).where(table[key].notna(), "none") + if table.duplicated(keys).any(): + raise ValueError(f"Dataset '{dataset}' contains duplicate resolved parameter combinations.") + merged = table if merged is None else merged.merge(table, on=keys, how="inner") + if merged is None or merged.empty: + raise ValueError("The datasets share no combination.") + scores = merged[datasets].to_numpy(dtype="float64") + merged["balanced"] = scores.mean(axis=1) + optimum = scores.max(axis=0) + merged["mean_relative_optimum"] = (scores / optimum).mean(axis=1) + merged["min_relative_optimum"] = (scores / optimum).min(axis=1) + if reference is not None: + mask = np.ones(len(merged), dtype=bool) + for key, value in reference.items(): + if key not in keys: + continue + column = merged[key] + wanted = "none" if value is None else value + mask &= np.array([_same(v, wanted) for v in column]) + if mask.sum() != 1: + raise ValueError(f"The reference combination matches {int(mask.sum())} rows, expected one: {reference}.") + base = scores[mask][0] + relative = scores / np.where(base > 0, base, np.nan) - 1.0 + absolute = scores - base + merged["balanced_gain"] = merged["balanced"] / base.mean() - 1.0 + merged["n_up"] = (absolute > 0).sum(axis=1) + merged["worst_relative"] = np.nanmin(relative, axis=1) + violates = (relative < ais.GATE["max_relative_loss"]) & (absolute < ais.GATE["max_absolute_loss"]) + merged["passed"] = ( + (merged["n_up"] >= len(datasets) - ais.GATE["max_down"]) & ~violates.any(axis=1) + & (merged["balanced_gain"] >= ais.GATE["min_balanced_gain"]) + ) + for index, dataset in enumerate(datasets): + merged[f"rel_{dataset}"] = relative[:, index] + return merged + + +def select_plateau(ranked: pd.DataFrame, tolerance: float = 0.001) -> pd.Series: + """Select the robust, cheaper member of the near-optimal balanced-mSA plateau.""" + if ranked.empty: + raise ValueError("Cannot select from an empty sweep ranking.") + best = float(ranked["balanced"].max()) + plateau = ranked[ranked["balanced"] >= best - tolerance].copy() + + def numeric_column(name: str) -> pd.Series: + values = plateau[name] if name in plateau else pd.Series(0, index=plateau.index) + return pd.to_numeric(values, errors="coerce").fillna(0) + + plateau["_n_iter"] = numeric_column("n_iter") + contact_weight = numeric_column("contact_weight") + contact_mask = plateau.get("contact_mask_threshold", pd.Series("none", index=plateau.index)) + plateau["_active_controls"] = ( + (contact_weight != 0).astype(int) + (~contact_mask.isin(("none", None))).astype(int) + ) + plateau["_contact_weight"] = contact_weight + return plateau.sort_values( + ["min_relative_optimum", "_n_iter", "_active_controls", "_contact_weight", "balanced"], + ascending=[False, True, True, True, False], + ).iloc[0] + + +def selected_config(row: pd.Series, name: str) -> Dict[str, object]: + """Turn one ranked sweep row into a run-compatible sparse configuration.""" + params = {} + for key in ais.SPARSE_KEYS: + if key not in row or row[key] == "none" or pd.isna(row[key]): + continue + value = row[key] + if isinstance(value, np.generic): + value = value.item() + if key in ("n_iter", "min_size"): + value = int(value) + elif key == "boundary_magnitude_max" and np.isinf(value): + value = ais.EXPLICIT_OFF + params[key] = value + return {"name": name, "mode": "sparse", "params_2d": {"sparse": params}} + + +def _same(a: object, b: object) -> bool: + try: + return bool(np.isclose(float(a), float(b))) + except (TypeError, ValueError): + return str(a) == str(b) + + +def main(argv: Optional[Sequence[str]] = None) -> int: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--grid", type=Path, nargs="+", required=True) + parser.add_argument("--kind", choices=ais.KINDS, default="v5", help="Manifest family the sweep ran on.") + parser.add_argument("--subset", nargs="+", default=["primary", "training_extra"]) + parser.add_argument("--datasets", nargs="*", default=None) + parser.add_argument("--data-root", type=Path, default=ais.DEFAULT_DATA_ROOT) + parser.add_argument("--output-root", type=Path, default=ais.DEFAULT_OUTPUT_ROOT) + parser.add_argument("--campaign-root", type=Path, default=ais.apg3d_manifest.CAMPAIGN_ROOT) + parser.add_argument("--model-type", default="hvit_t") + parser.add_argument("--joint-checkpoint", default="best") + parser.add_argument("--no-reference", action="store_true", help="Do not compare against the library defaults.") + parser.add_argument("--sort", choices=("balanced", "mean_relative_optimum", "balanced_gain"), default="balanced") + parser.add_argument("--top", type=int, default=25) + parser.add_argument("--output", type=Path, default=None) + parser.add_argument("--select-config", type=Path, default=None, + help="Write the plateau-selected row as a run-compatible configuration.") + parser.add_argument("--config-name", default="dice-reoptimized") + parser.add_argument("--plateau-tolerance", type=float, default=0.001) + args = parser.parse_args(argv) + + tables = load_sweep_tables_many( + args.grid, args.subset, args.output_root.resolve(), args.data_root.resolve(), args.campaign_root, + args.model_type, args.joint_checkpoint, args.datasets, kind=args.kind, + ) + reference = None if args.no_reference else ais.default_postprocessing(args.model_type, "sparse") + ranked = rank_shared(tables, reference) + order = [args.sort] + (["passed"] if "passed" in ranked else []) + ranked = ranked.sort_values(order, ascending=False).reset_index(drop=True) + keys = _parameter_columns(tables[sorted(tables)[0]]) + shown = keys + ["balanced", "mean_relative_optimum", "min_relative_optimum"] + if "passed" in ranked: + shown += ["balanced_gain", "n_up", "worst_relative", "passed"] + print(f"{int(ranked['passed'].sum())} of {len(ranked)} combinations pass the gate against the defaults.") + pd.set_option("display.width", 250) + print(ranked[shown].head(args.top).to_string(index=False, float_format=lambda v: f"{v:.4f}")) + if "passed" in ranked and ranked["passed"].any(): + best = ranked[ranked["passed"]].sort_values("balanced", ascending=False).iloc[0] + print("\nBest passing combination:", {k: best[k] for k in keys}) + print("Per-dataset change:", {d: f"{100 * best[f'rel_{d}']:+.1f}%" for d in sorted(tables)}) + if args.output is not None: + args.output.parent.mkdir(parents=True, exist_ok=True) + ranked.to_csv(args.output, index=False) + print(f"\nRanking: {args.output}") + if args.select_config is not None: + selected = select_plateau(ranked, args.plateau_tolerance) + config = selected_config(selected, args.config_name) + args.select_config.parent.mkdir(parents=True, exist_ok=True) + with open(args.select_config, "w") as f: + json.dump(config, f, indent=2, sort_keys=True) + f.write("\n") + print(f"Selected configuration: {args.select_config}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/finetuning/v2/evaluation/optimization/submit_optimization_jobs.py b/finetuning/v2/evaluation/optimization/submit_optimization_jobs.py index 32d01b513..55ce435d5 100644 --- a/finetuning/v2/evaluation/optimization/submit_optimization_jobs.py +++ b/finetuning/v2/evaluation/optimization/submit_optimization_jobs.py @@ -1,7 +1,8 @@ """Submit the APG optimization scripts to Slurm as array jobs, or run them locally. Every task is one 'tagcommand' line of a tasks file. One array script dispatches the lines by -SLURM_ARRAY_TASK_ID, retries in-process failures, and records the outcome of every task as a +SLURM_ARRAY_TASK_ID, optionally packs several commands into a full-node allocation, retries failures, +and records the outcome of every task as a '.done' or '.failed' marker beside the logs, so a dependent stage can wait for a marker rather than for a file that may still be half written. Preemption restarts the script from the top through '--requeue'; the marker check and the scripts' own per-sample resume make that idempotent. @@ -11,9 +12,8 @@ python submit_optimization_jobs.py submit --name smoke --preset 2d --tasks-file tasks.txt --local python submit_optimization_jobs.py status --tail 3 -The presets encode the cluster facts of grete: '2d' runs on a 10 GB MIG slice, '3d' on a 20 GB one. -Canonical timing trials must share one hardware identity, so run them with '--throttle 1' on a -fixed GRES type. +The presets encode the cluster facts of Grete and standard96s. Canonical timing trials must share one +hardware identity, so run them with '--throttle 1' on a fixed GRES type. """ from __future__ import annotations @@ -35,19 +35,21 @@ # The benchmark's DEFAULT_OUTPUT_ROOT, duplicated so this module does not import torch. OUTPUT_ROOT = Path("/mnt/vast-nhr/projects/cidas/cca/experiments/micro_sam2/apg_optimization") JOBS_ROOT = OUTPUT_ROOT / "jobs" -ENV = "super" +ENV = "new-stack" PARTITION = "grete:preemptible" CONSTRAINT = "inet" N_ATTEMPTS = 3 RETRY_SLEEP_SECONDS = 30 DEFAULT_THROTTLE = 8 +DEFAULT_TASKS_PER_JOB = 1 +PRESET_TASKS_PER_JOB = {"cpu-test": 48} # Read by common.py at call time; pinned into the job script so a job resolves the same checkpoints. PINNED_ENV_VARS = ("MICRO_SAM2_JOINT_CHECKPOINT_ROOT", "MICRO_SAM2_JOINT_EXPORT_ROOT") @dataclasses.dataclass(frozen=True) class SlurmResources: - gres: str + gres: Optional[str] mem: str time_limit: str qos: Optional[str] = None @@ -61,6 +63,12 @@ class SlurmResources: "2d-short": SlurmResources("1g.10gb:1", "16G", "02:00:00", qos="2h"), "3d": SlurmResources("2g.20gb:1", "32G", "12:00:00"), "3d-large": SlurmResources("2g.20gb:1", "64G", "12:00:00"), + # Cached AIS/APG sweeps and screens never load the model. The test partition has a hard one-hour + # limit, which is sufficient for the cache-aware 2d shards and avoids reserving an idle MIG slice. + "cpu-test": SlurmResources(None, "500G", "00:59:00", cpus=192, partition="standard96s:test"), + # Small cached screens do not justify an exclusive test node. + "cpu-shared": SlurmResources(None, "16G", "01:00:00", cpus=4, partition="standard96s:shared"), + # Legacy long-running CPU preset on the GPU partition; retained for existing campaign commands. "cpu": SlurmResources("1g.10gb:1", "64G", "04:00:00", cpus=16), } @@ -120,23 +128,30 @@ def env_exports() -> str: def render_job_script( name: str, job_dir: Path, n_tasks: int, resources: SlurmResources, throttle: int = DEFAULT_THROTTLE, dependency: Optional[str] = None, attempts: int = N_ATTEMPTS, + tasks_per_job: int = DEFAULT_TASKS_PER_JOB, ) -> str: """Render the Slurm array script. Every '#SBATCH' line precedes the first command.""" + if tasks_per_job < 1: + raise ValueError("tasks_per_job must be positive.") + n_array_jobs = (n_tasks + tasks_per_job - 1) // tasks_per_job header = [ "#!/bin/bash", f"#SBATCH --job-name={sanitize(name)}", f"#SBATCH -p {resources.partition}", - f"#SBATCH -G {resources.gres}", + ] + if resources.gres is not None: + header.append(f"#SBATCH -G {resources.gres}") + header.extend([ f"#SBATCH -c {resources.cpus}", f"#SBATCH --mem={resources.mem}", f"#SBATCH -t {resources.time_limit}", f"#SBATCH --constraint={CONSTRAINT}", "#SBATCH --requeue", "#SBATCH --open-mode=append", - f"#SBATCH --array=0-{n_tasks - 1}%{throttle}", + f"#SBATCH --array=0-{n_array_jobs - 1}%{throttle}", f"#SBATCH -o {job_dir}/logs/{sanitize(name)}_%A_%a.out", f"#SBATCH -e {job_dir}/logs/{sanitize(name)}_%A_%a.err", - ] + ]) if resources.qos: header.append(f"#SBATCH --qos={resources.qos}") if resources.account: @@ -174,9 +189,67 @@ def render_job_script( printf 'exit=0 elapsed=%s attempts=%s job=%s restarts=%s node=%s\\n' "$elapsed" "$attempt" "$SLURM_JOB_ID" \\ "${{SLURM_RESTART_COUNT:-0}}" "$SLURMD_NODENAME" > "$markers/$tag.done" else - printf 'exit=%s elapsed=%s attempts=%s job=%s\\n' "$rc" "$elapsed" "$attempt" "$SLURM_JOB_ID" > "$markers/$tag.failed" + printf 'exit=%s elapsed=%s attempts=%s job=%s\\n' "$rc" "$elapsed" "$attempt" "$SLURM_JOB_ID" \\ + > "$markers/$tag.failed" fi exit $rc +""" + if tasks_per_job > 1: + body = f""" +set -eo pipefail +source ~/.bashrc +set -u +micromamba activate {ENV} +cd {REPOSITORY_ROOT} +export PYTHONUNBUFFERED=1 +{env_exports()} +markers={job_dir}/logs + +run_task() {{ + local task_index="$1" + local line tag command started rc attempt elapsed + line=$(sed -n "$((task_index + 1))p" {job_dir}/tasks.txt) + tag=$(cut -f1 <<< "$line") + command=$(cut -f2- <<< "$line") + echo "[$(date -Is)] task $task_index '$tag' array $SLURM_ARRAY_TASK_ID job $SLURM_JOB_ID" \\ + "restart ${{SLURM_RESTART_COUNT:-0}} node $SLURMD_NODENAME" + if [ -f "$markers/$tag.done" ]; then echo "'$tag' is already done."; return 0; fi + rm -f "$markers/$tag.failed" + started=$SECONDS + rc=1 + attempt=0 + for attempt in $(seq 1 {attempts}); do + rc=0 + eval "$command" >> "$markers/$tag.out" 2>> "$markers/$tag.err" || rc=$? + [ $rc -eq 0 ] && break + echo "[$(date -Is)] attempt $attempt of '$tag' failed with exit $rc." + sleep {RETRY_SLEEP_SECONDS} + done + elapsed=$((SECONDS - started)) + if [ $rc -eq 0 ]; then + printf 'exit=0 elapsed=%s attempts=%s job=%s restarts=%s node=%s\\n' \\ + "$elapsed" "$attempt" "$SLURM_JOB_ID" "${{SLURM_RESTART_COUNT:-0}}" "$SLURMD_NODENAME" \\ + > "$markers/$tag.done" + else + printf 'exit=%s elapsed=%s attempts=%s job=%s\\n' "$rc" "$elapsed" "$attempt" "$SLURM_JOB_ID" \\ + > "$markers/$tag.failed" + fi + return "$rc" +}} + +first_task=$((SLURM_ARRAY_TASK_ID * {tasks_per_job})) +last_task=$((first_task + {tasks_per_job})) +[ "$last_task" -gt {n_tasks} ] && last_task={n_tasks} +pids=() +for ((task_index=first_task; task_index Optional[str]: def write_job_dir( name: str, tasks: Sequence[Task], resources: SlurmResources, jobs_root: Path = JOBS_ROOT, throttle: int = DEFAULT_THROTTLE, dependency: Optional[str] = None, attempts: int = N_ATTEMPTS, - argv: Optional[Sequence[str]] = None, + argv: Optional[Sequence[str]] = None, tasks_per_job: int = DEFAULT_TASKS_PER_JOB, ) -> Path: """Create '/_/' with tasks.txt, job.sh, logs/ and submit.json.""" stamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") job_dir = jobs_root / f"{stamp}_{sanitize(name)}" (job_dir / "logs").mkdir(parents=True, exist_ok=False) write_tasks_file(job_dir, tasks) - script = render_job_script(name, job_dir, len(tasks), resources, throttle, dependency, attempts) + script = render_job_script( + name, job_dir, len(tasks), resources, throttle, dependency, attempts, tasks_per_job, + ) (job_dir / "job.sh").write_text(script) record = { "name": name, "argv": list(argv) if argv is not None else sys.argv, "resources": dataclasses.asdict(resources), "n_tasks": len(tasks), + "tasks_per_job": tasks_per_job, + "n_array_jobs": (len(tasks) + tasks_per_job - 1) // tasks_per_job, "throttle": throttle, "dependency": dependency, "attempts": attempts, @@ -283,14 +360,18 @@ def submit_tasks( tasks: Sequence[Task], name: str, resources: SlurmResources, throttle: int = DEFAULT_THROTTLE, dependency: Optional[str] = None, attempts: int = N_ATTEMPTS, dry_run: bool = False, local: bool = False, jobs_root: Path = JOBS_ROOT, resume_from: Optional[Path] = None, argv: Optional[Sequence[str]] = None, + tasks_per_job: int = DEFAULT_TASKS_PER_JOB, ) -> Tuple[Optional[Path], Optional[str]]: """The Python entry point the job builders call. Returns (job_dir, job_id).""" tasks = filter_resume(tasks, resume_from) if not tasks: print("Nothing to do.") return None, None - job_dir = write_job_dir(name, tasks, resources, jobs_root, throttle, dependency, attempts, argv) - print(f"Job directory: {job_dir} ({len(tasks)} tasks)") + job_dir = write_job_dir( + name, tasks, resources, jobs_root, throttle, dependency, attempts, argv, tasks_per_job, + ) + n_array_jobs = (len(tasks) + tasks_per_job - 1) // tasks_per_job + print(f"Job directory: {job_dir} ({len(tasks)} tasks packed into {n_array_jobs} array jobs)") if dry_run: print((job_dir / "job.sh").read_text()) return job_dir, None @@ -366,9 +447,11 @@ def status(job_dir: Path, tail: int = 1) -> int: for index in _expand_array_ids(row["JobID"]): states[index] = row failing = 0 - name = sanitize(json.loads((job_dir / "submit.json").read_text())["name"]) + submission = json.loads((job_dir / "submit.json").read_text()) + name = sanitize(submission["name"]) + tasks_per_job = int(submission.get("tasks_per_job", 1)) for index, (tag, _) in enumerate(tasks): - row = states.get(index, {}) + row = states.get(index // tasks_per_job, {}) marker, marker_text = _marker_state(job_dir, tag) slurm_state = row.get("State", "-") log = job_dir / "logs" / f"{tag}.out" @@ -403,6 +486,10 @@ def add_submit_arguments(parser: argparse.ArgumentParser) -> None: parser.add_argument("--throttle", type=int, default=DEFAULT_THROTTLE, help="Concurrent array tasks.") parser.add_argument("--dependency", default=None, help="Slurm dependency, e.g. afterok:123.") parser.add_argument("--attempts", type=int, default=N_ATTEMPTS, help="In-process retries per task.") + parser.add_argument( + "--tasks-per-job", type=int, default=None, + help="Concurrent task-file commands per Slurm array element (preset-dependent by default).", + ) parser.add_argument("--dry-run", action="store_true", help="Write the job directory, do not submit.") parser.add_argument("--local", action="store_true", help="Run the tasks here, sequentially.") parser.add_argument("--jobs-root", type=Path, default=JOBS_ROOT) @@ -421,10 +508,15 @@ def resolve_resources(args: argparse.Namespace) -> SlurmResources: def submit_from_args(tasks: Sequence[Task], args: argparse.Namespace) -> Tuple[Optional[Path], Optional[str]]: if not args.local and not args.dry_run: warn_missing_env() + tasks_per_job = args.tasks_per_job + if tasks_per_job is None: + tasks_per_job = PRESET_TASKS_PER_JOB.get(args.preset, DEFAULT_TASKS_PER_JOB) + if tasks_per_job < 1: + raise ValueError("--tasks-per-job must be positive.") return submit_tasks( tasks, args.name, resolve_resources(args), throttle=args.throttle, dependency=args.dependency, attempts=args.attempts, dry_run=args.dry_run, local=args.local, jobs_root=args.jobs_root, - resume_from=args.resume_from, + resume_from=args.resume_from, tasks_per_job=tasks_per_job, ) diff --git a/finetuning/v2/evaluation/parameter_search.py b/finetuning/v2/evaluation/parameter_search.py index 135f8a134..80f7e86aa 100644 --- a/finetuning/v2/evaluation/parameter_search.py +++ b/finetuning/v2/evaluation/parameter_search.py @@ -38,9 +38,9 @@ from bioimage_cpp.segmentation import label as connected_components, watershed -from bioimage_py.evaluation import symmetric_best_dice_score - -from micro_sam.v2.postprocessing import watershed_heightmap, _compute_flow_density +from micro_sam.v2.postprocessing import ( + drop_instances_without_boundary_dip, lower_height_under_seeds, watershed_heightmap, _compute_flow_density, +) from common import ( DATASETS_3D, DATASETS_DENSE, DATASET_SPACING, VAL_SPLITS, VAL_Z_RANGE, @@ -254,10 +254,18 @@ def score_image_sparse_cached( (None where a combo failed), aligned with params_list. """ foreground = prediction[0] - directed = prediction[1:] + directed = prediction[1:4] + contact = prediction[4] if prediction.shape[0] > 4 else None ndim = foreground.ndim if directed.shape[0] > ndim: directed = directed[-ndim:] + if contact is not None and contact.shape != foreground.shape: + raise ValueError(f"The contact map {contact.shape} must have the shape of the foreground {foreground.shape}.") + if contact is None and any( + params.get("contact_weight") is not None or params.get("contact_mask_threshold") is not None + for params in params_list + ): + raise ValueError("'contact_weight' and 'contact_mask_threshold' need prediction channel 4.") # The convergence densities and the height maps are built up front, so the scoring below only reads them. fg_mask_cache, density_cache, hmap_cache = {}, {}, {} @@ -272,18 +280,34 @@ def score_image_sparse_cached( n_threads=n_threads, ) fw = params["foreground_weight"] - if fw not in hmap_cache: - hmap_cache[fw] = watershed_heightmap(foreground, directed, fw) + contact_weight = params.get("contact_weight") + hmap_key = (fw, contact_weight) + if hmap_key not in hmap_cache: + hmap = watershed_heightmap(foreground, directed, fw) + if contact is not None and contact_weight is not None and contact_weight != 0: + hmap = np.ascontiguousarray( + hmap + np.float32(contact_weight) * np.clip(contact, 0, 1), dtype="float32", + ) + hmap_cache[hmap_key] = hmap # The base watershed does not depend on min_size, so all min_size values of a combo reuse it. base_cache, base_lock = {}, threading.Lock() - def base_segmentation(key, fg_mask, density, density_threshold, hmap): + def base_segmentation(key, fg_mask, density, density_threshold, hmap, seed_floor, contact_mask_threshold): with base_lock: cached = base_cache.get(key) if cached is None: seeds = connected_components(density > density_threshold) - cached = watershed(hmap, markers=seeds, mask=fg_mask) + hmap = lower_height_under_seeds(hmap, seeds, seed_floor) + if contact is not None and contact_mask_threshold is not None: + open_mask = fg_mask & ~(contact > contact_mask_threshold) + first = watershed( + hmap, markers=np.where(open_mask, seeds, 0).astype(seeds.dtype), mask=open_mask, + ) + segmentation = watershed(hmap, markers=first, mask=fg_mask) + else: + segmentation = watershed(hmap, markers=seeds, mask=fg_mask) + cached = (segmentation, hmap) with base_lock: base_cache[key] = cached return cached @@ -291,11 +315,19 @@ def base_segmentation(key, fg_mask, density, density_threshold, hmap): def score(params): ft, sigma, n_iter, dt = (params[k] for k in FLOW_DENSITY_KEYS) fw, density_threshold = params["foreground_weight"], params["density_threshold"] + contact_weight = params.get("contact_weight") + contact_mask_threshold = params.get("contact_mask_threshold") + seed_floor = params.get("seed_floor", "none") fg_mask = fg_mask_cache[ft] - hmap = hmap_cache[fw] try: - key = (ft, sigma, n_iter, dt, density_threshold, fw) - seg = base_segmentation(key, fg_mask, density_cache[(ft, sigma, n_iter, dt)], density_threshold, hmap) + key = ( + ft, sigma, n_iter, dt, density_threshold, fw, seed_floor, + contact_weight, contact_mask_threshold, + ) + seg, hmap = base_segmentation( + key, fg_mask, density_cache[(ft, sigma, n_iter, dt)], density_threshold, + hmap_cache[(fw, contact_weight)], seed_floor, contact_mask_threshold, + ) min_size = params["min_size"] if min_size > 0: seg = seg.copy() @@ -303,6 +335,9 @@ def score(params): discard = ids[(sizes < min_size) & (ids > 0)] seg[np.isin(seg, discard)] = 0 seg = watershed(hmap, markers=seg, mask=fg_mask) + max_median = params.get("boundary_magnitude_max") + if max_median is not None and np.isfinite(max_median): + seg = drop_instances_without_boundary_dip(seg, directed, max_median) return compute_metrics(seg.astype("uint32"), labels, "sparse", border_min_size) except Exception as e: warnings.warn(f"Sparse postprocessing failed for {params}: {e}") @@ -700,11 +735,7 @@ def tune_parameters( PARTITION = "grete:preemptible" # The micro-sam2 environment on grete; every array task activates it. -ENV = "super" - -# Which joint training version a task sweeps. Pinned into the array script, so a queued task sweeps -# the weights the submission chose rather than whatever the environment holds when it starts. -JOINT_ENV_VARS = ("MICRO_SAM2_JOINT_CHECKPOINT_ROOT", "MICRO_SAM2_JOINT_EXPORT_ROOT") +ENV = "new-stack" CPUS = 4 # A 2d task took 54 min at worst as a shard and 62 min unsharded, with the slow histopathology datasets # sharded (REGISTRY_2D_SHARDS). A longer limit only keeps the task out of the backfill window. @@ -947,7 +978,7 @@ def write_array_script(job_folder, name, tasks_path, n_tasks, gpu, memory, time_ source ~/.bashrc micromamba activate {ENV} -{env_exports()} + line=$(sed -n "$((SLURM_ARRAY_TASK_ID + 1))p" {tasks_path}) tag=$(cut -f1 <<< "$line") command=$(cut -f2- <<< "$line") diff --git a/finetuning/v2/evaluation/submit_all_evaluations.py b/finetuning/v2/evaluation/submit_all_evaluations.py index c9ac526ef..0ae0ef236 100644 --- a/finetuning/v2/evaluation/submit_all_evaluations.py +++ b/finetuning/v2/evaluation/submit_all_evaluations.py @@ -84,18 +84,11 @@ ("interactive", "microsam_vol"): {"ndim": (3,), "modality": ("lm",)}, } -# The data that one model of a method can run on, on top of METHOD_SUPPORT. The key is (method, model), since model -# names repeat across methods. The CellPose 3 generalists are not histopathology models. -MODEL_SUPPORT = {("cellpose", "cyto3"): {"modality": ("lm", "em")}, ("cellpose", "nuclei"): {"modality": ("lm", "em")}} - -# Use --env to override the method-specific environments. StarDist runs in its own because it needs -# TensorFlow, which does not belong next to torch in the main environment. -METHOD_ENV = {"stardist": "stardist"} - -# cyto3 and nuclei are CellPose 3 checkpoints, which the CellPose 4 of the main environment cannot load. -MODEL_ENV = {("cellpose", "cyto3"): "cellpose3", ("cellpose", "nuclei"): "cellpose3"} - -DEFAULT_ENV = "super" +# Methods whose packages do not live in the default environment. The names are per machine, so +# --env overrides them and a missing one is reported before anything is submitted. 'new-stack' is +# the micro-sam2 environment on grete; the earlier default 'super' does not exist there. +METHOD_ENV = {"cellpose": "cp3", "stardist": "sd"} +DEFAULT_ENV = "new-stack" # Slurm resources per job. Only the grete partitions are available. 'grete:preemptible' is usually # free and starts within minutes, where the shared pools queue for days. It is MIG only, so the GPU @@ -230,13 +223,10 @@ def build_command( command.append("--skip_tuning") if args.tuning_root is not None: command.extend(["--tuning_root", args.tuning_root]) - if args.apg_params is not None and mode == "apg": - command.extend(["--apg_params", args.apg_params]) - - if args.n_samples is not None: - command.extend(["--n_samples", str(args.n_samples)]) - if sample_index is not None: - command.extend(["--sample_index", str(sample_index)]) + if args.ais_params is not None and mode == "ais": + command.extend(["--ais_params", args.ais_params]) + if args.result_tag is not None: + command.extend(["--result_tag", args.result_tag]) if args.segmentation_type == "interactive": command.extend(["-p", args.prompt_choice, "-iter", str(args.n_iterations)]) @@ -340,10 +330,10 @@ def main(): help="Automatic only. Submit one array task per sample. The task that finds all rows writes the result.", ) parser.add_argument("--tuning_root", type=str, default=None, help="Where parameter_search.py wrote its sweeps.") - parser.add_argument( - "--apg_params", type=str, default=None, - help="A JSON configuration of APG parameters, passed to every micro-sam2 APG task.", - ) + parser.add_argument("--ais_params", type=str, default=None, + help="AIS benchmark configuration passed to every automatic AIS job (see " + "evaluate_automatic_segmentation.py --ais_params).") + parser.add_argument("--result_tag", type=str, default=None, help="Result tag passed to every automatic job.") parser.add_argument("-p", "--prompt_choice", type=str, default="box", choices=("box", "point")) parser.add_argument("-iter", "--n_iterations", type=int, default=8, help="Iterative prompting rounds.") parser.add_argument("--min_size", type=int, default=0, diff --git a/finetuning/v2/generalist/ais_decoder/__init__.py b/finetuning/v2/generalist/ais_decoder/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/finetuning/v2/generalist/ais_decoder/ais_decoder_lib.py b/finetuning/v2/generalist/ais_decoder/ais_decoder_lib.py new file mode 100644 index 000000000..10693e69a --- /dev/null +++ b/finetuning/v2/generalist/ais_decoder/ais_decoder_lib.py @@ -0,0 +1,437 @@ +"""Building blocks of the AIS decoder training campaign (2026-09). + +Decoder-only training of the UniSAM2 automatic branch: the SAM2 image encoder stays frozen at the weights of +the joint/v4 geodesic checkpoint, the UNETR decoder is warm-started from the same checkpoint and trained on +the train splits of the AIS tuning datasets. Everything the trainer pickles into its checkpoints (datasets, +wrappers, the model class) lives in this importable module, so that the staging step can re-open a +checkpoint from another process. + +Four variants: 'baseline' (the current four channel target and loss), 'contact' (a fifth output channel +trained on the touching boundaries), 'fgcal' (the foreground trained with Dice plus a boundary-weighted +cross entropy) and 'both'. + +Datasets: livecell, tissuenet, dynamicnuclearnet, deepbacs, dic_hepg2, neurips_cellseg, yeaz, puma, tnbc +(train splits). deepseas is excluded (binary masks: connected components merge touching cells, which would +corrupt the contact target and the geodesic field) and so is covid_if (no split; 44 of its 49 files are +production-scored). The trainer's validation set is a deterministic tail of every train file list, so the +evaluation manifests (val splits) and the test splits stay untouched. +""" + +import math +import os +from functools import partial +from glob import glob +from typing import Dict, List, Optional, Sequence, Tuple + +import numpy as np +import torch +import torch_em +from torch_em.data import ConcatDataset, MinInstanceSampler, datasets + +from micro_sam.v2.datasets.generalist_loader import _configure_training_normalization, _prepare_data_loader +from micro_sam.v2.datasets.wrapper import UniDataWrapper +from micro_sam.v2.models.util import UniSAM2 +from micro_sam.v2.transforms.labels import GeodesicHybridDistanceTransform +from micro_sam.v2.transforms.raw import _identity, _normalize_percentile, _to_8bit + +DATA_ROOT = "/mnt/vast-nhr/projects/cidas/cca/data" +CAMPAIGN_ROOT = "/mnt/vast-nhr/projects/cidas/cca/experiments/micro_sam2/apg_optimization/ais_decoder_training" +V4_CHECKPOINT = ( + "/mnt/vast-nhr/projects/cidas/cca/models/micro_sam2/joint/v4/checkpoints/" + "joint_sam2_hvit_t_geodesic_multi_gpu/best.pt" +) +MODEL_TYPE = "hvit_t" +INITIAL_FEATURES = 32 +PATCH_SHAPE = (512, 512) +VAL_FRACTION = 0.05 +MIN_VAL_FILES = 2 + +# The loss settings of the four variants. +VARIANTS: Dict[str, Dict] = { + "baseline": {"contact": False, "boundary_weight": None}, + "contact": {"contact": True, "boundary_weight": None}, + "fgcal": {"contact": False, "boundary_weight": 4.0}, + "both": {"contact": True, "boundary_weight": 4.0}, + # Second round (2026-09-07 evening): the fifth channel holds the full inner boundary of every object + # (contact_mode "all") instead of the touching boundaries only. + "boundary": {"contact": True, "contact_mode": "all", "boundary_weight": None}, + "boundary_fgcal": {"contact": True, "contact_mode": "all", "boundary_weight": 4.0}, +} +BOUNDARY_RADIUS = 2 + +# Samples per epoch (before scaling) and validation samples per dataset group. +TRAIN_SAMPLES = { + "livecell": 25, # per cell type, eight types + "tissuenet": 200, "dynamicnuclearnet": 200, "neurips_cellseg": 150, "dic_hepg2": 120, "deepbacs": 120, + "yeaz_bf": 80, "yeaz_phc": 20, "yeaz_phc_stacks": 20, "puma": 100, "tnbc": 60, +} +VAL_SAMPLES = {"livecell": 3, "yeaz_phc": 2, "yeaz_phc_stacks": 2, "tnbc": 2} +DEFAULT_VAL_SAMPLES = 20 + + +def n_output_channels(variant: str) -> int: + return 4 + int(VARIANTS[variant]["contact"]) + + +# ---------------------------------------------------------------------------------------------- +# model + + +class FrozenEncoderUniSAM2(UniSAM2): + """UniSAM2 whose image encoder stays in eval mode while the decoder trains. + + The Hiera encoder has neither dropout nor batch norm, so this is hygiene rather than a numerical + necessity; the freezing itself is done by `freeze_encoder`. + """ + + def train(self, mode: bool = True): + super().train(mode) + self.encoder.eval() + return self + + +def freeze_encoder(model: torch.nn.Module) -> None: + for parameter in model.encoder.parameters(): + parameter.requires_grad_(False) + model.encoder.eval() + + +def decoder_parameters(model: torch.nn.Module) -> List[torch.nn.Parameter]: + """The trainable parameters: everything outside the encoder (the filter `train_joint_sam2` uses).""" + return [p for name, p in model.named_parameters() if not name.startswith("encoder")] + + +def _alias_legacy_modules() -> None: + """Make the module paths pickled into old joint checkpoints importable.""" + import sys + evaluation_dir = os.path.join(os.path.dirname(__file__), "..", "..", "evaluation") + sys.path.insert(0, os.path.abspath(evaluation_dir)) + import common # noqa: F401 (registers the aliases on import when needed) + if hasattr(common, "_alias_micro_sam2_modules"): + common._alias_micro_sam2_modules() + + +def load_lean_v4_states(v4_checkpoint: str = V4_CHECKPOINT, cache_dir: str = CAMPAIGN_ROOT) -> Dict[str, Dict]: + """The 'model_state' (SAM2) and 'unetr_state' (UniSAM2) of the v4 joint checkpoint, without the pickled + trainer state. Cached as a lean file, because the full checkpoint takes minutes to unpickle.""" + os.makedirs(cache_dir, exist_ok=True) + cache_path = os.path.join(cache_dir, "v4_lean_states.pt") + if os.path.exists(cache_path): + return torch.load(cache_path, map_location="cpu", weights_only=True) + try: + state = torch.load(v4_checkpoint, map_location="cpu", weights_only=False) + except (ModuleNotFoundError, AttributeError): + _alias_legacy_modules() + state = torch.load(v4_checkpoint, map_location="cpu", weights_only=False) + + def strip(state_dict): + return {(k[len("module."):] if k.startswith("module.") else k): v for k, v in state_dict.items()} + + lean = {"model_state": strip(state["model_state"]), "unetr_state": strip(state["unetr_state"])} + tmp_path = f"{cache_path}.tmp.{os.getpid()}" + torch.save(lean, tmp_path) + os.replace(tmp_path, cache_path) + return lean + + +def build_model(variant: str, device, unetr_state: Optional[Dict[str, torch.Tensor]]) -> torch.nn.Module: + """Build the (frozen encoder) UniSAM2 for a variant, warm-started from the v4 decoder state if given. + + A five channel decoder takes the four pretrained output rows and keeps the fresh initialisation of the + contact row. + """ + model = FrozenEncoderUniSAM2( + encoder=MODEL_TYPE, output_channels=n_output_channels(variant), initial_features=INITIAL_FEATURES, + device=device, + ) + if unetr_state is not None: + state = dict(unetr_state) + if model.out_channels == state["out_conv.weight"].shape[0]: + model.load_state_dict(state, strict=True) + else: + weight, bias = state.pop("out_conv.weight"), state.pop("out_conv.bias") + missing, unexpected = model.load_state_dict(state, strict=False) + assert sorted(missing) == ["out_conv.bias", "out_conv.weight"] and not unexpected, (missing, unexpected) + with torch.no_grad(): + model.out_conv.weight[:weight.shape[0]].copy_(weight) + model.out_conv.bias[:bias.shape[0]].copy_(bias) + freeze_encoder(model) + return model + + +# ---------------------------------------------------------------------------------------------- +# data + + +def _is_sampler_failure(error: Exception) -> bool: + """torch_em raises this when the min-instance sampler rejects every crop of a file (a 512^2 file with fewer + than three objects can never pass), which must not end the training.""" + return "Could not sample a valid batch" in str(error) + + +class RandomSubsetDataset(torch.utils.data.Dataset): + """A fixed number of random draws from a dataset, redrawing when a file cannot satisfy the sampler. + + torch_em splits 'n_samples' uniformly over the files of a segmentation dataset, so a small sample count + over many files would only ever read the first files, and it retries the same file when the sampler + rejects its crops. This wrapper draws a random index per access and moves on to another file when the + sampler gives up. It exposes 'datasets' so the normalization configuration recurses into the wrapped dataset. + """ + + def __init__(self, dataset, n_samples: int, max_draws: int = 50): + self.datasets = (dataset,) + self.n_samples = int(n_samples) + self.max_draws = int(max_draws) + self.ndim = getattr(dataset, "ndim", 2) + + def __len__(self): + return self.n_samples + + def __getitem__(self, index): + dataset = self.datasets[0] + last_error = None + for _ in range(self.max_draws): + try: + return dataset[np.random.randint(len(dataset))] + except RuntimeError as error: + if not _is_sampler_failure(error): + raise + last_error = error + raise RuntimeError(f"No valid sample in {self.max_draws} random draws.") from last_error + + +class FixedSubsetDataset(torch.utils.data.Dataset): + """The first 'n_samples' indices of a dataset, falling back to the following index when the sampler + rejects a file. Deterministic (validation), like `UniDataWrapper(max_samples=...)` but robust.""" + + def __init__(self, dataset, n_samples: int, max_draws: int = 50): + self.datasets = (dataset,) + self.n_samples = min(int(n_samples), len(dataset)) + self.max_draws = int(max_draws) + self.ndim = getattr(dataset, "ndim", 2) + + def __len__(self): + return self.n_samples + + def __getitem__(self, index): + dataset = self.datasets[0] + last_error = None + for offset in range(self.max_draws): + try: + return dataset[(index + offset) % len(dataset)] + except RuntimeError as error: + if not _is_sampler_failure(error): + raise + last_error = error + raise RuntimeError(f"No valid sample in {self.max_draws} consecutive files from index {index}.") from last_error + + +def _sorted_pairs(raw_paths: Sequence[str], label_paths: Sequence[str]) -> Tuple[List[str], List[str]]: + if len(raw_paths) != len(label_paths): + raise RuntimeError(f"Expect as many raw as label paths, got {len(raw_paths)} and {len(label_paths)}.") + pairs = sorted(zip(raw_paths, label_paths), key=lambda pair: str(pair[0])) + return [str(p[0]) for p in pairs], [str(p[1]) for p in pairs] + + +def split_tail(paths: Sequence[str], fraction: float = VAL_FRACTION, minimum: int = MIN_VAL_FILES): + """Deterministic train / validation split: the tail of the (already sorted) list is the validation set.""" + n_val = max(minimum, int(math.ceil(fraction * len(paths)))) + if n_val >= len(paths): + raise RuntimeError(f"Cannot hold out {n_val} of {len(paths)} files.") + return list(paths[:-n_val]), list(paths[-n_val:]) + + +def _common_kwargs(label_transform, sampler=None): + return { + "patch_shape": PATCH_SHAPE, + "label_transform2": label_transform, + "sampler": MinInstanceSampler(min_num_instances=3, exclude_ids=[0]) if sampler is None else sampler, + "label_dtype": torch.float32, + "ndim": 2, + } + + +def _image_dataset(raw_paths, label_paths, kwargs, raw_transform, n_samples): + """Image / label file pairs (tif, png, ...); the subset wrappers draw the files, so n_samples stays None.""" + return torch_em.default_segmentation_dataset( + raw_paths=raw_paths, raw_key=None, label_paths=label_paths, label_key=None, is_seg_dataset=False, + raw_transform=raw_transform, n_samples=n_samples, **kwargs, + ) + + +def _container_dataset(paths, raw_key, label_key, kwargs, raw_transform, with_channels, patch_shape=None): + """zarr / h5 / tif-stack files read with keys; one patch per file, randomised by `RandomSubsetDataset`.""" + kwargs = dict(kwargs) + if patch_shape is not None: + kwargs["patch_shape"] = patch_shape + return torch_em.default_segmentation_dataset( + raw_paths=paths, raw_key=raw_key, label_paths=paths, label_key=label_key, is_seg_dataset=True, + with_channels=with_channels, raw_transform=raw_transform, n_samples=None, **kwargs, + ) + + +def _wrap(dataset, n_samples: int, is_val: bool): + """Training leaves draw 'n_samples' random samples per epoch; validation leaves read their first samples. + Both skip files the sampler cannot satisfy instead of ending the run.""" + subset = FixedSubsetDataset(dataset, n_samples) if is_val else RandomSubsetDataset(dataset, n_samples) + return UniDataWrapper(subset, source_ndim=2) + + +def _train_count(name: str, scale: float) -> int: + return max(1, int(round(TRAIN_SAMPLES[name] * scale))) + + +def _val_count(name: str) -> int: + return VAL_SAMPLES.get(name, DEFAULT_VAL_SAMPLES) + + +def _tif_is_stack(path: str) -> bool: + import tifffile + with tifffile.TiffFile(path) as f: + return len(f.series[0].shape) == 3 + + +def build_datasets( + data_root: str, label_transform, scale: float = 1.0, +) -> Tuple[List[UniDataWrapper], List[UniDataWrapper], Dict[str, Dict[str, List[str]]]]: + """The training and validation leaves of the nine datasets and the file lists behind them.""" + kwargs = _common_kwargs(label_transform) + train_leaves, val_leaves, manifest = [], [], {} + + def record(name, train_raw, val_raw): + manifest[name] = {"train": list(map(str, train_raw)), "val": list(map(str, val_raw))} + + def add_images(name, raw, labels, raw_transform, sampler_kwargs=None, count_name=None): + count_name = count_name or name + raw, labels = _sorted_pairs(raw, labels) + train_raw, val_raw = split_tail(raw) + train_labels, val_labels = split_tail(labels) + this_kwargs = kwargs if sampler_kwargs is None else _common_kwargs(label_transform, **sampler_kwargs) + train_leaves.append(_wrap( + _image_dataset(train_raw, train_labels, this_kwargs, raw_transform, None), + _train_count(count_name, scale), is_val=False, + )) + val_leaves.append(_wrap( + _image_dataset(val_raw, val_labels, this_kwargs, raw_transform, None), _val_count(count_name), is_val=True, + )) + record(name, train_raw, val_raw) + + def add_containers( + name, paths, raw_key, label_key, raw_transform, with_channels, patch_shape=None, count_name=None, + ): + count_name = count_name or name + paths = sorted(map(str, paths)) + train_paths, val_paths = split_tail(paths) + train_leaves.append(_wrap( + _container_dataset(train_paths, raw_key, label_key, kwargs, raw_transform, with_channels, patch_shape), + _train_count(count_name, scale), is_val=False, + )) + val_leaves.append(_wrap( + _container_dataset(val_paths, raw_key, label_key, kwargs, raw_transform, with_channels, patch_shape), + _val_count(count_name), is_val=True, + )) + record(name, train_paths, val_paths) + + # 1. LIVECell, one dataset per cell type; images that also appear in the val split are dropped. + livecell_root = os.path.join(data_root, "livecell") + for cell_type in datasets.livecell.CELL_TYPES: + raw, labels = datasets.livecell.get_livecell_paths(livecell_root, split="train", cell_types=[cell_type]) + val_raw, _ = datasets.livecell.get_livecell_paths(livecell_root, split="val", cell_types=[cell_type]) + val_names = {os.path.basename(p) for p in val_raw} + keep = [i for i, p in enumerate(raw) if os.path.basename(p) not in val_names] + raw, labels = [raw[i] for i in keep], [labels[i] for i in keep] + add_images( + f"livecell_{cell_type}", raw, labels, _identity, + sampler_kwargs={"sampler": MinInstanceSampler(min_num_instances=6, exclude_ids=[0])}, count_name="livecell", + ) + + # 2. TissueNet: the rgb composite (nucleus, cell, empty) with per-channel normalization, cell labels. + add_containers( + "tissuenet", datasets.tissuenet.get_tissuenet_paths(os.path.join(data_root, "tissuenet"), split="train"), + "raw/rgb", "labels/cell", partial(_normalize_percentile, axis=(1, 2)), with_channels=True, + ) + + # 3. DynamicNuclearNet. + add_containers( + "dynamicnuclearnet", + datasets.dynamicnuclearnet.get_dynamicnuclearnet_paths( + os.path.join(data_root, "dynamicnuclearnet"), split="train", + ), + "raw", "labels", _identity, with_channels=False, + ) + + # 4. DeepBacs (mixed): source / target folders. + image_folder, label_folder = datasets.deepbacs.get_deepbacs_paths( + os.path.join(data_root, "deepbacs"), bac_type="mixed", split="train", + ) + add_images( + "deepbacs", sorted(glob(os.path.join(image_folder, "*.tif"))), + sorted(glob(os.path.join(label_folder, "*.tif"))), _to_8bit, + ) + + # 5. DIC HepG2 (rgb png, three distinct channels). + raw, labels = datasets.dic_hepg2.get_dic_hepg2_paths(os.path.join(data_root, "dic_hepg2"), split="train") + add_images("dic_hepg2", raw, labels, _identity) + + # 6. NeurIPS CellSeg (mixed formats; `_identity` converts to rgb like the getter's make_rgb). + raw, labels = datasets.neurips_cell_seg.get_neurips_cellseg_paths( + os.path.join(data_root, "neurips_cellseg"), split="train", + ) + add_images("neurips_cellseg", raw, labels, _identity) + + # 7. YeaZ: bright field (2d), phase contrast 2d images and phase contrast frame stacks. + yeaz_root = os.path.join(data_root, "yeaz") + bf_raw, bf_labels = datasets.yeaz.get_yeaz_paths(yeaz_root, choice="bf", split="train") + phc_raw, phc_labels = datasets.yeaz.get_yeaz_paths(yeaz_root, choice="phc", split="train") + phc_raw, phc_labels = _sorted_pairs(phc_raw, phc_labels) + is_stack = [_tif_is_stack(p) for p in phc_raw] + phc_2d = ([p for p, s in zip(phc_raw, is_stack) if not s], [p for p, s in zip(phc_labels, is_stack) if not s]) + phc_stacks = ([p for p, s in zip(phc_raw, is_stack) if s], [p for p, s in zip(phc_labels, is_stack) if s]) + groups = { + "yeaz_bf": (_sorted_pairs(bf_raw, bf_labels), PATCH_SHAPE), + "yeaz_phc": (phc_2d, PATCH_SHAPE), + "yeaz_phc_stacks": (phc_stacks, (1,) + PATCH_SHAPE), + } + for name, ((raw, labels), patch_shape) in groups.items(): + train_raw, val_raw = split_tail(raw) + train_labels, val_labels = split_tail(labels) + for split_raw, split_labels, is_val in ((train_raw, train_labels, False), (val_raw, val_labels, True)): + dataset = torch_em.default_segmentation_dataset( + raw_paths=split_raw, raw_key=None, label_paths=split_labels, label_key=None, is_seg_dataset=True, + raw_transform=_identity, n_samples=None, **{**kwargs, "patch_shape": patch_shape}, + ) + leaves = val_leaves if is_val else train_leaves + count = _val_count(name) if is_val else _train_count(name, scale) + leaves.append(_wrap(dataset, count, is_val)) + record(name, train_raw, val_raw) + + # 8. PUMA nuclei (rgb h5). + add_containers( + "puma", datasets.puma.get_puma_paths(os.path.join(data_root, "puma"), split="train", annotations="nuclei"), + "raw", "labels/instances/nuclei", _identity, with_channels=True, + ) + + # 9. TNBC (rgb h5, channel-first). + add_containers( + "tnbc", datasets.tnbc.get_tnbc_paths(os.path.join(data_root, "tnbc"), split="train"), + "raw", "labels/instances", _identity, with_channels=True, + ) + + _configure_training_normalization(train_leaves, val_leaves) + return train_leaves, val_leaves, manifest + + +def build_loaders( + variant: str, data_root: str, batch_size: int, n_workers: int, val_workers: int, scale: float = 1.0, +): + """The train and validation loaders of a variant plus the file manifest.""" + settings = VARIANTS[variant] + label_transform = GeodesicHybridDistanceTransform( + contact=settings["contact"], contact_mode=settings.get("contact_mode", "touching"), + ) + train_leaves, val_leaves, manifest = build_datasets(data_root, label_transform, scale=scale) + train_loader = _prepare_data_loader(ConcatDataset(*train_leaves), batch_size, shuffle=True, num_workers=n_workers) + val_loader = _prepare_data_loader( + ConcatDataset(*val_leaves), batch_size, shuffle=False, num_workers=val_workers, deterministic=True, + ) + return train_loader, val_loader, manifest diff --git a/finetuning/v2/generalist/ais_decoder/evaluate_ais_decoder.sh b/finetuning/v2/generalist/ais_decoder/evaluate_ais_decoder.sh new file mode 100755 index 000000000..3101b6636 --- /dev/null +++ b/finetuning/v2/generalist/ais_decoder/evaluate_ais_decoder.sh @@ -0,0 +1,45 @@ +#!/bin/bash +# Evaluate one trained AIS decoder variant on the AIS benchmarks: stage the checkpoint, cache the predictions +# of the 2d manifests (primary, training_extra, holdout) and the 3d crop manifests (primary, holdout) on the +# cluster, then run the library defaults (and the two contact configurations for five-channel decoders) on the +# caches, one CPU task per (subset, configuration), chained with SLURM dependencies. +# +# bash evaluate_ais_decoder.sh [best|latest] +# +# Afterwards: python optimization/report_ais_decoders.py --subsets primary training_extra --ndim 2 [--output ...] +set -eo pipefail +VARIANT=${1:?variant} +WHICH=${2:-best} +ROOT=/mnt/vast-nhr/projects/cidas/cca/experiments/micro_sam2/apg_optimization +REPO=/mnt/vast-nhr/home/pape41/u12086/Work/my_projects/micro-sam +OPT=$REPO/finetuning/v2/evaluation/optimization +PY=/mnt/vast-nhr/home/pape41/u12086/Work/software/micromamba/envs/envs/new-stack/bin/python +export MICRO_SAM2_JOINT_CHECKPOINT_ROOT=$ROOT/ais_decoder_training/staged +export MICRO_SAM2_JOINT_EXPORT_ROOT=$ROOT/model_exports + +$PY $REPO/finetuning/v2/generalist/ais_decoder/stage_ais_decoder_checkpoint.py --variant "$VARIANT" --which "$WHICH" +CHANNELS=$($PY -c "import json; print(json.load(open('$ROOT/ais_decoder_training/staged/joint_sam2_hvit_t_multi_gpu/$VARIANT.json'))['output_channels'])") +echo "staged $VARIANT ($WHICH): $CHANNELS output channels" + +CONFIGS="" +if [ "$CHANNELS" -gt 4 ]; then + CONFIGS="--configs $OPT/configs/ais_contact_ridge.json $OPT/configs/ais_contact_mask.json" +fi + +newest_job_id() { cat "$(ls -td $ROOT/jobs/*_"$1" | head -1)/job_id.txt"; } + +cd $OPT +$PY ais_campaign_tasks.py predict --name "dec_${VARIANT}_predict2d" --preset 2d --kind v5 \ + --subsets primary training_extra holdout --extra "--joint-checkpoint $VARIANT" +J2D=$(newest_job_id "dec_${VARIANT}_predict2d") +$PY ais_campaign_tasks.py predict --name "dec_${VARIANT}_predict3d" --preset 3d --kind apg3d \ + --subsets primary holdout --extra "--joint-checkpoint $VARIANT" +J3D=$(newest_job_id "dec_${VARIANT}_predict3d") +echo "predict jobs: 2d $J2D, 3d $J3D" + +$PY ais_campaign_tasks.py screen --name "dec_${VARIANT}_screen2d" --preset cpu --kind v5 \ + --subsets primary training_extra holdout --extra "--joint-checkpoint $VARIANT --ndim 2" $CONFIGS \ + --dependency "afterok:$J2D" +$PY ais_campaign_tasks.py screen --name "dec_${VARIANT}_screen3d" --preset cpu --kind apg3d \ + --subsets primary holdout --extra "--joint-checkpoint $VARIANT" $CONFIGS --dependency "afterok:$J3D" +echo "screens submitted (afterok the predictions); check with: squeue -u \$USER" diff --git a/finetuning/v2/generalist/ais_decoder/finalize_ais_decoder_reports.sh b/finetuning/v2/generalist/ais_decoder/finalize_ais_decoder_reports.sh new file mode 100755 index 000000000..27e0a0908 --- /dev/null +++ b/finetuning/v2/generalist/ais_decoder/finalize_ais_decoder_reports.sh @@ -0,0 +1,60 @@ +#!/bin/bash +# Wait until the screens of every variant have finished (all tasks of the newest dec__screen{2d,3d} job +# directories carry a .done marker), then write the final comparison tables and the field diagnostics. +# +# bash finalize_ais_decoder_reports.sh [max_wait_seconds] +# +# Outputs: /ais/reports/decoders_final_{dev,holdout}{,_datasets,_mechanisms}.csv, +# /ais/reports/decoders_final_3d*.csv, /ais/reports/decoder_fields_*.csv +set -o pipefail +VARIANTS=${VARIANTS:-"baseline contact fgcal both"} # override: VARIANTS="boundary boundary_fgcal" bash ... +MAX_WAIT=${1:-32400} +ROOT=/mnt/vast-nhr/projects/cidas/cca/experiments/micro_sam2/apg_optimization +REPO=/mnt/vast-nhr/home/pape41/u12086/Work/my_projects/micro-sam +OPT=$REPO/finetuning/v2/evaluation/optimization +PY=/mnt/vast-nhr/home/pape41/u12086/Work/software/micromamba/envs/envs/new-stack/bin/python +V4=$ROOT/v4_geodesic_checkpoints/joint_sam2_hvit_t_multi_gpu/best.pt +EPOCH=856a433c4b33348e1d85c4c13278f057 +export MICRO_SAM2_JOINT_EXPORT_ROOT=$ROOT/model_exports + +screens_done() { # all tasks of the newest job dir of this name have a .done marker + local dir + dir=$(ls -td "$ROOT"/jobs/*_"$1" 2>/dev/null | head -1) + [ -n "$dir" ] || return 1 + local n_tasks n_done + n_tasks=$(wc -l < "$dir/tasks.txt") + n_done=$(ls "$dir"/logs/*.done 2>/dev/null | wc -l) + [ "$n_done" -ge "$n_tasks" ] +} + +waited=0 +while true; do + pending="" + for v in $VARIANTS; do + for kind in screen2d screen3d; do + screens_done "dec_${v}_${kind}" || pending="$pending dec_${v}_${kind}" + done + done + if [ -z "$pending" ]; then echo "$(date +%H:%M) all screens done"; break; fi + if [ "$waited" -ge "$MAX_WAIT" ]; then echo "$(date +%H:%M) giving up waiting for:$pending"; break; fi + echo "$(date +%H:%M) waiting for:$pending" + sleep 300; waited=$((waited + 300)) +done + +cd "$OPT" +export MICRO_SAM2_JOINT_CHECKPOINT_ROOT=$ROOT/ais_decoder_training/staged +$PY report_ais_decoders.py --variants $VARIANTS --production-checkpoint "$V4" \ + --configs current-defaults contact-ridge contact-mask --subsets primary training_extra --ndim 2 --epoch $EPOCH \ + --output "$ROOT/ais/reports/decoders_final_dev" 2>&1 | grep -v "Warning\|warnings.warn" +$PY report_ais_decoders.py --variants $VARIANTS --production-checkpoint "$V4" \ + --configs current-defaults contact-ridge contact-mask --subsets holdout --ndim 2 --epoch $EPOCH \ + --output "$ROOT/ais/reports/decoders_final_holdout" 2>&1 | grep -v "Warning\|warnings.warn" +$PY report_ais_decoders.py --variants $VARIANTS --production-checkpoint "$V4" \ + --configs current-defaults contact-ridge --kind apg3d --subsets primary holdout --ndim 3 --epoch $EPOCH \ + --output "$ROOT/ais/reports/decoders_final_3d" 2>&1 | grep -v "Warning\|warnings.warn" +for v in $VARIANTS; do + [ -f "$ROOT/ais_decoder_training/staged/joint_sam2_hvit_t_multi_gpu/$v.pt" ] || continue + $PY diagnose_decoder_fields.py --joint-checkpoint "$v" --subset primary training_extra --ndim 2 \ + --output "$ROOT/ais/reports/decoder_fields_$v.csv" 2>&1 | grep -v "Warning\|warnings.warn" | tail -14 +done +echo "$(date +%H:%M) finalisation done" diff --git a/finetuning/v2/generalist/ais_decoder/finalize_round2_reports.sh b/finetuning/v2/generalist/ais_decoder/finalize_round2_reports.sh new file mode 100644 index 000000000..fd619abb8 --- /dev/null +++ b/finetuning/v2/generalist/ais_decoder/finalize_round2_reports.sh @@ -0,0 +1,106 @@ +#!/bin/bash +# Round 2 of the AIS decoder campaign (boundary / boundary_fgcal), unattended: +# 1. screen the shared tuned configuration `dec-top1` (plain, ridge 1, ridge 2 + mask 0.3) on the 2d caches +# of the two new decoders, chained `afterok` on their prediction jobs, +# 2. wait until every round-2 screen has finished (screen2d / screen3d from evaluate_ais_decoder.sh, +# top_screen from step 1, contact_screen from launch_tuning_after_caches.sh), +# 3. write the conclusive overview of all six decoders and the field diagnostics of the two new ones. +# +# bash finalize_round2_reports.sh [max_wait_seconds_for_the_caches] +# +# Submit it with --dependency afterany on the two ais_eval_ jobs, and run a frozen copy: bash +# re-reads a running script by byte offset, so editing this file while a job sleeps in a wait loop breaks it. +# +# Outputs under /ais/reports/: decoders_all_defaults_{dev,holdout}*.csv, decoders_all_tuned_{dev,holdout}*.csv, +# decoders_all_3d*.csv, decoders_boundary_contact_dev*.csv, decoder_fields_{boundary,boundary_fgcal}*.csv +set -o pipefail +NEW=${NEW_VARIANTS:-"boundary boundary_fgcal"} +ALL=${VARIANTS:-"baseline contact fgcal both boundary boundary_fgcal"} +MAX_WAIT=${1:-7200} +ROOT=/mnt/vast-nhr/projects/cidas/cca/experiments/micro_sam2/apg_optimization +REPO=/mnt/vast-nhr/home/pape41/u12086/Work/my_projects/micro-sam +OPT=$REPO/finetuning/v2/evaluation/optimization +PY=/mnt/vast-nhr/home/pape41/u12086/Work/software/micromamba/envs/envs/new-stack/bin/python +V4=$ROOT/v4_geodesic_checkpoints/joint_sam2_hvit_t_multi_gpu/best.pt +EPOCH=856a433c4b33348e1d85c4c13278f057 +export MICRO_SAM2_JOINT_CHECKPOINT_ROOT=$ROOT/ais_decoder_training/staged +export MICRO_SAM2_JOINT_EXPORT_ROOT=$ROOT/model_exports +TOP_CONFIGS="$OPT/configs/ais_dec_top1.json $OPT/configs/ais_dec_top1_ridge1.json $OPT/configs/ais_dec_top1_ridge2_mask0.3.json" +REPORTS=$ROOT/ais/reports + +job_dir() { ls -td "$ROOT"/jobs/*_"$1" 2>/dev/null | head -1; } +tasks_done() { # all tasks of the newest job dir of this name carry a .done marker + local dir; dir=$(job_dir "$1"); [ -n "$dir" ] || return 1 + [ "$(ls "$dir"/logs/*.done 2>/dev/null | wc -l)" -ge "$(wc -l < "$dir/tasks.txt")" ] +} +wait_for() { # wait_for + local limit=$1; shift; local waited=0 + while true; do + local pending="" + for name in "$@"; do tasks_done "$name" || pending="$pending $name"; done + [ -z "$pending" ] && { echo "$(date +%H:%M) all done"; return 0; } + [ "$waited" -ge "$limit" ] && { echo "$(date +%H:%M) timeout waiting for:$pending"; return 1; } + echo "$(date +%H:%M) waiting for:$pending"; sleep 300; waited=$((waited + 300)) + done +} + +cd "$OPT" || exit 1 + +# 1. The dec-top1 screens of the two new decoders, on their 2d prediction jobs. +for v in $NEW; do + waited=0 + while [ -z "$(job_dir "dec_${v}_predict2d")" ]; do + [ "$waited" -ge "$MAX_WAIT" ] && { echo "$(date +%H:%M) no dec_${v}_predict2d job dir, skipping its top screen"; break; } + echo "$(date +%H:%M) waiting for the dec_${v}_predict2d job dir"; sleep 120; waited=$((waited + 120)) + done + d=$(job_dir "dec_${v}_predict2d"); [ -n "$d" ] || continue + if [ -n "$(job_dir "dec_${v}_top_screen")" ]; then echo "$(date +%H:%M) dec_${v}_top_screen exists already"; continue; fi + j=$(cat "$d/job_id.txt") + echo "$(date +%H:%M) submitting dec_${v}_top_screen (afterok:$j)" + $PY ais_campaign_tasks.py screen --name "dec_${v}_top_screen" --preset cpu --kind v5 \ + --subsets primary training_extra holdout --no-defaults --configs $TOP_CONFIGS \ + --extra "--joint-checkpoint $v --ndim 2" --dependency "afterok:$j" +done + +# 2. Every round-2 screen (the contact screens come from launch_tuning_after_caches.sh). +names="" +for v in $NEW; do + names="$names dec_${v}_screen2d dec_${v}_screen3d dec_${v}_top_screen dec_${v}_contact_screen" +done +wait_for 28800 $names || true + +# 3. The conclusive overview of all six decoders. +echo "$(date +%H:%M) writing the overview" +$PY report_ais_decoders.py --variants $ALL --production-checkpoint "$V4" --baseline-variant baseline \ + --configs current-defaults contact-ridge contact-mask --subsets primary training_extra --ndim 2 --epoch $EPOCH \ + --output "$REPORTS/decoders_all_defaults_dev" 2>&1 | grep -v "Warning\|warnings.warn" +$PY report_ais_decoders.py --variants $ALL --production-checkpoint "$V4" --baseline-variant baseline \ + --configs current-defaults contact-ridge contact-mask --subsets holdout --ndim 2 --epoch $EPOCH \ + --output "$REPORTS/decoders_all_defaults_holdout" 2>&1 | grep -v "Warning\|warnings.warn" +for pair in dev:"primary training_extra" holdout:holdout; do + $PY report_ais_decoders.py --variants $ALL --production-checkpoint "$V4" --baseline-variant baseline \ + --baseline-config dec-top1 --configs current-defaults dec-top1 dec-fgcal-top1 dec-top1-ridge1 dec-top1-ridge2-mask0.3 \ + --subsets ${pair##*:} --ndim 2 --epoch $EPOCH \ + --output "$REPORTS/decoders_all_tuned_${pair%%:*}" 2>&1 | grep -v "Warning\|warnings.warn" +done +$PY report_ais_decoders.py --variants $ALL --production-checkpoint "$V4" --baseline-variant baseline \ + --configs current-defaults contact-ridge --kind apg3d --subsets primary holdout --ndim 3 --epoch $EPOCH \ + --output "$REPORTS/decoders_all_3d" 2>&1 | grep -v "Warning\|warnings.warn" +# The boundary channel's ridge and mask settings against the decoder's own defaults. +for v in $NEW; do + $PY report_ais_decoders.py --variants "$v" --baseline-variant "$v" --configs current-defaults \ + contact-ridge-w0.5 contact-ridge contact-ridge-w2.0 contact-ridge-w4.0 contact-mask-t0.3 contact-mask \ + contact-mask-t0.7 contact-ridge1-mask0.5 --subsets primary training_extra --ndim 2 --epoch $EPOCH \ + --output "$REPORTS/decoders_${v}_contact_dev" 2>&1 | grep -v "Warning\|warnings.warn" +done +# Field diagnostics. --contact-mode must match the training target of the fifth channel, otherwise the head's +# precision is scored against a target that calls its correct pixels negative; `both` is rescored in the +# touching mode so that the round-1 reference carries the new recall_touching / recall_bg_boundary columns. +for pair in boundary:all boundary_fgcal:all both:touching; do + v=${pair%%:*}; mode=${pair##*:} + case " $NEW both " in *" $v "*) ;; *) continue ;; esac + [ -f "$ROOT/ais_decoder_training/staged/joint_sam2_hvit_t_multi_gpu/$v.pt" ] || continue + $PY diagnose_decoder_fields.py --joint-checkpoint "$v" --subset primary training_extra --ndim 2 \ + --contact-mode "$mode" --output "$REPORTS/decoder_fields_$v.csv" 2>&1 | grep -v "Warning\|warnings.warn" | tail -14 +done +echo "$(date +%H:%M) round-2 finalisation done" diff --git a/finetuning/v2/generalist/ais_decoder/launch_tuning_after_caches.sh b/finetuning/v2/generalist/ais_decoder/launch_tuning_after_caches.sh new file mode 100755 index 000000000..d51d9cd2b --- /dev/null +++ b/finetuning/v2/generalist/ais_decoder/launch_tuning_after_caches.sh @@ -0,0 +1,85 @@ +#!/bin/bash +# Once the 2d prediction caches of baseline and contact exist, submit their grid sweeps (and the contact +# configuration screens for the five-channel 'contact' decoder); then, when every variant's sweeps are done, +# rank each sweep (report_ais_sweep.py) into /ais/reports/dec__sweep_dev.csv. +# +# bash launch_tuning_after_caches.sh [max_wait_seconds] [--wait V...] [--rank V...] +# +# WARNING: pass the variants as ARGUMENTS, never through the environment. `SBATCH_EXPORT=none` is set on this +# system, so `sbatch` does NOT propagate the submitting environment and the `WAIT_VARIANTS` / `VARIANTS` +# variables silently fall back to the round-1 defaults below - which is exactly what happened on 2026-09-08 +# (job 15777359 re-ran the four round-1 sweeps and never submitted the round-2 ones). The environment variables +# are still honoured when the script is run directly in a shell. +set -o pipefail +WAIT_ARGS=""; RANK_ARGS=""; POSITIONAL=""; mode="" +for a in "$@"; do + case "$a" in + --wait) mode=wait ;; + --rank) mode=rank ;; + *) case "$mode" in wait) WAIT_ARGS="$WAIT_ARGS $a" ;; rank) RANK_ARGS="$RANK_ARGS $a" ;; + *) POSITIONAL="$POSITIONAL $a" ;; esac ;; + esac +done +set -- $POSITIONAL +[ -n "$WAIT_ARGS" ] && WAIT_VARIANTS="$WAIT_ARGS" +[ -n "$RANK_ARGS" ] && VARIANTS="$RANK_ARGS" +VARIANTS=${VARIANTS:-"baseline contact fgcal both"} # override: ... --rank boundary boundary_fgcal +MAX_WAIT=${1:-32400} +ROOT=/mnt/vast-nhr/projects/cidas/cca/experiments/micro_sam2/apg_optimization +OPT=/mnt/vast-nhr/home/pape41/u12086/Work/my_projects/micro-sam/finetuning/v2/evaluation/optimization +PY=/mnt/vast-nhr/home/pape41/u12086/Work/software/micromamba/envs/envs/new-stack/bin/python +export MICRO_SAM2_JOINT_CHECKPOINT_ROOT=$ROOT/ais_decoder_training/staged +export MICRO_SAM2_JOINT_EXPORT_ROOT=$ROOT/model_exports +PRIMARY="livecell tissuenet dynamicnuclearnet deepbacs dic_hepg2" +EXTRA="yeaz neurips_cellseg deepseas puma tnbc covid_if" +CONTACT_CONFIGS="$OPT/configs/ais_contact_ridge_w0.5.json $OPT/configs/ais_contact_ridge_w2.0.json $OPT/configs/ais_contact_ridge_w4.0.json $OPT/configs/ais_contact_mask_t0.3.json $OPT/configs/ais_contact_mask_t0.7.json $OPT/configs/ais_contact_ridge1_mask0.5.json $OPT/configs/ais_contact_ridge.json $OPT/configs/ais_contact_mask.json" + +tasks_done() { # all tasks of the newest job dir of this name carry a .done marker + local dir; dir=$(ls -td "$ROOT"/jobs/*_"$1" 2>/dev/null | head -1); [ -n "$dir" ] || return 1 + [ "$(ls "$dir"/logs/*.done 2>/dev/null | wc -l)" -ge "$(wc -l < "$dir/tasks.txt")" ] +} +wait_for() { # wait_for + local limit=$1; shift; local waited=0 + while true; do + local pending="" + for name in "$@"; do tasks_done "$name" || pending="$pending $name"; done + [ -z "$pending" ] && return 0 + [ "$waited" -ge "$limit" ] && { echo "$(date +%H:%M) timeout waiting for:$pending"; return 1; } + echo "$(date +%H:%M) waiting for:$pending"; sleep 300; waited=$((waited + 300)) + done +} + +cd "$OPT" +declare -A launched +while true; do + for v in ${WAIT_VARIANTS:-baseline contact}; do + [ -n "${launched[$v]}" ] && continue + if tasks_done "dec_${v}_predict2d"; then + echo "$(date +%H:%M) caches of $v ready, submitting sweeps" + $PY ais_campaign_tasks.py sweep --name "dec_${v}_sweep_primary" --preset cpu --kind v5 --subsets primary \ + --grid configs/ais_grid_lm_v4.json --datasets $PRIMARY --num-shards 1 --extra "--joint-checkpoint $v" + $PY ais_campaign_tasks.py sweep --name "dec_${v}_sweep_extra" --preset cpu --kind v5 --subsets training_extra \ + --grid configs/ais_grid_lm_v4.json --datasets $EXTRA --num-shards 1 --extra "--joint-checkpoint $v" + if [ "$v" = "contact" ] || [ "$v" = "boundary" ] || [ "$v" = "boundary_fgcal" ]; then + $PY ais_campaign_tasks.py screen --name "dec_${v}_contact_screen" --preset cpu --kind v5 \ + --subsets primary training_extra holdout --no-defaults --configs $CONTACT_CONFIGS \ + --extra "--joint-checkpoint $v --ndim 2" + fi + launched[$v]=1 + fi + done + all_launched=1; for v in ${WAIT_VARIANTS:-baseline contact}; do [ -n "${launched[$v]}" ] || all_launched=0; done + [ "$all_launched" = 1 ] && break + [ "$MAX_WAIT" -le 0 ] && { echo "$(date +%H:%M) gave up waiting for the caches"; break; } + sleep 300; MAX_WAIT=$((MAX_WAIT - 300)) +done + +names="" +for v in $VARIANTS; do names="$names dec_${v}_sweep_primary dec_${v}_sweep_extra"; done +wait_for 14400 $names || true +for v in $VARIANTS; do + tasks_done "dec_${v}_sweep_primary" && tasks_done "dec_${v}_sweep_extra" || { echo "sweeps of $v incomplete, skipping the ranking"; continue; } + $PY report_ais_sweep.py --grid configs/ais_grid_lm_v4.json --subset primary training_extra --joint-checkpoint "$v" \ + --top 25 --output "$ROOT/ais/reports/dec_${v}_sweep_dev.csv" 2>&1 | grep -v "Warning\|warnings.warn" | tail -40 +done +echo "$(date +%H:%M) tuning launcher done" diff --git a/finetuning/v2/generalist/ais_decoder/stage_ais_decoder_checkpoint.py b/finetuning/v2/generalist/ais_decoder/stage_ais_decoder_checkpoint.py new file mode 100644 index 000000000..cdfa858ac --- /dev/null +++ b/finetuning/v2/generalist/ais_decoder/stage_ais_decoder_checkpoint.py @@ -0,0 +1,76 @@ +"""Stage a trained AIS decoder as a lean joint-format checkpoint for the evaluation harness. + +Writes /joint_sam2_hvit_t_multi_gpu/.pt with the v4 SAM2 'model_state' (the frozen +encoder equals the v4 encoder, so the interactive half is the production one) and the trained 'unetr_state'. +Then `export MICRO_SAM2_JOINT_CHECKPOINT_ROOT=` and select the model with +`--joint-checkpoint ` in benchmark_ais_optimization.py; the checkpoint id is the file's checksum. + + python stage_ais_decoder_checkpoint.py --variant contact [--which best|latest] +""" + +import argparse +import json +import os +import subprocess +import sys + +import torch + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import ais_decoder_lib as lib # noqa: E402,F401 (registers the classes the trainer pickled) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--variant", required=True, choices=sorted(lib.VARIANTS)) + parser.add_argument("--which", default="best", choices=["best", "latest"]) + parser.add_argument("--save-root", default=lib.CAMPAIGN_ROOT) + parser.add_argument("--staged-root", default=None, help="Default /staged.") + parser.add_argument("--name", default=None, help="Checkpoint name, default ais_decoder_.") + parser.add_argument("--v4-checkpoint", default=lib.V4_CHECKPOINT) + args = parser.parse_args() + + name = args.name or f"ais_decoder_{args.variant}" + checkpoint_path = os.path.join(args.save_root, "checkpoints", name, f"{args.which}.pt") + staged_root = args.staged_root or os.path.join(args.save_root, "staged") + staged_dir = os.path.join(staged_root, f"joint_sam2_{lib.MODEL_TYPE}_multi_gpu") + os.makedirs(staged_dir, exist_ok=True) + staged_path = os.path.join(staged_dir, f"{args.variant}.pt") + + trained = torch.load(checkpoint_path, map_location="cpu", weights_only=False) + unetr_state = {k: v for k, v in trained["model_state"].items()} + v4 = lib.load_lean_v4_states(args.v4_checkpoint, args.save_root) + # The encoder was frozen: the trained state must carry the v4 encoder unchanged. + for key, value in v4["unetr_state"].items(): + if key.startswith("encoder.") and not torch.equal(value, unetr_state[key]): + raise RuntimeError(f"The encoder weights changed during training ({key}); refusing to stage.") + try: + revision = subprocess.check_output( + ["git", "rev-parse", "HEAD"], cwd=os.path.dirname(__file__), text=True, + ).strip() + except Exception: # noqa: BLE001 + revision = None + lean = { + "model_state": v4["model_state"], + "unetr_state": unetr_state, + "source": { + "variant": args.variant, "checkpoint": checkpoint_path, "which": args.which, + "iteration": int(trained.get("iteration", -1)), "epoch": int(trained.get("epoch", -1)), + "best_epoch": int(trained.get("best_epoch", -1)), + "best_metric": float(trained.get("best_metric", float("nan"))), + "current_metric": float(trained.get("current_metric", float("nan"))), "git_revision": revision, + "output_channels": int(unetr_state["out_conv.weight"].shape[0]), + }, + } + tmp_path = f"{staged_path}.tmp.{os.getpid()}" + torch.save(lean, tmp_path) + os.replace(tmp_path, staged_path) + with open(os.path.join(staged_dir, f"{args.variant}.json"), "w") as f: + json.dump(lean["source"], f, indent=2) + print(f"staged {checkpoint_path} -> {staged_path}") + print(json.dumps(lean["source"], indent=2)) + print(f"export MICRO_SAM2_JOINT_CHECKPOINT_ROOT={staged_root}") + + +if __name__ == "__main__": + main() diff --git a/finetuning/v2/generalist/ais_decoder/submit_ais_decoder_training.py b/finetuning/v2/generalist/ais_decoder/submit_ais_decoder_training.py new file mode 100644 index 000000000..2cf0411ae --- /dev/null +++ b/finetuning/v2/generalist/ais_decoder/submit_ais_decoder_training.py @@ -0,0 +1,104 @@ +"""Submit the AIS decoder trainings as single-GPU SLURM jobs on grete:shared. + + python submit_ais_decoder_training.py --variants baseline contact fgcal --iterations 30000 --batch-size 8 --dry + python submit_ais_decoder_training.py --variants both --iterations 30000 --batch-size 8 --after 1234 1235 1236 + +Writes /jobs/_.sh and a submit.json with the job ids. `--after` makes the job +start only after the listed jobs have started (SLURM 'after' dependency), which is how the fourth model waits +for the other three. +""" + +import argparse +import datetime +import json +import os +import subprocess +import sys + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import ais_decoder_lib as lib # noqa: E402 + +REPOSITORY_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..", "..")) +TRAIN_SCRIPT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "train_ais_decoder.py") + +TEMPLATE = """#!/bin/bash +#SBATCH --job-name=ais_decoder_{variant} +#SBATCH -p {partition} +#SBATCH -G {gres} +#SBATCH -c {cpus} +#SBATCH --mem={mem} +#SBATCH -t {time} +#SBATCH --constraint=inet +#SBATCH -A {account} +#SBATCH -o {log_dir}/ais_decoder_{variant}_%j.out +#SBATCH -e {log_dir}/ais_decoder_{variant}_%j.err +{dependency} +set -eo pipefail +source ~/.bashrc +set -u +micromamba activate {env} +cd {repository} +export PYTHONUNBUFFERED=1 +export OMP_NUM_THREADS=1 +export MKL_NUM_THREADS=1 +nvidia-smi --query-gpu=name,memory.total --format=csv +python {script} --variant {variant} --iterations {iterations} --batch-size {batch_size} \\ + --n-workers {n_workers} --val-workers {val_workers} --lr {lr} --epoch-scale {epoch_scale} \\ + --save-root {save_root} {extra} +""" + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--variants", nargs="+", required=True, choices=sorted(lib.VARIANTS)) + parser.add_argument("--iterations", type=int, required=True) + parser.add_argument("--batch-size", type=int, default=8) + parser.add_argument("--n-workers", type=int, default=12) + parser.add_argument("--val-workers", type=int, default=3) + parser.add_argument("--lr", type=float, default=5e-5) + parser.add_argument("--epoch-scale", type=float, default=1.0) + parser.add_argument("--save-root", default=lib.CAMPAIGN_ROOT) + parser.add_argument("--partition", default="grete:shared") + parser.add_argument("--gres", default="A100:1") + parser.add_argument("--cpus", type=int, default=16) + parser.add_argument("--mem", default="64G") + parser.add_argument("--time", default="12:00:00") + parser.add_argument("--account", default="nim00007") + parser.add_argument("--env", default="new-stack") + parser.add_argument("--after", nargs="*", default=None, help="Job ids this job waits for (start after they start).") + parser.add_argument("--extra", default="", help="Extra arguments for train_ais_decoder.py, verbatim.") + parser.add_argument("--dry", action="store_true") + args = parser.parse_args() + + stamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") + job_dir = os.path.join(args.save_root, "jobs") + log_dir = os.path.join(args.save_root, "logs", "slurm") + os.makedirs(job_dir, exist_ok=True) + os.makedirs(log_dir, exist_ok=True) + dependency = f"#SBATCH --dependency=after:{':'.join(args.after)}" if args.after else "" + submitted = {} + for variant in args.variants: + script = TEMPLATE.format( + variant=variant, partition=args.partition, gres=args.gres, cpus=args.cpus, mem=args.mem, time=args.time, + account=args.account, log_dir=log_dir, dependency=dependency, env=args.env, repository=REPOSITORY_ROOT, + script=TRAIN_SCRIPT, iterations=args.iterations, batch_size=args.batch_size, n_workers=args.n_workers, + val_workers=args.val_workers, lr=args.lr, epoch_scale=args.epoch_scale, save_root=args.save_root, + extra=args.extra, + ) + script_path = os.path.join(job_dir, f"{stamp}_{variant}.sh") + with open(script_path, "w") as f: + f.write(script) + if args.dry: + print(f"--- {script_path} ---\n{script}") + continue + job_id = subprocess.check_output(["sbatch", "--parsable", script_path], text=True).strip().split(";")[0] + submitted[variant] = job_id + print(f"submitted {variant}: job {job_id} ({script_path})") + if submitted: + record = {"timestamp": stamp, "argv": sys.argv, "jobs": submitted} + with open(os.path.join(job_dir, f"{stamp}_submit.json"), "w") as f: + json.dump(record, f, indent=2) + + +if __name__ == "__main__": + main() diff --git a/finetuning/v2/generalist/ais_decoder/train_ais_decoder.py b/finetuning/v2/generalist/ais_decoder/train_ais_decoder.py new file mode 100644 index 000000000..bcb9bbf37 --- /dev/null +++ b/finetuning/v2/generalist/ais_decoder/train_ais_decoder.py @@ -0,0 +1,170 @@ +"""Train one AIS decoder variant (decoder only, encoder frozen at the v4 joint weights). + +Example (smoke test on the session GPU, then a full run): + python train_ais_decoder.py --variant contact --smoke 30 --batch-size 4 --n-workers 1 + python train_ais_decoder.py --variant contact --iterations 30000 --batch-size 8 --n-workers 12 + +Checkpoints: /checkpoints/ais_decoder_/{best,latest}.pt (torch_em layout), the file +lists behind the loaders: /checkpoints/ais_decoder_/data_manifest.json. +""" + +import argparse +import json +import os +import sys +import time + +import torch +import torch.multiprocessing + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import ais_decoder_lib as lib # noqa: E402 + +from micro_sam.util import training_autocast_dtype # noqa: E402 +from micro_sam.v2.datasets.util import check_loader # noqa: E402 +from micro_sam.v2.loss import DirectedDistanceLoss # noqa: E402 +from micro_sam.v2.training.sam2_trainer import UniSAM2Logger, UniSAM2Trainer # noqa: E402 + + +def parse_args(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--variant", required=True, choices=sorted(lib.VARIANTS)) + parser.add_argument("--iterations", type=int, default=None, help="Training iterations (required unless --smoke).") + parser.add_argument("--batch-size", type=int, default=8) + parser.add_argument("--n-workers", type=int, default=12, help="Train loader workers.") + parser.add_argument("--val-workers", type=int, default=3) + parser.add_argument("--lr", type=float, default=5e-5) + parser.add_argument("--epoch-scale", type=float, default=1.0, + help="Multiplier of the per-dataset samples per epoch (base ~1450 samples).") + parser.add_argument("--save-root", default=lib.CAMPAIGN_ROOT) + parser.add_argument("--data-root", default=lib.DATA_ROOT) + parser.add_argument("--init-checkpoint", default=lib.V4_CHECKPOINT, help="Joint checkpoint to warm-start from.") + parser.add_argument("--no-warm-start", action="store_true", help="Random decoder init (not used in the campaign).") + parser.add_argument("--name", default=None, help="Checkpoint name, default ais_decoder_.") + parser.add_argument("--log-image-interval", type=int, default=100) + parser.add_argument("--resume", default=None, help="Trainer checkpoint to resume from.") + parser.add_argument("--smoke", type=int, default=None, + help="Smoke test: time the loader, run this many iterations plus one validation, report.") + parser.add_argument("--device", default=None) + return parser.parse_args() + + +def build_trainer(args, model, train_loader, val_loader, device, name): + settings = lib.VARIANTS[args.variant] + loss = DirectedDistanceLoss( + mask_distances_in_bg=True, contact=settings["contact"], boundary_weight=settings["boundary_weight"], + boundary_radius=lib.BOUNDARY_RADIUS, + ) + optimizer = torch.optim.AdamW(lib.decoder_parameters(model), lr=args.lr, weight_decay=0.1) + scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode="min", factor=0.9, patience=10) + return UniSAM2Trainer( + name=name, model=model, train_loader=train_loader, val_loader=val_loader, loss=loss, metric=loss, + optimizer=optimizer, device=device, lr_scheduler=scheduler, + mixed_precision=training_autocast_dtype(device) is not None, mixed_precision_dtype="bfloat16", + early_stopping=None, log_image_interval=args.log_image_interval, logger=UniSAM2Logger, logger_kwargs=None, + id_=None, save_root=args.save_root, compile_model=False, rank=None, + ) + + +def time_loader(loader, n_batches): + started = time.perf_counter() + n_samples = 0 + for index, (x, y) in enumerate(loader): + n_samples += x.shape[0] + if index + 1 >= n_batches: + break + seconds = time.perf_counter() - started + return seconds, n_samples + + +def time_gpu_step(trainer, x, y, n_steps): + """The GPU-only cost of one training step on a fixed batch (forward, loss, backward, optimizer step).""" + trainer.model.train() + x, y = x.to(trainer.device), y.to(trainer.device) + dtype = torch.bfloat16 if trainer.mixed_precision else None + times = [] + for step in range(n_steps + 3): + torch.cuda.synchronize() + started = time.perf_counter() + trainer.optimizer.zero_grad() + with torch.autocast(device_type="cuda", dtype=dtype, enabled=dtype is not None): + prediction = trainer.model(x) + loss = trainer.loss(prediction, y) + loss.backward() + trainer.optimizer.step() + torch.cuda.synchronize() + if step >= 3: + times.append(time.perf_counter() - started) + return sum(times) / len(times), float(loss.detach()) + + +def main(): + args = parse_args() + if args.iterations is None and args.smoke is None: + raise SystemExit("Pass --iterations or --smoke.") + # Python 3.14 starts worker processes through a fork server by default; every loader worker then re-imports + # the whole environment (30-60 s each, serialised) and the validation workers do so every epoch. Forking + # copies the parent instead; the workers never touch CUDA, so forking after the model was built is safe. + torch.multiprocessing.set_start_method("fork", force=True) + torch.set_num_threads(2) + device = torch.device(args.device or ("cuda" if torch.cuda.is_available() else "cpu")) + name = args.name or (f"smoke_{args.variant}" if args.smoke else f"ais_decoder_{args.variant}") + n_channels = lib.n_output_channels(args.variant) + print(f"variant {args.variant}: {n_channels} output channels, loss settings {lib.VARIANTS[args.variant]}") + + train_loader, val_loader, manifest = lib.build_loaders( + args.variant, args.data_root, args.batch_size, args.n_workers, args.val_workers, scale=args.epoch_scale, + ) + print(f"train: {len(train_loader.dataset)} samples per epoch in {len(train_loader)} iterations of " + f"{args.batch_size}; validation: {len(val_loader.dataset)} samples") + for dataset, lists in manifest.items(): + print(f" {dataset:22s} train files {len(lists['train']):5d} val files {len(lists['val']):3d}") + + unetr_state = None + if not args.no_warm_start: + unetr_state = lib.load_lean_v4_states(args.init_checkpoint, args.save_root)["unetr_state"] + model = lib.build_model(args.variant, device, unetr_state) + n_trainable = sum(p.numel() for p in lib.decoder_parameters(model)) + n_frozen = sum(p.numel() for p in model.encoder.parameters()) + print(f"model: {n_trainable / 1e6:.2f} M trainable decoder parameters, " + f"{n_frozen / 1e6:.2f} M frozen encoder parameters") + + trainer = build_trainer(args, model, train_loader, val_loader, device, name) + checkpoint_dir = os.path.join(args.save_root, "checkpoints", name) + os.makedirs(checkpoint_dir, exist_ok=True) + with open(os.path.join(checkpoint_dir, "data_manifest.json"), "w") as f: + json.dump({"variant": args.variant, "args": vars(args), "datasets": manifest}, f, indent=2) + + if args.smoke: + check_loader(train_loader, n_samples=3, n_target_channels=n_channels) + seconds, n_samples = time_loader(train_loader, n_batches=max(2, args.smoke // 5)) + print(f"[smoke] loader: {n_samples / seconds:.2f} samples/s with {args.n_workers} workers " + f"({seconds / max(1, n_samples // args.batch_size):.2f} s per batch of {args.batch_size})") + x, y = next(iter(train_loader)) + if device.type == "cuda": + torch.cuda.reset_peak_memory_stats(device) + step_seconds, loss_value = time_gpu_step(trainer, x, y, n_steps=max(5, args.smoke // 3)) + print(f"[smoke] gpu step: {step_seconds:.3f} s per iteration at batch {args.batch_size} " + f"(loss {loss_value:.4f})") + if device.type == "cuda": + print(f"[smoke] peak memory after gpu steps: {torch.cuda.max_memory_allocated(device) / 2**30:.2f} GiB " + f"allocated, {torch.cuda.max_memory_reserved(device) / 2**30:.2f} GiB reserved") + started = time.perf_counter() + trainer.fit(iterations=args.smoke, overwrite_training=True) + fit_seconds = time.perf_counter() - started + print(f"[smoke] fit: {args.smoke} iterations + validation ({len(val_loader)} batches) in {fit_seconds:.1f} s") + if device.type == "cuda": + print(f"[smoke] peak memory overall: {torch.cuda.max_memory_allocated(device) / 2**30:.2f} GiB allocated, " + f"{torch.cuda.max_memory_reserved(device) / 2**30:.2f} GiB reserved") + return + + started = time.perf_counter() + trainer.fit(iterations=args.iterations, overwrite_training=args.resume is None, load_from_checkpoint=args.resume) + print(f"training finished after {(time.perf_counter() - started) / 3600:.2f} h") + if device.type == "cuda": + print(f"[peak-memory] {torch.cuda.max_memory_allocated(device) / 2**30:.2f} GiB allocated, " + f"{torch.cuda.max_memory_reserved(device) / 2**30:.2f} GiB reserved") + + +if __name__ == "__main__": + main() diff --git a/micro_sam/v2/automatic_prompt_generation.py b/micro_sam/v2/automatic_prompt_generation.py index f71b3bcb4..eb37c41ad 100644 --- a/micro_sam/v2/automatic_prompt_generation.py +++ b/micro_sam/v2/automatic_prompt_generation.py @@ -43,7 +43,6 @@ from sam2.utils.amg import calculate_stability_score -from bioimage_cpp.utils import Blocking from bioimage_cpp.segmentation import label # Only the tiled stitching in 'TiledAutomaticPromptGenerator.generate' uses this, so a missing @@ -1286,8 +1285,6 @@ def generate( batch_size: int = DEFAULT_PROMPT_GENERATION["batch_size"], n_threads: int = DEFAULT_PROMPT_GENERATION["n_threads"], verbose: bool = False, - pbar_init: Optional[Callable] = None, - pbar_update: Optional[Callable] = None, ) -> np.ndarray: """Derive prompts from the stored predictions, apply them and merge the masks. @@ -1338,8 +1335,6 @@ def generate( batch_size: Number of prompts per forward pass. n_threads: Number of threads for the flow integration the candidates come from. verbose: Whether to show progress over the propagation passes of a volume. - pbar_init: Initialize an external progress stage with its total and description. - pbar_update: Advance the external progress bar after completed work. Returns: The instance segmentation, uint32 array with the spatial shape of the prediction. @@ -1365,16 +1360,12 @@ def generate( components = resolved = None if refinement is not None: components, resolved = _parse_refinement(refinement, refinement_kwargs, is_volume=True) - if pbar_init is not None: - pbar_init(1, "APG: deriving volume prompts") prompts = derive_volume_prompts( - self._prediction[0], self._prediction[1:], model_type=self._model_type, + self._prediction[0], self._prediction[1:4], model_type=self._model_type, candidate_threshold=candidate_threshold, foreground_threshold=foreground_threshold, n_iter=n_iter, dt=dt, sigma=sigma, spacing=spacing, min_candidate_size=min_candidate_size, n_threads=n_threads, ) - if pbar_update is not None: - pbar_update(1) if prompts is None: self._last_generation_stats = { "proposed_candidates": 0, @@ -1396,7 +1387,6 @@ def generate( prompts, multimasking=multimasking, batch_size=batch_size, score_threshold=score_threshold, max_overlap=max_overlap, components=components, refinement_kwargs=resolved, - pbar_init=pbar_init, pbar_update=pbar_update, ) else: with autocast(self._predictor.device): @@ -1404,7 +1394,6 @@ def generate( prompts, multimasking=multimasking, batch_size=batch_size, score_threshold=score_threshold, max_overlap=max_overlap, components=components, refinement_kwargs=resolved, - pbar_init=pbar_init, pbar_update=pbar_update, ) self._last_generation_stats["scored_candidates"] = len(candidates) records = self._propagate_candidates( @@ -1414,21 +1403,16 @@ def generate( ) # Tiled records arrive grouped by tile and need their halo overlaps resolved, which # '_merge' does polymorphically; an untiled volume merges them flat. - if pbar_init is not None: - pbar_init(1, "APG: merging volume masks") segmentation, _ = self._merge( records, shape, score_threshold=score_threshold, max_overlap=max_overlap, min_size=min_size, max_size_factor=max_size_factor, ) - if pbar_update is not None: - pbar_update(1) return segmentation proposals = self.propose( candidate_threshold=candidate_threshold, foreground_threshold=foreground_threshold, n_iter=n_iter, dt=dt, sigma=sigma, min_candidate_size=min_candidate_size, multimasking=multimasking, batch_size=batch_size, n_threads=n_threads, - pbar_init=pbar_init, pbar_update=pbar_update, ) return self.select( proposals, score_threshold=score_threshold, max_overlap=max_overlap, min_size=min_size, @@ -1447,8 +1431,6 @@ def propose( multimasking: bool = DEFAULT_PROMPT_GENERATION["multimasking"], batch_size: int = DEFAULT_PROMPT_GENERATION["batch_size"], n_threads: int = DEFAULT_PROMPT_GENERATION["n_threads"], - pbar_init: Optional[Callable] = None, - pbar_update: Optional[Callable] = None, ) -> list: """Derive the prompts and turn them into scored mask proposals, without selecting any of them. @@ -1467,8 +1449,6 @@ def propose( multimasking: Whether to predict several masks per point and keep the best scoring one. batch_size: Number of prompts per forward pass. n_threads: Number of threads for the flow integration the candidates come from. - pbar_init: Initialize each progress stage with its total and description. - pbar_update: Advance progress after deriving prompts and each prediction batch. Returns: The proposals, to be passed to `select`. Their layout is an implementation detail of the @@ -1479,10 +1459,8 @@ def propose( if self._prediction.ndim == 4: raise ValueError("Proposals can only be reused for an image, because a volume gates its propagation.") - if pbar_init is not None: - pbar_init(1, "APG: deriving prompts") prompts = derive_point_prompts( - self._prediction[0], self._prediction[1:], model_type=self._model_type, + self._prediction[0], self._prediction[1:4], model_type=self._model_type, candidate_threshold=candidate_threshold, foreground_threshold=foreground_threshold, n_iter=n_iter, dt=dt, sigma=sigma, min_candidate_size=min_candidate_size, n_threads=n_threads, ) @@ -1490,11 +1468,7 @@ def propose( pbar_update(1) if prompts is None: return [] - if pbar_init is not None: - pbar_init((len(prompts["points"]) + batch_size - 1) // batch_size, "APG: prompting batches") - return self._apply( - prompts, multimasking=multimasking, batch_size=batch_size, pbar_update=pbar_update, - ) + return self._apply(prompts, multimasking=multimasking, batch_size=batch_size) def select( self, @@ -1561,11 +1535,9 @@ def _region_box(self, key) -> tuple: def _set_region(self, key) -> None: """Point the predictor at the region. Its image is already set for a single one.""" - def _apply(self, prompts: dict, multimasking: bool, batch_size: int, pbar_update=None) -> list: + def _apply(self, prompts: dict, multimasking: bool, batch_size: int) -> list: """Turn the prompts into mask proposals.""" - return self._apply_prompts( - self._predictor, prompts, multimasking=multimasking, batch_size=batch_size, pbar_update=pbar_update, - ) + return self._apply_prompts(self._predictor, prompts, multimasking=multimasking, batch_size=batch_size) def _merge( self, proposals: list, shape: tuple, score_threshold: float, max_overlap: float, min_size: int, @@ -1806,9 +1778,7 @@ def _predict_prompt_batch( combined = (scores.float() * stability.float()).cpu().numpy() return [(mask, float(score)) for mask, score in zip(masks, combined)] - def _apply_prompts( - self, predictor, prompts, multimasking: bool, batch_size: int, pbar_update=None, - ) -> List[Dict[str, Any]]: + def _apply_prompts(self, predictor, prompts, multimasking: bool, batch_size: int) -> List[Dict[str, Any]]: """Prompt the interactive branch in batches, returning records for the merge. Takes the predictor rather than reading `self._predictor`, so the volumetric scoring can hand @@ -1860,14 +1830,11 @@ def _apply_prompts( # The prompt as (x, y); the refinement groups the first round's prompts by it. "point": (float(batch_points[offset, 0, 0]), float(batch_points[offset, 0, 1])), }) - if pbar_update is not None: - pbar_update(1) return records def _score_candidates( self, prompts: dict, multimasking: bool, batch_size: int, score_threshold: float, max_overlap: float, components: Optional[tuple] = None, refinement_kwargs: Optional[dict] = None, - pbar_init=None, pbar_update=None, ) -> List[dict]: """Prompt every candidate in 2d on its anchor slice, and keep the strong, non-duplicate ones. @@ -1938,8 +1905,7 @@ def finish(candidate, record): "frame": int(frame), "segmentation": segmentation, "records": records, "matches": matches, "points": points[indices][:, 0, :], } - with autocast(predictor.device): - refined = self._refine_anchors(context, components, refinement_kwargs, batch_size) + refined = self._refine_anchors(context, components, refinement_kwargs, batch_size) return [ finish(candidate, records[record_index]) for candidate, record_index in zip(refined, matches.values()) diff --git a/micro_sam/v2/batched_inference.py b/micro_sam/v2/batched_inference.py index a727f5ab9..7331d39b8 100644 --- a/micro_sam/v2/batched_inference.py +++ b/micro_sam/v2/batched_inference.py @@ -1314,6 +1314,11 @@ def _resolve_z_blocking(z_block: Optional[int], z_halo: Optional[int]) -> Tuple[ return z_block, z_halo +def _n_output_channels(model) -> int: + """The decoder's output channel count: 4 (foreground and three distances) unless the model says otherwise.""" + return int(getattr(model, "out_channels", 4)) + + def _decode_volume_embeddings( model: torch.nn.Module, image_embeddings: Dict, @@ -1362,7 +1367,7 @@ def _decode_volume_embeddings( z_block, z_halo = _resolve_z_blocking(z_block, z_halo) original_size = tuple(int(value) for value in np.asarray(image_embeddings["original_size"]).reshape(-1)[:2]) - output = np.zeros((4, n_slices, *original_size), dtype="float32") + output = np.zeros((_n_output_channels(model), n_slices, *original_size), dtype="float32") jobs = [] for z0 in range(0, n_slices, z_block): z1 = min(z0 + z_block, n_slices) @@ -1431,7 +1436,7 @@ def _decode_tiled_2d_embeddings( The stitched decoder predictions, shape (4, Y, X): foreground and the three distance channels. """ features, shape, halo, tiling = _tiled_metadata(image_embeddings, is_3d=False) - output = np.zeros((4, *shape), dtype="float32") + output = np.zeros((_n_output_channels(model), *shape), dtype="float32") jobs = [] for tile_id in range(tiling.number_of_blocks): tile_features = features[str(tile_id)] @@ -1521,7 +1526,7 @@ def _decode_tiled_3d_embeddings( n_slices = shape[0] z_block, z_halo = _resolve_z_blocking(z_block, z_halo) jobs = _tiled_3d_jobs(features, tiling, n_slices, z_block, z_halo) - output = np.zeros((4, *shape), dtype="float32") + output = np.zeros((_n_output_channels(model), *shape), dtype="float32") if pbar_init is not None: pbar_init(tiling.number_of_blocks * n_slices, "Automatic segmentation (tiles)") @@ -1588,7 +1593,7 @@ def _decode_tiled_3d_slice( if not 0 <= index < n_slices: raise ValueError(f"The slice index must be in [0, {n_slices}), got {index}.") - output = np.zeros((4, *shape[1:]), dtype="float32") + output = np.zeros((_n_output_channels(model), *shape[1:]), dtype="float32") jobs = [] for tile_id in range(tiling.number_of_blocks): tile_features = features[str(tile_id)] diff --git a/micro_sam/v2/instance_segmentation.py b/micro_sam/v2/instance_segmentation.py index afcf0b13c..e937c7d00 100644 --- a/micro_sam/v2/instance_segmentation.py +++ b/micro_sam/v2/instance_segmentation.py @@ -731,7 +731,8 @@ def _check_decoder_width(model, initial_features): def get_unisam2_model( - checkpoint_path, device=None, encoder=_DEFAULT_MODEL, output_channels=4, peft_kwargs=None, encoder_model_type=None + checkpoint_path, device=None, encoder=_DEFAULT_MODEL, output_channels=None, peft_kwargs=None, + encoder_model_type=None, ): """Load a UniSAM2 model for automatic segmentation from a checkpoint. @@ -741,7 +742,8 @@ def get_unisam2_model( encoder: The SAM2 encoder to build the decoder on. Either the backbone name to build from scratch, e.g. 'hvit_t', or a prebuilt SAM2 image-encoder module to reuse (which avoids rebuilding / downloading the base backbone). Its weights are (re)defined by the checkpoint. - output_channels: The number of output channels (foreground + directed distances). + output_channels: The number of output channels (foreground, directed distances and optional auxiliary + channels). By default it is read off the checkpoint's output layer. peft_kwargs: The arguments for `PEFT_Sam2`. The function uses the saved arguments by default. encoder_model_type: The SAM2 model type for a prebuilt PEFT encoder. You must set this argument for modules. @@ -781,8 +783,10 @@ def get_unisam2_model( sam2_model = PEFT_Sam2(sam2_model, **peft_kwargs).sam encoder = sam2_model.image_encoder - # The decoder width is not recorded in the checkpoint, so read it off 'out_conv'. + # Neither the decoder width nor the channel count is recorded in the checkpoint, so read them off 'out_conv'. initial_features = model_state["out_conv.weight"].shape[1] + if output_channels is None: + output_channels = model_state["out_conv.weight"].shape[0] model = UniSAM2(encoder=encoder, output_channels=output_channels, initial_features=initial_features, device=device) _check_decoder_width(model, initial_features) @@ -1015,8 +1019,9 @@ def _segment_from_predictions(prediction: np.ndarray, mode: str = "sparse", **kw """Convert UniSAM2 predictions into an instance segmentation. Args: - prediction: The UniSAM2 predictions, shape (4, *spatial). Channel 0 is the foreground - probability and channels 1-3 are the directed distances. + prediction: The UniSAM2 predictions, shape (4, *spatial) or (5, *spatial). Channel 0 is the foreground + probability, channels 1-3 are the directed distances and the optional channel 4 is the contact + probability, which the sparse mode forwards as 'contact'. mode: The segmentation mode. 'sparse' uses flow-based segmentation (LM data, 2d and 3d), 'dense' uses multicut-based segmentation (EM data, 2d and 3d). kwargs: Additional parameters forwarded to the postprocessing function @@ -1039,7 +1044,9 @@ def _segment_from_predictions(prediction: np.ndarray, mode: str = "sparse", **kw else: seg = run_multicut(boundary_map, distances, **kwargs) else: - seg = flow_instance_segmentation(foreground, prediction[1:], **kwargs) + if prediction.shape[0] > 4: + kwargs = {"contact": prediction[4], **kwargs} + seg = flow_instance_segmentation(foreground, prediction[1:4], **kwargs) return seg.astype("uint32") @@ -1125,12 +1132,13 @@ def _predict_probe(this_model, inputs): desc = "Automatic segmentation (volume)" if is_3d else "Automatic segmentation" pbar_init(n_blocks, desc) + n_channels = int(getattr(self._model, "out_channels", 4)) if is_3d: input_ = raw[np.newaxis].astype("float32") - output = np.zeros((4, *raw.shape), dtype="float32") + output = np.zeros((n_channels, *raw.shape), dtype="float32") else: input_ = raw[np.newaxis, np.newaxis].astype("float32") - output = np.zeros((4, 1, *raw.shape), dtype="float32") + output = np.zeros((n_channels, 1, *raw.shape), dtype="float32") img_size = getattr(getattr(self._model, "encoder", None), "img_size", 1024) resize_model = ResizeLongestSideWrapper(self._model, img_size) diff --git a/micro_sam/v2/loss/directed_distance_based.py b/micro_sam/v2/loss/directed_distance_based.py index e56cdabbb..17c395b76 100644 --- a/micro_sam/v2/loss/directed_distance_based.py +++ b/micro_sam/v2/loss/directed_distance_based.py @@ -1,3 +1,5 @@ +from typing import Optional + import torch import torch.nn as nn import torch.nn.functional as F @@ -18,47 +20,87 @@ def _masked_mse(prediction: torch.Tensor, target: torch.Tensor, mask: torch.Tens return (error.sum(dims) / mask.sum(dims).clamp_min(1.0)).mean() +def _weighted_bce( + prediction: torch.Tensor, target: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6, +) -> torch.Tensor: + """Binary cross entropy on probabilities, weighted per pixel and normalized per sample by the weight sum. + + The probabilities are cast to float32 before the logarithm: in bfloat16 a value close to one rounds to + exactly one and the log of its complement would be infinite. + """ + prediction = prediction.float().clamp(eps, 1.0 - eps) + target = target.float() + error = -(target * torch.log(prediction) + (1.0 - target) * torch.log1p(-prediction)) * weight + dims = tuple(range(1, error.ndim)) + return (error.sum(dims) / weight.sum(dims).clamp_min(1.0)).mean() + + +def boundary_band(foreground: torch.Tensor, radius: int) -> torch.Tensor: + """The pixels within 'radius' of a transition between foreground and background. + + Computed in-plane with a max and a min pooling of the binary foreground, so the band has the same width + on either side of every object boundary. + + Args: + foreground: The binary foreground target, shape (B, 1, Z, Y, X). + radius: The half width of the band in pixels. + + Returns: + The band as a float tensor of the foreground's shape and dtype (1 inside the band). + """ + kernel, padding = (1, 2 * radius + 1, 2 * radius + 1), (0, radius, radius) + upper = F.max_pool3d(foreground, kernel, stride=1, padding=padding) + lower = -F.max_pool3d(-foreground, kernel, stride=1, padding=padding) + return (upper != lower).to(foreground.dtype) + + class DirectedDistanceLoss(nn.Module): """Loss for directed distance based instance segmentation. - The inputs contain foreground, three directed distances, and an optional fifth boundary channel. - The boundary loss combines Dice and binary cross entropy (BCE) with ``boundary_dice_weight``. + Expects input and targets with four channels, foreground and three distance channels (in z, y and x), + plus a fifth contact channel when ``contact=True``. The foreground is trained with ``foreground_loss`` + (Dice by default); ``boundary_weight`` adds a per-pixel binary cross entropy whose weight rises to + ``1 + boundary_weight`` within ``boundary_radius`` pixels of every object boundary, which calibrates the + predicted extent to the annotated boundary instead of rewarding a wide, soft foreground. The distances are + trained with a masked mean squared error, the contact channel with Dice plus binary cross entropy. Args: - mask_distances_in_bg: The flag to exclude background voxels from the distance loss. - foreground_loss: The loss for foreground predictions and targets. - with_boundaries: The flag for a fifth boundary channel in the inputs and targets. - boundary_dice_weight: The Dice weight in the boundary loss. One selects Dice only. - Zero selects BCE only. Values between zero and one mix the two losses. + mask_distances_in_bg: Whether to mask the loss for distance predictions in the background. + foreground_loss: The loss for comparing foreground predictions and target. Dice by default. + contact: Whether the fifth channel holds the contact (touching boundary) probability. + contact_weight: The weight of the contact term. + boundary_weight: The extra weight of the foreground cross entropy in the boundary band. None disables + the cross entropy term altogether (the default, Dice only). + boundary_radius: The half width of the boundary band in pixels. """ def __init__( self, mask_distances_in_bg: bool = True, - foreground_loss: nn.Module = DiceLoss(), - with_boundaries: bool = False, - boundary_dice_weight: float = 1.0, + foreground_loss: Optional[nn.Module] = None, + contact: bool = False, + contact_weight: float = 1.0, + boundary_weight: Optional[float] = None, + boundary_radius: int = 2, ) -> None: super().__init__() - if not 0.0 <= boundary_dice_weight <= 1.0: - raise ValueError(f"boundary_dice_weight must be between zero and one, got {boundary_dice_weight}.") - - self.foreground_loss = foreground_loss + self.foreground_loss = DiceLoss() if foreground_loss is None else foreground_loss self.mask_distances_in_bg = mask_distances_in_bg - self.with_boundaries = with_boundaries - self.boundary_dice_weight = boundary_dice_weight - self.boundary_loss = DiceLoss() if with_boundaries else None + self.contact = contact + self.contact_weight = contact_weight + self.boundary_weight = boundary_weight + self.boundary_radius = boundary_radius + self.contact_loss = DiceLoss() if contact else None self.init_kwargs = { - "mask_distances_in_bg": mask_distances_in_bg, - "with_boundaries": with_boundaries, - "boundary_dice_weight": boundary_dice_weight, + "mask_distances_in_bg": mask_distances_in_bg, "contact": contact, "contact_weight": contact_weight, + "boundary_weight": boundary_weight, "boundary_radius": boundary_radius, } @property def n_channels(self) -> int: - """Return the number of prediction and target channels.""" - return 4 + int(self.with_boundaries) + """The number of prediction and target channels the loss expects.""" + return 4 + int(self.contact) def forward(self, input_: torch.Tensor, target: torch.Tensor) -> torch.Tensor: assert input_.shape == target.shape, (input_.shape, target.shape) @@ -70,12 +112,10 @@ def forward(self, input_: torch.Tensor, target: torch.Tensor) -> torch.Tensor: # and treats it differently (sums over it independently). # This will lead to a very large dice loss that dominates over everything else. fg_input, fg_target = input_[:, 0:1], target[:, 0:1] - - # Voxels without ground truth carry FOREGROUND_IGNORE_VALUE (-1) in the foreground channel. Zeroing both - # tensors there is a Dice loss mask, and the zeroed fg_target also drops them from the distance masks below. - valid = (fg_target != FOREGROUND_IGNORE_VALUE).to(fg_target.dtype) - fg_target = fg_target * valid - fg_loss = self.foreground_loss(fg_input * valid, fg_target) + fg_loss = self.foreground_loss(fg_input, fg_target) + if self.boundary_weight is not None: + weight = 1.0 + self.boundary_weight * boundary_band(fg_target, self.boundary_radius) + fg_loss = fg_loss + _weighted_bce(fg_input, fg_target, weight) # Check whether the input is 2d or not. # For 2d inputs, we avoid computing gradients for masked (pseudo) z-distances. @@ -93,21 +133,11 @@ def forward(self, input_: torch.Tensor, target: torch.Tensor) -> torch.Tensor: xdist_loss = _masked_mse(input_[:, 3:4], target[:, 3:4], yx_mask) overall_loss = fg_loss + zdist_loss + ydist_loss + xdist_loss - if self.with_boundaries: - boundary_input, boundary_target = input_[:, 4:5], target[:, 4:5] - dice_loss = self.boundary_loss(boundary_input * valid, boundary_target * valid) - if self.boundary_dice_weight == 1.0: - boundary_loss = dice_loss - else: - # CUDA autocast prohibits BCE on probabilities. - with torch.autocast(device_type=boundary_input.device.type, enabled=False): - # Clamp rounded sigmoid outputs to keep them away from zero and one. - probability = boundary_input.float().clamp(1e-6, 1.0 - 1e-6) - error = F.binary_cross_entropy(probability, boundary_target.float(), reduction="none") - # Normalize over valid voxels per sample, as for the distance terms. - boundary_valid = valid.float() - dims = tuple(range(1, error.ndim)) - bce_loss = ((error * boundary_valid).sum(dims) / boundary_valid.sum(dims).clamp_min(1.0)).mean() - boundary_loss = self.boundary_dice_weight * dice_loss + (1.0 - self.boundary_dice_weight) * bce_loss - overall_loss = overall_loss + boundary_loss + + if self.contact: + contact_input, contact_target = input_[:, 4:5], target[:, 4:5] + contact_loss = self.contact_loss(contact_input, contact_target) + contact_loss = contact_loss + _weighted_bce(contact_input, contact_target, torch.ones_like(contact_target)) + overall_loss = overall_loss + self.contact_weight * contact_loss + return overall_loss diff --git a/micro_sam/v2/models/util.py b/micro_sam/v2/models/util.py index 91026da7c..eeb596e27 100644 --- a/micro_sam/v2/models/util.py +++ b/micro_sam/v2/models/util.py @@ -10,7 +10,9 @@ class CustomActivation(nn.Module): - """Apply sigmoid to foreground and optional auxiliary channels, and tanh to distances.""" + """Applies 'Sigmoid' to channel 0 (the foreground) and to every channel from 4 on (auxiliary probabilities + such as the contact channel), and 'Tanh' to channels 1-3 (the directed distances). + """ def forward(self, x: torch.Tensor) -> torch.Tensor: return torch.cat([torch.sigmoid(x[:, :1]), torch.tanh(x[:, 1:4]), torch.sigmoid(x[:, 4:])], dim=1) @@ -60,36 +62,3 @@ def __init__( **kwargs ) self.to(device) - - -class SemanticSAM2(UNETR3D): - """UNETR-based model for semantic (2d + 3d) segmentation. - - The model has no final activation, so it returns the raw class logits that the semantic losses expect. - """ - def __init__( - self, - encoder: Union[str, nn.Module] = "hvit_t", - num_classes: int = 3, - img_size: int = 1024, - device: Optional[Union[str, torch.device]] = None, - **kwargs, - ): - device = torch.device("cpu") if device is None else torch.device(get_device(device)) - - # One encoder type for both callers, so the weights land under the same keys either way. - if isinstance(encoder, str): - encoder = get_sam2_model(model_type=encoder, input_type="images", device=device).image_encoder - - super().__init__( - img_size=img_size, - backbone="sam2", - encoder=SAM2EncoderAdapter(encoder, img_size=img_size), - final_activation=None, - out_channels=num_classes, - use_sam_stats=True, - embed_dim=256, - use_strip_pooling=True, - **kwargs - ) - self.to(device) diff --git a/micro_sam/v2/postprocessing.py b/micro_sam/v2/postprocessing.py index bce579f4c..d8cb30d4c 100644 --- a/micro_sam/v2/postprocessing.py +++ b/micro_sam/v2/postprocessing.py @@ -23,49 +23,71 @@ # Per (model_type, mode) defaults from the registry parameter search: the best-average-rank # combination across every dataset that shares that mode's grid, computed separately for each of the # 4 registry backbones. +# 'boundary_magnitude_max' is the instance filter of `flow_instance_segmentation`; None keeps it off. +# 'seed_floor' lowers the height map under the seeds before the watershed ('none', 'zero' or 'ring'). +# 'sparse_volume' holds the keys whose default differs for a volume (a size floor counts voxels, not +# pixels); it is layered over 'sparse' by `default_postprocessing(..., ndim=3)`. +# +# The hvit_t entry is the result of the 2026-09 AIS optimization on the joint/v4 geodesic checkpoint +# (finetuning/v2/evaluation/optimization/notes/AIS_V4_OPTIMIZATION.md). Images: against the registry values +# (min_size 100, sigma 0.5, no filter) the wider density smoothing, the ground-truth-like size floor and the +# boundary filter gain +2.4 % balanced mSA on eleven 2d development datasets (9 up, worst -0.8 %) and +# +4.3 % on the 2d holdout. Volumes keep the registry values and add the filter only (+4.7 % / +9.7 % on +# the 3d tuning crops, +3.4 % on the seven test-only 3d datasets, none down); the stronger volume settings +# that won on the tuning crops did not carry over to the test datasets. DEFAULT_POSTPROCESSING = { "hvit_t": { "sparse": { - "foreground_threshold": 0.5, "density_threshold": 10.0, "min_size": 100, - "sigma": 0.5, "n_iter": 50, "dt": 0.5, "foreground_weight": 0.5, + "foreground_threshold": 0.5, "density_threshold": 10.0, "min_size": 50, + "sigma": 1.0, "n_iter": 50, "dt": 0.5, "foreground_weight": 0.5, "boundary_magnitude_max": 0.4, + "seed_floor": "none", }, + "sparse_volume": {"min_size": 100, "sigma": 0.5}, "dense": {"beta": 0.5, "density_threshold": 5.0, "sigma": 0.5, "n_iter": 50, "dt": 0.5}, }, "hvit_s": { "sparse": { "foreground_threshold": 0.5, "density_threshold": 20.0, "min_size": 100, - "sigma": 0.25, "n_iter": 50, "dt": 0.5, "foreground_weight": 0.75, + "sigma": 0.25, "n_iter": 50, "dt": 0.5, "foreground_weight": 0.75, "boundary_magnitude_max": None, + "seed_floor": "none", }, + "sparse_volume": {}, "dense": {"beta": 0.5, "density_threshold": 3.0, "sigma": 0.5, "n_iter": 25, "dt": 0.5}, }, "hvit_b": { "sparse": { "foreground_threshold": 0.5, "density_threshold": 20.0, "min_size": 100, - "sigma": 0.25, "n_iter": 50, "dt": 0.5, "foreground_weight": 0.65, + "sigma": 0.25, "n_iter": 50, "dt": 0.5, "foreground_weight": 0.65, "boundary_magnitude_max": None, + "seed_floor": "none", }, + "sparse_volume": {}, "dense": {"beta": 0.5, "density_threshold": 5.0, "sigma": 0.5, "n_iter": 50, "dt": 0.5}, }, "hvit_l": { "sparse": { "foreground_threshold": 0.4, "density_threshold": 10.0, "min_size": 50, - "sigma": 0.5, "n_iter": 50, "dt": 0.25, "foreground_weight": 0.65, + "sigma": 0.5, "n_iter": 50, "dt": 0.25, "foreground_weight": 0.65, "boundary_magnitude_max": None, + "seed_floor": "none", }, + "sparse_volume": {}, "dense": {"beta": 0.5, "density_threshold": 5.0, "sigma": 1.0, "n_iter": 50, "dt": 0.5}, }, } -def default_postprocessing(model_type: str = DEFAULT_MODEL, mode: str = "sparse") -> dict: - """The default postprocessing parameters for one model type and mode. +def default_postprocessing(model_type: str = DEFAULT_MODEL, mode: str = "sparse", ndim: int = 2) -> dict: + """The default postprocessing parameters for one model type, mode and dimensionality. Args: model_type: The SAM2 backbone, e.g. 'hvit_t', or a finetuned model built on one, e.g. 'hvit_t_cells' (only the backbone prefix is used to look up the table). Must be one of the 4 registry backbones. mode: 'sparse' (`flow_instance_segmentation`) or 'dense' (`run_multicut`). + ndim: The number of spatial dimensions of the data, 2 or 3. A volume takes the + '_volume' overrides of the table on top of the mode's defaults. Returns: - The default parameter dict for that model type and mode. + The default parameter dict for that model type, mode and dimensionality. """ backbone = model_type[:6] if backbone not in DEFAULT_POSTPROCESSING: @@ -73,7 +95,11 @@ def default_postprocessing(model_type: str = DEFAULT_MODEL, mode: str = "sparse" f"No default postprocessing parameters for model type '{model_type}'. " f"Choose one built on a backbone in {sorted(DEFAULT_POSTPROCESSING)}." ) - return DEFAULT_POSTPROCESSING[backbone][mode] + table = DEFAULT_POSTPROCESSING[backbone] + defaults = dict(table[mode]) + if ndim == 3: + defaults.update(table.get(f"{mode}_volume", {})) + return defaults def _compute_flow_density( @@ -135,6 +161,97 @@ def watershed_heightmap( return np.ascontiguousarray(hmap, dtype="float32") +def lower_height_under_seeds(heightmap: np.ndarray, seeds: np.ndarray, mode: str) -> np.ndarray: + """Lower the height map under the seeds so that a seed's own height does not hold its front back. + + The watershed floods monotonically: a front never drops below the height it started from. Every proper + seed sits on a peak of the inverted-magnitude height map (the predicted magnitude dips at the object's + centre), so a seed whose centre dip is deeper than the contact dip to its neighbour loses the object to + the neighbour's front. 'zero' sets the height under every seed to zero; 'ring' sets it to the minimum + height of a ring of 2-3 pixels around the seed, so that a seed inherits the level of its own basin + and a seed on a high plateau (a spurious one) keeps a high floor. + + Args: + heightmap: The watershed height map, shape (*spatial). + seeds: The seed components, integer labels, same shape. + mode: 'none' (return the height map unchanged), 'zero' or 'ring'. + + Returns: + The height map with the seeds lowered, float32 and C-contiguous. + """ + if mode == "none": + return heightmap + out = np.array(heightmap, dtype="float32", copy=True) + if mode == "zero": + out[seeds != 0] = 0.0 + return np.ascontiguousarray(out) + if mode != "ring": + raise ValueError(f"Unknown seed floor '{mode}'; expected 'none', 'zero' or 'ring'.") + from scipy.ndimage import grey_dilation, minimum as labelled_minimum + + inner = grey_dilation(seeds, size=(3,) * seeds.ndim) + outer = grey_dilation(seeds, size=(7,) * seeds.ndim) + ring = np.where((outer != 0) & (inner == 0), outer, 0) + ids = np.unique(ring) + ids = ids[ids != 0] + if len(ids) == 0: + return np.ascontiguousarray(out) + floors = np.zeros(int(seeds.max()) + 1, dtype="float32") + floors[ids] = labelled_minimum(heightmap, labels=ring, index=ids) + core = inner != 0 + out[core] = np.minimum(out[core], floors[inner[core]]) + return np.ascontiguousarray(out) + + +def drop_instances_without_boundary_dip( + segmentation: np.ndarray, directed_distances: np.ndarray, max_median: float +) -> np.ndarray: + """Drop the instances whose boundary shows no dip of the distance magnitude. + + The magnitude of the directed distances falls to (almost) zero along the boundary of every object + the decoder recognised, because the distance to the object's boundary is what it predicts. A false + foreground region carries no such structure: its boundary runs through the decoder's background + output (magnitude about one) or through the interior of a field that belongs to something else. An + instance whose median boundary magnitude exceeds 'max_median' is therefore removed. The rule is + label-free and scale-free, and a real object passes it at any size. + + Args: + segmentation: The instance segmentation, shape (*spatial). + directed_distances: Distance channels stacked along axis 0, shape (ndim, *spatial). + max_median: Instances whose median boundary magnitude exceeds this value are dropped. + + Returns: + The filtered segmentation, same dtype and shape. + """ + # The inner boundary: instance pixels with an axis neighbour of another label (or background). + boundary = np.zeros(segmentation.shape, dtype=bool) + for axis in range(segmentation.ndim): + lower = [slice(None)] * segmentation.ndim + upper = [slice(None)] * segmentation.ndim + lower[axis], upper[axis] = slice(None, -1), slice(1, None) + differs = segmentation[tuple(lower)] != segmentation[tuple(upper)] + boundary[tuple(lower)] |= differs + boundary[tuple(upper)] |= differs + boundary &= segmentation != 0 + if not boundary.any(): + return segmentation + labels = segmentation[boundary] + values = np.linalg.norm(directed_distances[(slice(None),) + np.nonzero(boundary)], axis=0) + # One sort over the boundary pixels gives every instance's median (the mean of the two middle values + # for an even count, like `scipy.ndimage.median`). + order = np.lexsort((values, labels)) + labels, values = labels[order], values[order] + starts = np.flatnonzero(np.r_[True, labels[1:] != labels[:-1]]) + counts = np.diff(np.r_[starts, len(labels)]) + upper_middle = values[starts + counts // 2] + lower_middle = values[starts + (counts - 1) // 2] + medians = 0.5 * (upper_middle + lower_middle) + drop = labels[starts][medians > max_median] + if drop.size == 0: + return segmentation + return np.where(np.isin(segmentation, drop), 0, segmentation).astype(segmentation.dtype) + + def flow_instance_segmentation( foreground: np.ndarray, directed_distances: np.ndarray, @@ -148,6 +265,11 @@ def flow_instance_segmentation( min_size: Optional[int] = None, foreground_weight: Optional[float] = None, n_threads: int = 8, + boundary_magnitude_max: Optional[float] = None, + seed_floor: Optional[str] = None, + contact: Optional[np.ndarray] = None, + contact_weight: Optional[float] = None, + contact_mask_threshold: Optional[float] = None, ) -> np.ndarray: """Instance segmentation from directed-distance predictions via flow following. @@ -156,8 +278,9 @@ def flow_instance_segmentation( watershed. Works for both 2D and 3D inputs. If 3 distance channels are supplied for a 2D foreground map the leading - z-channel is automatically dropped, so you can always pass ``out[1:]`` - regardless of dimensionality. + z-channel is automatically dropped, so you can always pass the three distance + channels ``out[1:4]`` regardless of dimensionality. Any other channel count raises, + so that an auxiliary channel appended to the prediction is never read as a distance. Args: foreground: Foreground probability map, shape (Y, X) or (Z, Y, X). @@ -175,13 +298,29 @@ def flow_instance_segmentation( foreground_weight: Weight of the foreground term in the watershed heightmap, see `watershed_heightmap`. n_threads: Number of threads for the flow computation. + boundary_magnitude_max: Drop instances whose median boundary magnitude exceeds this value, see + `drop_instances_without_boundary_dip`. None takes the per-model default, which may itself be + None (no filtering); pass ``float("inf")`` to disable a default filter explicitly. + seed_floor: How the height map is lowered under the seeds before the watershed, see + `lower_height_under_seeds`. None takes the per-model default. + contact: The predicted contact (touching boundary) probability, same shape as the foreground, from a + decoder with a fifth output channel. Only used through the two keywords below. + contact_weight: Adds ``contact_weight * contact`` to the watershed height map, so that the fronts of + two touching objects meet on the predicted contact line. None or 0 leaves the height map unchanged. + contact_mask_threshold: Excludes the pixels with ``contact > threshold`` from the first seeded watershed + and assigns them afterwards by flooding from the resulting instances, so that no instance grows + across a contact line. None disables the exclusion. Returns: Instance segmentation, uint32 array, same spatial shape as foreground. """ - defaults = default_postprocessing(model_type, "sparse") + defaults = default_postprocessing(model_type, "sparse", ndim=foreground.ndim) if foreground_threshold is None: foreground_threshold = defaults["foreground_threshold"] + if boundary_magnitude_max is None: + boundary_magnitude_max = defaults.get("boundary_magnitude_max") + if seed_floor is None: + seed_floor = defaults.get("seed_floor", "none") if n_iter is None: n_iter = defaults["n_iter"] if dt is None: @@ -196,11 +335,17 @@ def flow_instance_segmentation( foreground_weight = defaults["foreground_weight"] ndim = foreground.ndim - if directed_distances.shape[0] > ndim: - directed_distances = directed_distances[-ndim:] - assert directed_distances.shape[0] == ndim, ( - f"Expected {ndim} distance channels, got {directed_distances.shape[0]}." - ) + if directed_distances.shape[0] == 3 and ndim == 2: + directed_distances = directed_distances[1:] # Drop the (pseudo) z channel of a 2d prediction. + if directed_distances.shape[0] != ndim: + raise ValueError( + f"Expected {ndim} distance channels (or 3 for 2d input), got {directed_distances.shape[0]}. Pass the " + "three distance channels 'prediction[1:4]'; an auxiliary channel goes into 'contact'." + ) + if contact is None and (contact_weight is not None or contact_mask_threshold is not None): + raise ValueError("'contact_weight' and 'contact_mask_threshold' need the predicted contact map 'contact'.") + if contact is not None and contact.shape != foreground.shape: + raise ValueError(f"The contact map {contact.shape} must have the shape of the foreground {foreground.shape}.") fg_mask = foreground > foreground_threshold @@ -210,7 +355,17 @@ def flow_instance_segmentation( seeds = label(density > density_threshold) hmap = watershed_heightmap(foreground, directed_distances, foreground_weight) - seg = watershed(hmap, markers=seeds, mask=fg_mask) + if contact is not None and contact_weight is not None and contact_weight != 0: + # The contact line becomes a ridge, so the fronts of two touching objects meet on it. + hmap = np.ascontiguousarray(hmap + np.float32(contact_weight) * np.clip(contact, 0, 1), dtype="float32") + hmap = lower_height_under_seeds(hmap, seeds, seed_floor) + if contact is not None and contact_mask_threshold is not None: + # Flood everything but the contact pixels first, then let the instances claim the contact pixels. + open_mask = fg_mask & ~(contact > contact_mask_threshold) + first = watershed(hmap, markers=np.where(open_mask, seeds, 0).astype(seeds.dtype), mask=open_mask) + seg = watershed(hmap, markers=first, mask=fg_mask) + else: + seg = watershed(hmap, markers=seeds, mask=fg_mask) if min_size > 0: ids, sizes = np.unique(seg, return_counts=True) @@ -218,6 +373,10 @@ def flow_instance_segmentation( seg[np.isin(seg, discard)] = 0 seg = watershed(hmap, markers=seg, mask=fg_mask) + # After the size filter, so that a dropped region is not refilled by its neighbours. + if boundary_magnitude_max is not None and np.isfinite(boundary_magnitude_max): + seg = drop_instances_without_boundary_dip(seg, directed_distances, boundary_magnitude_max) + return seg.astype("uint32") diff --git a/micro_sam/v2/transforms/labels.py b/micro_sam/v2/transforms/labels.py index e7c3391ea..5f9adbd3e 100644 --- a/micro_sam/v2/transforms/labels.py +++ b/micro_sam/v2/transforms/labels.py @@ -3,8 +3,7 @@ import numpy as np -from scipy.ndimage import binary_dilation - +from scipy.ndimage import binary_dilation, maximum_filter, minimum_filter from skimage.measure import regionprops from skimage.segmentation import find_boundaries @@ -277,7 +276,7 @@ def _joint_em_cell_label_trafo(y, label_trafo, ignore_label=None): """EM label transform for joint training - keeps instance IDs as channel 0. Like :func:`_em_cell_label_trafo` but returns - ``[instance_ids, expected_fg, d_x, d_y, d_z]`` (5 channels) instead of + ``[instance_ids, expected_fg, d_z, d_y, d_x]`` (5 channels) instead of dropping the instance channel. ``label_trafo`` must produce a 5-channel array (i.e. be a :class:`_JointLabelTransform` / ``instances=True``). """ @@ -294,30 +293,74 @@ def _joint_em_cell_label_trafo(y, label_trafo, ignore_label=None): return np.concatenate([instances[None], expected_fg[None], y[2:]], axis=0) -def object_boundaries(labels: np.ndarray) -> np.ndarray: - """Return a dilated mask of all object boundaries. +def touching_boundaries(labels: np.ndarray, radius: int = 1, dilation: int = 1) -> np.ndarray: + """The contact lines between touching objects. + + A pixel is a contact pixel if its ``(2 * radius + 1)`` neighbourhood holds two different non-zero labels. + Directly touching objects therefore contribute their two facing boundary lines, and a one pixel annotation + gap between two objects contributes the gap itself. Object interiors and background away from any pair of + objects are never contacts. The mask is then dilated by ``dilation`` pixels, so that the target is a few + pixels wide and learnable. + + Args: + labels: The instance segmentation, 2d or 3d, any integer dtype. + radius: The neighbourhood radius in pixels. + dilation: The number of binary dilation passes applied to the contact mask. + + Returns: + The boolean contact mask with the shape of ``labels``. + """ + labels = np.asarray(labels).astype("int64") + size = 2 * radius + 1 + highest = maximum_filter(labels, size=size, mode="nearest") + # Background must not count as a label: send it above every id, so the minimum picks the smallest object id. + sentinel = labels.max() + 1 + lowest = minimum_filter(np.where(labels > 0, labels, sentinel), size=size, mode="nearest") + # A neighbourhood with at least one object has a real minimum id; two different ids give lowest < highest. + contact = (highest > 0) & (lowest != highest) + if dilation > 0 and contact.any(): + contact = binary_dilation(contact, iterations=dilation) + return contact + + +def object_boundaries(labels: np.ndarray, dilation: int = 1) -> np.ndarray: + """The inner boundaries of every object, to a neighbour and to the background alike. + + The classical boundary target: ``find_boundaries(mode="inner")`` dilated by ``dilation`` pixels, so it is + defined identically on every object (a few percent of the pixels rather than the sub-percent contact class) + and coincides with the zero level set of the geodesic distance channels. + + Args: + labels: The instance segmentation, 2d or 3d, any integer dtype. + dilation: The number of binary dilation passes applied to the boundary mask. - The transform dilates each inner boundary once. The target includes isolated objects and objects that touch. + Returns: + The boolean boundary mask with the shape of ``labels``. """ - boundary = find_boundaries(np.asarray(labels), mode="inner") - if boundary.any(): - boundary = binary_dilation(boundary, iterations=1) + labels = np.asarray(labels).astype("int64") + boundary = find_boundaries(labels, mode="inner") + if dilation > 0 and boundary.any(): + boundary = binary_dilation(boundary, iterations=dilation) return boundary class DirectedPerObjectBoundaryDistanceTransform: - """Compute directed-distance targets with an optional boundary channel. + """Per object directed distances with optional foreground, instance and contact channels. - The channel layout is ``[instance_ids?, foreground?, d_z, d_y, d_x, boundaries?]``. + Output layout along the channel axis: ``[instance_ids?, foreground?, d_z, d_y, d_x, contact?]``, i.e. the + optional instance channel comes first, the foreground mask second, then the three distance channels in axis + order and finally the optional contact channel (see :func:`touching_boundaries`). Args: - min_size: The minimum object size. The transform removes smaller objects. - foreground: The flag to prepend the binary foreground mask. - instances: The flag to prepend the instance IDs. - apply_label: The flag to relabel the input with connected components. + min_size: Objects smaller than this are removed before the transform. + foreground: Whether to prepend the binary foreground mask. + instances: Whether to prepend the instance ids (joint training). + apply_label: Whether to relabel the input with connected components. sampling: The voxel spacing for anisotropic data. - with_boundaries: The flag to append the full object-boundary mask. - n_threads: The number of threads for distance computation across objects. + contact: Whether to append the contact channel, the touching boundaries between objects. + contact_dilation: The dilation of the contact lines in pixels, see :func:`touching_boundaries`. + contact_mode: What the contact channel holds: "touching" (the boundaries between touching objects, + :func:`touching_boundaries`) or "all" (the inner boundary of every object, :func:`object_boundaries`). """ eps = 1e-7 @@ -328,9 +371,12 @@ def __init__( instances: bool = False, apply_label: bool = True, sampling: Optional[Tuple[float, ...]] = None, - with_boundaries: bool = False, - n_threads: int = 1, + contact: bool = False, + contact_dilation: int = 1, + contact_mode: str = "touching", ): + if contact_mode not in ("touching", "all"): + raise ValueError(f"Unknown contact_mode '{contact_mode}'; expected 'touching' or 'all'.") self.min_size = min_size self.n_threads = n_threads self.distance_fill_value = 1 @@ -338,7 +384,9 @@ def __init__( self.instances = instances self.apply_label = apply_label self.sampling = sampling - self.with_boundaries = with_boundaries + self.contact = contact + self.contact_dilation = contact_dilation + self.contact_mode = contact_mode def compute_normalized_directed_distances(self, labels, label_id, boundaries, bb, distances): """@private @@ -425,9 +473,13 @@ def compute(prop): to_channel_first = (ndim,) + tuple(range(ndim)) distances = distances.transpose(to_channel_first) - if self.with_boundaries: - boundaries = object_boundaries(labels).astype("float32") - distances = np.concatenate([distances, boundaries[None]], axis=0) + # Append the contact channel (touching boundaries) after the distances if specified. + if self.contact: + if self.contact_mode == "all": + contact = object_boundaries(labels, dilation=self.contact_dilation).astype("float32") + else: + contact = touching_boundaries(labels, radius=1, dilation=self.contact_dilation).astype("float32") + distances = np.concatenate([distances, contact[None]], axis=0) # Add the foreground mask as first channel if specified. if self.foreground: @@ -519,9 +571,9 @@ def compute_normalized_directed_distances(self, labels, label_id, boundaries, bb class _JointLabelTransform(DirectedPerObjectBoundaryDistanceTransform): """Distance transform for joint interactive + automatic training. - This transform sets ``instances=True`` by default. - The output layout is ``[instance_ids, foreground_mask, d_z, d_y, d_x, boundaries?]``. - Set ``with_boundaries=True`` to append the sixth channel. + Identical to :class:`DirectedPerObjectBoundaryDistanceTransform` but + defaults to ``instances=True`` so the output always has 5 channels: + ``[instance_ids, foreground_mask, d_z, d_y, d_x]`` (6 with ``contact=True``). The interactive branch uses channel 0 (cast to int64 as instance IDs) and the automatic branch uses channels 1 onward. @@ -534,8 +586,11 @@ def __init__(self, instances: bool = True, **kwargs): class _JointGeodesicLabelTransform(GeodesicHybridDistanceTransform): """Geodesic hybrid distance transform for joint interactive + automatic training. - The output layout is ``[instance_ids, foreground_mask, d_z, d_y, d_x, boundaries?]``. - The directed distances come from the geodesic field around each object's center. + The :class:`GeodesicHybridDistanceTransform` counterpart of + :class:`_JointLabelTransform`: same 5-channel output + ``[instance_ids, foreground_mask, d_z, d_y, d_x]``, but the directed distances come from + the geodesic field around each object's center instead of the euclidean vector to the + nearest boundary. """ def __init__(self, instances: bool = True, **kwargs): diff --git a/test/test_ais_checkpoint_comparison.py b/test/test_ais_checkpoint_comparison.py new file mode 100644 index 000000000..9a48d1613 --- /dev/null +++ b/test/test_ais_checkpoint_comparison.py @@ -0,0 +1,91 @@ +import json +import sys +from pathlib import Path + +import numpy as np +import pandas as pd +import pytest + + +OPTIMIZATION_ROOT = Path(__file__).parents[1] / "finetuning/v2/evaluation/optimization" +sys.path.insert(0, str(OPTIMIZATION_ROOT)) + +import report_ais_checkpoint_comparison as comparison # noqa + + +def _paired_rows(delta=0.02): + rows = [] + for dataset, strata in {"plain": [""], "stratified": ["x", "y"]}.items(): + for stratum in strata: + for index in range(4): + row = { + "sample_id": f"{dataset}:{stratum}:{index}", "dataset": dataset, "stratum": stratum, + "msa_baseline": 0.5 + 0.01 * index, "msa_boundary": 0.5 + 0.01 * index + delta, + "gt_objects_baseline": 10, "gt_objects_boundary": 10, + } + for column in comparison.EXTENT_COLUMNS: + row[f"{column}_baseline"] = 0.7 + row[f"{column}_boundary"] = 0.72 + for column in comparison.TIME_COLUMNS: + row[f"{column}_baseline"] = 1.0 + row[f"{column}_boundary"] = 1.1 + for column in (*comparison.FATE_COLUMNS, *comparison.OTHER_COUNT_COLUMNS): + row[f"{column}_baseline"] = 1 + row[f"{column}_boundary"] = 1 + rows.append(row) + return pd.DataFrame(rows) + + +def test_balanced_scores_equal_weight_strata(): + paired = _paired_rows() + group = paired[paired["dataset"] == "stratified"].copy() + group.loc[group["stratum"] == "x", "msa_baseline"] = 0.1 + group.loc[group["stratum"] == "y", "msa_baseline"] = 0.9 + assert comparison.balanced_scores(group)[0] == pytest.approx(0.5) + + +def test_hierarchical_bootstrap_and_domain_table_are_paired(): + paired = _paired_rows(delta=0.03) + overall, intervals = comparison.bootstrap(paired, n_bootstrap=500, seed=7) + assert overall["absolute_ci_low"] == pytest.approx(0.03) + assert overall["absolute_ci_high"] == pytest.approx(0.03) + assert overall["probability_boundary_better"] == 1.0 + domains = comparison.dataset_table(paired, intervals) + assert domains["improved"].all() and not domains["material_loss"].any() + assert np.allclose(domains["absolute_delta"], 0.03) + + +def test_manifest_coverage_rejects_partial_or_wrong_stratum(): + paired = _paired_rows() + manifest = {"samples": [ + {"sample_id": row.sample_id, "dataset": row.dataset, "stratum": row.stratum} + for row in paired.itertuples() + ]} + comparison.validate_manifest_coverage(paired, manifest) + with pytest.raises(ValueError, match="1 missing"): + comparison.validate_manifest_coverage(paired.iloc[:-1], manifest) + wrong = paired.copy() + wrong.loc[0, "stratum"] = "wrong" + with pytest.raises(ValueError, match="1 missing and 1 unexpected"): + comparison.validate_manifest_coverage(wrong, manifest) + + +def test_training_disjointness_audit_detects_dataset_alias(tmp_path): + data_root = tmp_path / "data" + data_root.mkdir() + raw = data_root / "ood/raw.tif" + label = data_root / "ood/label.tif" + raw.parent.mkdir() + raw.touch() + label.touch() + manifest = {"samples": [{"dataset": "vicar", "raw_path": "ood/raw.tif", "label_path": "ood/label.tif"}]} + training = tmp_path / "training.json" + boundary_training = tmp_path / "boundary_training.json" + training.write_text(json.dumps({"variant": "baseline", "datasets": {"train": {"train": []}}})) + boundary_training.write_text(json.dumps({"variant": "boundary", "datasets": {"train": {"train": []}}})) + audit = comparison.audit_training_disjointness(manifest, data_root, [training, boundary_training]) + assert audit["passed"] and not audit["dataset_overlap"] + + training.write_text(json.dumps({"variant": "baseline", "datasets": {"vicar_cells": {"train": []}}})) + with pytest.raises(RuntimeError, match="overlaps decoder training"): + comparison.audit_training_disjointness(manifest, data_root, [training, boundary_training]) diff --git a/test/test_ais_optimization.py b/test/test_ais_optimization.py new file mode 100644 index 000000000..14d5925aa --- /dev/null +++ b/test/test_ais_optimization.py @@ -0,0 +1,648 @@ +import json +import sys +from pathlib import Path + +import numpy as np +import pandas as pd +import pytest + +EVALUATION_ROOT = Path(__file__).parents[1] / "finetuning/v2/evaluation" +sys.path.insert(0, str(EVALUATION_ROOT)) +sys.path.insert(0, str(EVALUATION_ROOT / "optimization")) + +import benchmark_ais_optimization as ais # noqa +from benchmark_apg_3d import object_counts as reference_object_counts # noqa +from common import unmatched_objects # noqa + + +def _blobs(shape, centers, radii): + labels = np.zeros(shape, dtype="uint32") + grid = np.indices(shape) + for index, (center, radius) in enumerate(zip(centers, radii), start=1): + distance = sum(((g - c) / r) ** 2 for g, c, r in zip(grid, center, radius)) + labels[(distance <= 1) & (labels == 0)] = index + return labels + + +@pytest.fixture(scope="module") +def geodesic_prediction(): + """A noisy geodesic hybrid field of three touching-ish 2d objects, as the v4 decoder would predict it.""" + from micro_sam.v2.transforms.labels import GeodesicHybridDistanceTransform + + labels = _blobs((128, 160), [(40, 50), (40, 95), (95, 110)], [(25, 22), (25, 24), (20, 30)]) + target = GeodesicHybridDistanceTransform(foreground=True)(labels).astype("float32") + rng = np.random.default_rng(0) + prediction = target + rng.normal(0, 0.02, target.shape).astype("float32") + prediction[0] = np.clip(prediction[0], 0, 1) + return prediction, labels + + +def test_resolve_postprocessing_fills_library_defaults(): + from micro_sam.v2.postprocessing import default_postprocessing + + resolved = ais.resolve_postprocessing({}, "hvit_t") + assert resolved["sparse"] == default_postprocessing("hvit_t", "sparse") + assert resolved["dense"] == default_postprocessing("hvit_t", "dense") + volumes = ais.resolve_postprocessing({}, "hvit_t", ndim=3) + assert volumes["sparse"] == default_postprocessing("hvit_t", "sparse", ndim=3) + + flat = ais.resolve_postprocessing({"n_iter": 200, "dt": 1.0}, "hvit_t") + assert flat["sparse"]["n_iter"] == 200 and flat["sparse"]["dt"] == 1.0 + assert flat["dense"] == resolved["dense"] + + nested = ais.resolve_postprocessing({"sparse": {"sigma": 1.0}, "dense": {"beta": 0.7}}, "hvit_t") + assert nested["sparse"]["sigma"] == 1.0 and nested["dense"]["beta"] == 0.7 + + with pytest.raises(ValueError, match="Unknown AIS parameters"): + ais.resolve_postprocessing({"candidate_threshold": 1.0}, "hvit_t") + with pytest.raises(ValueError, match="only contain 'sparse' and 'dense'"): + ais.resolve_postprocessing({"sparse": {}, "n_iter": 50}, "hvit_t") + + +def test_resolve_postprocessing_null_uses_default_and_off_disables_filter(): + defaults = ais.resolve_postprocessing({}, "hvit_t")["sparse"] + resolved = ais.resolve_postprocessing({"boundary_magnitude_max": None}, "hvit_t")["sparse"] + assert resolved["boundary_magnitude_max"] == defaults["boundary_magnitude_max"] + assert np.isinf( + ais.resolve_postprocessing({"boundary_magnitude_max": "off"}, "hvit_t")["sparse"][ + "boundary_magnitude_max" + ] + ) + with pytest.raises(ValueError, match="only valid for boundary_magnitude_max"): + ais.resolve_postprocessing({"seed_floor": "off"}, "hvit_t") + + +def test_load_config_defaults_and_file(tmp_path): + name, mode, params_2d, params_3d = ais.load_config(None, "hvit_t") + assert (name, mode) == ("current-defaults", "auto") + assert params_2d == ais.resolve_postprocessing({}, "hvit_t", ndim=2) + assert params_3d == ais.resolve_postprocessing({}, "hvit_t", ndim=3) + assert params_3d["sparse"]["min_size"] == 100 and params_2d["sparse"]["min_size"] == 50 + + path = tmp_path / "candidate.json" + path.write_text(json.dumps({"name": "travel", "params_2d": {"n_iter": 400}, "params_3d": {"n_iter": 100}})) + name, mode, params_2d, params_3d = ais.load_config(path, "hvit_t") + assert name == "travel" and params_2d["sparse"]["n_iter"] == 400 and params_3d["sparse"]["n_iter"] == 100 + + # A volume takes the image overrides when it has none of its own. + path.write_text(json.dumps({"name": "shared", "mode": "sparse", "params_2d": {"sigma": 2.0}})) + _, mode, params_2d, params_3d = ais.load_config(path, "hvit_t") + assert mode == "sparse" and params_3d["sparse"]["sigma"] == 2.0 + + path.write_text(json.dumps({"name": "bad", "mode": "flow"})) + with pytest.raises(ValueError, match="Unknown mode"): + ais.load_config(path, "hvit_t") + + +def test_prediction_cache_validates_checkpoint_sample_and_shapes(tmp_path): + cache = ais.PredictionCache(tmp_path, "checkpoint-a", "manifest-a") + sample = {"sample_id": "toy:0"} + prediction = np.zeros((4, 8, 9), dtype="float32") + labels = np.zeros((8, 9), dtype="uint32") + record = { + "checkpoint_checksum": "checkpoint-a", "sample_id": "toy:0", "shape": list(prediction.shape), + } + cache.store(sample, prediction, labels, None, record) + loaded, loaded_labels, valid, loaded_record = cache.load(sample) + assert np.array_equal(loaded, prediction) and np.array_equal(loaded_labels, labels) + assert valid is None and loaded_record == record + + _, record_path = cache.paths(sample) + bad = dict(record, checkpoint_checksum="checkpoint-b") + record_path.write_text(json.dumps(bad)) + with pytest.raises(RuntimeError, match="different checkpoint"): + cache.load(sample) + + +def test_sparse_pipeline_matches_library(geodesic_prediction): + from micro_sam.v2.postprocessing import flow_instance_segmentation + + prediction, labels = geodesic_prediction + params = ais.resolve_postprocessing( + {"min_size": 20, "n_iter": 200, "dt": 0.5, "density_threshold": 5.0}, "hvit_t", + )["sparse"] + expected = flow_instance_segmentation(prediction[0], prediction[1:], model_type="hvit_t", n_threads=2, **params) + intermediates = ais.sparse_pipeline(prediction, params, None, 2) + assert np.array_equal(intermediates["segmentation"], expected) + assert intermediates["seeds"].max() == 3 + assert set(intermediates) >= {"before_min_size", "fg_mask", "density", "heightmap"} + assert len(np.unique(expected)) - 1 == 3 + + +def test_segment_prediction_matches_postprocess_unisam2(geodesic_prediction): + from common import postprocess_unisam2 + + prediction, _ = geodesic_prediction + params = ais.resolve_postprocessing({"min_size": 20}, "hvit_t")["sparse"] + mine = ais.segment_prediction(prediction, params, dense=False, spacing=None, model_type="hvit_t", n_threads=2) + reference = postprocess_unisam2(prediction, "livecell", "hvit_t", params={"min_size": 20}) + assert np.array_equal(mine, reference) + + +def test_matched_ids_agrees_with_unmatched_objects(geodesic_prediction): + prediction, labels = geodesic_prediction + params = ais.resolve_postprocessing({"min_size": 20}, "hvit_t")["sparse"] + segmentation = ais.sparse_pipeline(prediction, params, None, 2)["segmentation"] + # Delete one instance and shave another so that a match fails on IoU rather than on absence. + segmentation[segmentation == 1] = 0 + rows = np.where(segmentation == 2)[0] + segmentation[rows.min():rows.min() + 30][segmentation[rows.min():rows.min() + 30] == 2] = 0 + matched = set(ais.matched_ids(labels, segmentation).tolist()) + unmatched = set(np.unique(unmatched_objects(labels, segmentation)).tolist()) - {0} + assert matched | unmatched == {1, 2, 3} and not (matched & unmatched) + assert 1 in unmatched + + +def test_object_counts_agree_with_reference_for_volumes(): + labels = _blobs((12, 64, 64), [(6, 20, 20), (6, 40, 44), (1, 50, 12), (10, 12, 50)], + [(4, 10, 10), (5, 12, 9), (1, 8, 8), (0.5, 6, 6)]) + assert labels.max() == 4 + segmentation = labels.copy() + segmentation[segmentation == 2] = 0 # a miss + segmentation[labels == 3] = 7 # matched under another id + segmentation[:, 30:34, :] = 0 # shave everything + mine = ais.object_counts(labels, segmentation) + reference = reference_object_counts(labels, segmentation) + assert mine["gt_objects"] == reference["gt_objects"] == 4 + assert mine["matched"] == reference["merged"] + assert mine["unmatched"] == reference["unmatched"] + assert mine["severed_objects"] == reference["severed_objects"] >= 1 + assert mine["genuine_misses"] == reference["genuine_misses"] + assert mine["predicted_objects"] == 3 + + +def test_object_counts_for_images_report_no_severed_objects(): + labels = _blobs((64, 64), [(20, 20), (44, 44)], [(10, 10), (12, 9)]) + counts = ais.object_counts(labels, labels) + assert counts == { + "gt_objects": 2, "severed_objects": 0, "matched": 2, "unmatched": 0, "genuine_misses": 0, + "predicted_objects": 2, + } + empty = ais.object_counts(labels, np.zeros_like(labels)) + assert empty["matched"] == 0 and empty["unmatched"] == 2 and empty["genuine_misses"] == 2 + + +def test_seed_diagnostics_count_misses_splits_and_background_seeds(): + labels = _blobs((64, 96), [(20, 20), (20, 60), (48, 40)], [(10, 10), (10, 12), (9, 20)]) + seeds = np.zeros_like(labels, dtype="uint64") + seeds[20, 20] = 1 # object 1: one seed + seeds[18, 58] = 2 + seeds[22, 64] = 3 # object 2: split + seeds[5, 90] = 4 # background + seeds[60, 5] = 5 # background + segmentation = labels.copy() + segmentation[labels == 3] = 0 # object 3 (no seed) is missing from the result + intermediates = {"seeds": seeds, "fg_mask": labels != 0, "before_min_size": labels} + diagnostics = ais.seed_diagnostics(intermediates, labels, segmentation) + assert diagnostics["n_seeds"] == 5 + assert diagnostics["gt_with_0_seeds"] == 1 + assert diagnostics["gt_with_1_seed"] == 1 + assert diagnostics["gt_with_2plus_seeds"] == 1 + assert diagnostics["background_seeds"] == 2 + assert diagnostics["seeded_unmatched"] == 0 + assert diagnostics["unseeded_missing"] == 1 and diagnostics["unseeded_absorbed"] == 0 + assert diagnostics["matched_before_min_size"] == 3 + assert diagnostics["fg_iou"] == 1.0 + assert diagnostics["matched_iou"] == 1.0 + + # A seeded object the watershed then loses is lost at the assignment; here object 1 is undersized + # (only a quarter of it survives) and object 2, with two seeds, is a split. + segmentation = labels.copy() + segmentation[labels == 1] = 0 + segmentation[16:24, 16:24][labels[16:24, 16:24] == 1] = 1 + columns = np.indices(labels.shape)[1] + segmentation[(labels == 2) & (columns >= 56) & (columns < 64)] = 9 # three parts, none above IoU 0.5 + segmentation[(labels == 2) & (columns >= 64)] = 10 + diagnostics = ais.seed_diagnostics(intermediates, labels, segmentation) + assert diagnostics["seeded_unmatched"] == 2 + assert diagnostics["seeded_undersized"] == 1 and diagnostics["seeded_split"] == 1 + assert diagnostics["seeded_merged"] == 0 and diagnostics["seeded_oversized"] == 0 + + # One instance covering all three objects: object 1 (one seed) is merged, object 2 (two seeds) is a + # split, and the unseeded object 3 is absorbed (less than half of the instance is its own). + segmentation = np.where(labels != 0, 1, 0).astype("uint32") + diagnostics = ais.seed_diagnostics(intermediates, labels, segmentation) + assert diagnostics["seeded_merged"] == 1 and diagnostics["seeded_split"] == 1 + assert diagnostics["unseeded_absorbed"] == 1 and diagnostics["unseeded_missing"] == 0 + + +def test_object_fates_reports_iou_and_flags(): + labels = _blobs((64, 96), [(20, 20), (20, 60), (48, 40)], [(10, 10), (10, 12), (9, 20)]) + fates = ais.object_fates(labels, labels) + assert fates["ids"].tolist() == [1, 2, 3] + assert np.allclose(fates["iou"], 1.0) and fates["absorbed"].all() and not fates["merged"].any() + assert not fates["undersized"].any() + fates = ais.object_fates(labels, np.zeros_like(labels)) + assert np.allclose(fates["iou"], 0.0) and not fates["absorbed"].any() + + +def _sample_rows(datasets, msa_by_dataset, family=None, seen=""): + rows = [] + for dataset in datasets: + for index, msa in enumerate(msa_by_dataset[dataset]): + rows.append({ + "sample_id": f"{dataset}:{index}", "dataset": dataset, "ndim": 2, + "family": family.get(dataset, dataset) if family else dataset, "seen_in_training": seen, + "metric_mode": "sparse", "postprocessing_mode": "sparse", "initialization_seconds": 1.0, + "generation_seconds": 0.5, "total_seconds": 1.5, "peak_cuda_memory_bytes": 10 + index, + "msa": msa, "gt_objects": 4, "predicted_objects": 3, "matched": 3, "unmatched": 1, + "severed_objects": 0, "genuine_misses": 1, "matched_before_min_size": 3, "n_seeds": 3, + "gt_with_0_seeds": 1, "gt_with_1_seed": 3, "gt_with_2plus_seeds": 0, "background_seeds": 0, + "seeded_unmatched": 0, "fg_iou": 0.9, "pipeline_mismatch": 0, + }) + return pd.DataFrame(rows) + + +def test_summarize_reports_means_sums_and_balanced_row(): + samples = _sample_rows(["a", "b"], {"a": [0.2, 0.4], "b": [0.8, 0.8, 0.8]}) + summary = ais.summarize(samples).set_index("dataset") + assert summary.loc["a", "msa_mean"] == pytest.approx(0.3) + assert summary.loc["b", "n_samples"] == 3 and summary.loc["b", "matched"] == 9 + assert summary.loc[ais.BALANCED_ROW, "msa_mean"] == pytest.approx(0.55) + assert summary.loc[ais.BALANCED_ROW, "total_seconds"] == pytest.approx(7.5) + assert summary.loc["b", "peak_cuda_memory_bytes"] == 12 + assert "__family_macro__" not in summary.index + + +def test_summarize_adds_family_macros_for_crop_manifests(): + samples = _sample_rows( + ["cremi", "cremi_seen", "gonuclear"], {"cremi": [0.1], "cremi_seen": [0.3], "gonuclear": [0.6]}, + family={"cremi": "cremi", "cremi_seen": "cremi", "gonuclear": "gonuclear"}, + ) + samples.loc[samples["dataset"] == "cremi_seen", "seen_in_training"] = "True" + samples.loc[samples["dataset"] != "cremi_seen", "seen_in_training"] = "False" + summary = ais.summarize(samples).set_index("dataset") + assert summary.loc["__dataset_balanced__", "msa_mean"] == pytest.approx((0.1 + 0.3 + 0.6) / 3) + assert summary.loc["__family_macro__", "msa_mean"] == pytest.approx((0.2 + 0.6) / 2) + assert summary.loc["__unseen_macro__", "msa_mean"] == pytest.approx((0.1 + 0.6) / 2) + + +def test_gate_table_applies_the_generalization_rule(): + baseline = pd.Series({"a": 0.5, "b": 0.4, "c": 0.3, "d": 0.2}) + verdict = ais.gate_table(baseline, pd.Series({"a": 0.53, "b": 0.42, "c": 0.31, "d": 0.21})) + assert verdict["passed"] and verdict["n_up"] == 4 + # One dataset below both loss limits fails, however large the balanced gain. + verdict = ais.gate_table(baseline, pd.Series({"a": 0.9, "b": 0.9, "c": 0.9, "d": 0.18})) + assert not verdict["checks"]["no_dataset_below_loss_limits"] + # A tiny absolute loss on a near-zero score is tolerated by the absolute limit. + verdict = ais.gate_table(pd.Series({"a": 0.5, "b": 0.01}), pd.Series({"a": 0.6, "b": 0.008})) + assert verdict["checks"]["no_dataset_below_loss_limits"] + # Too many datasets down fails. + verdict = ais.gate_table(baseline, pd.Series({"a": 0.9, "b": 0.39, "c": 0.29, "d": 0.19})) + assert not verdict["checks"]["up_on_all_but_two"] + # Below the balanced gain fails. + verdict = ais.gate_table(baseline, pd.Series({"a": 0.501, "b": 0.401, "c": 0.301, "d": 0.201})) + assert not verdict["checks"]["balanced_gain_at_least_2_percent"] + + +def test_dataset_scores_use_negated_cremi_on_dense_data(): + samples = _sample_rows(["a"], {"a": [0.2, 0.4]}) + dense = samples.copy() + dense["dataset"], dense["metric_mode"], dense["cremi"] = "snemi", "dense", [0.9, 0.7] + scores = ais.dataset_scores(pd.concat([samples, dense], ignore_index=True)) + assert scores["a"] == pytest.approx(0.3) and scores["snemi"] == pytest.approx(-0.8) + + +def test_report_joins_subsets_and_flags_the_gate(tmp_path): + def write_run(name, subset, msa_by_dataset): + run_dir = tmp_path / f"{name}-{subset}" + run_dir.mkdir() + samples = _sample_rows(sorted(msa_by_dataset), msa_by_dataset) + samples.to_csv(run_dir / "samples.csv", index=False) + (run_dir / "metadata.json").write_text(json.dumps({"status": "complete", "config_name": name})) + return run_dir + + runs = { + "current-defaults": [ + write_run("current-defaults", "primary", {"a": [0.4], "b": [0.5]}), + write_run("current-defaults", "extra", {"c": [0.6]}), + ], + "candidate": [ + write_run("candidate", "primary", {"a": [0.44], "b": [0.55]}), + write_run("candidate", "extra", {"c": [0.63]}), + ], + } + table, details = ais.report(runs, "current-defaults") + table = table.set_index("config") + assert table.loc["candidate", "passed"] and table.loc["candidate", "n_datasets"] == 3 + assert not table.loc["current-defaults", "passed"] + assert details.query("config == 'candidate' and dataset == 'a'")["relative"].iloc[0] == pytest.approx(0.1) + + +def test_grid_combinations_deduplicate_flow_travel(): + grid = {"n_iter": [50, 100], "dt": [0.5, 1.0], "sigma": [0.5]} + combinations = ais.grid_combinations(grid, "sparse") + travels = sorted(round(c["n_iter"] * c["dt"], 6) for c in combinations) + assert travels == [25.0, 50.0, 100.0] + with pytest.raises(ValueError, match="Unknown sparse grid parameters"): + ais.grid_combinations({"beta": [0.5]}, "sparse") + + explicit = ais.grid_combinations({"combinations": [{"n_iter": 100}, {"n_iter": 100}, {"n_iter": 200}]}, "sparse") + assert explicit == [{"n_iter": 100}, {"n_iter": 200}] + with pytest.raises(ValueError, match="at least one"): + ais.grid_combinations({"combinations": []}, "sparse") + with pytest.raises(ValueError, match="Unknown sparse grid parameters"): + ais.grid_combinations({"combinations": [{"beta": 0.5}]}, "sparse") + + families = ais.grid_combinations({ + "shared": {"n_iter": [400], "dt": [0.5]}, + "families": {"base": {}, "ridge": {"contact_weight": [0.5, 1.0]}}, + }, "sparse") + assert len(families) == 3 + assert [combo["mechanism_family"] for combo in families] == ["base", "ridge", "ridge"] + with pytest.raises(ValueError, match="redefines shared"): + ais.grid_combinations({ + "shared": {"n_iter": [400]}, "families": {"bad": {"n_iter": [800]}}, + }, "sparse") + + +def test_sweep_shards_keep_expensive_cache_groups_together(): + grid = { + "foreground_threshold": [0.4, 0.5], "sigma": [0.5], "n_iter": [400, 800], "dt": [0.5], + "density_threshold": [5.0, 10.0], "min_size": [25, 50], "foreground_weight": [0.5, 1.0], + } + combinations = [ + ais.resolve_postprocessing({"sparse": combo}, "hvit_t")["sparse"] + for combo in ais.grid_combinations(grid, "sparse") + ] + shards = [ais.shard_combinations(combinations, "sparse", index, 3) for index in range(3)] + assert sum(map(len, shards)) == len(combinations) + assert {json.dumps(combo, sort_keys=True) for shard in shards for combo in shard} == { + json.dumps(combo, sort_keys=True) for combo in combinations + } + flow_keys = ais.SWEEP_CACHE_KEYS["sparse"] + groups = [{tuple(combo[key] for key in flow_keys) for combo in shard} for shard in shards] + assert all(not (first & second) for index, first in enumerate(groups) for second in groups[index + 1:]) + work = [sum(group[2] for group in shard_groups) for shard_groups in groups] + assert max(work) - min(work) <= max(group[2] for shard_groups in groups for group in shard_groups) + with pytest.raises(ValueError, match="Invalid shard"): + ais.shard_combinations(combinations, "sparse", 3, 3) + with pytest.raises(ValueError, match="only 4 distinct"): + ais.shard_combinations(combinations, "sparse", 0, 5) + + +def test_shared_configuration_ranks_by_mean_relative_optimum(tmp_path): + grid = pd.DataFrame({"sigma": [0.5, 1.0, 2.0], "n_iter": [50, 50, 50]}) + for dataset, scores in {"a": [0.5, 0.4, 0.2], "b": [0.3, 0.6, 0.3]}.items(): + table = grid.copy() + table["n_images"], table["msa_mean"], table["msa_std"] = 3, scores, 0.0 + table.to_csv(tmp_path / f"{dataset}.csv", index=False) + shared = ais.shared_configuration(tmp_path, ["a", "b"]) + assert list(shared.columns[:2]) == ["sigma", "n_iter"] + assert shared.iloc[0]["sigma"] == 1.0 # 0.8 + 1.0 over 1.0 + 0.5 + assert shared.iloc[0]["mean_relative"] == pytest.approx(0.9) + assert shared.iloc[0]["balanced"] == pytest.approx(0.5) + assert (tmp_path / "shared_config.csv").exists() + + +def test_gt_seed_markers_and_ridge_heightmap(): + labels = _blobs((64, 96), [(20, 20), (20, 60), (48, 40)], [(10, 10), (10, 12), (9, 20)]) + markers = ais.gt_seed_markers(labels) + assert markers.dtype == np.uint64 + ids, counts = np.unique(markers[markers != 0], return_counts=True) + assert ids.tolist() == [1, 2, 3] and counts.tolist() == [9, 9, 9] + # Every marker sits inside its own object. + for index in ids: + assert set(labels[markers == index].tolist()) == {index} + # A marker never leaks into a neighbouring object or the background, even for a one-pixel object. + tiny = np.zeros((8, 8), dtype="uint32") + tiny[2:6, 2:6] = 1 + tiny[3, 3] = 2 + markers = ais.gt_seed_markers(tiny) + assert (tiny[markers == 2] == 2).all() and (markers == 2).sum() == 1 + ridge = ais.gt_ridge_heightmap(labels) + assert ridge.dtype == np.float32 and ridge.flags["C_CONTIGUOUS"] + assert set(np.unique(ridge).tolist()) == {0.0, 1.0} + assert (ridge[labels == 0] == 0).all() + + +def test_oracle_sample_recovers_ground_truth_with_gt_seeds_and_foreground(geodesic_prediction): + prediction, labels = geodesic_prediction + sample = {"sample_id": "toy:0", "dataset": "toy", "ndim": 2} + context = {"ndim": 2, "metric_mode": "sparse", "postprocessing_mode": "sparse", "spacing": None, + "border_min_size": 0} + params = ais.resolve_postprocessing({"min_size": 20}, "hvit_t") + row = ais.oracle_sample(sample, context, prediction, labels, None, params, n_threads=2) + assert set(f"msa_{name}" for name in ais.ORACLES) <= set(row) + assert row["msa_gt_seeds_gt_fg"] >= row["msa_baseline"] + assert row["msa_gt_seeds_gt_fg"] > 0.95 and row["matched_gt_seeds_gt_fg"] == 3 + summary = ais.summarize_oracles(pd.DataFrame([row, {**row, "sample_id": "toy:1"}])).set_index("dataset") + assert summary.loc[ais.BALANCED_ROW, "msa_baseline"] == pytest.approx(row["msa_baseline"]) + assert summary.loc["toy", "gain_gt_seeds_gt_fg"] == pytest.approx( + row["msa_gt_seeds_gt_fg"] / row["msa_baseline"] - 1.0 + ) + + +def test_rank_shared_flags_gate_against_the_reference(): + import report_ais_sweep as rs + + grid = pd.DataFrame({"sigma": [0.5, 1.0, 0.5, 1.0], "boundary_magnitude_max": [np.nan, np.nan, 0.4, 0.4]}) + tables = {} + msa = {"a": [0.50, 0.48, 0.55, 0.54], "b": [0.30, 0.31, 0.33, 0.30], "c": [0.20, 0.22, 0.22, 0.10]} + for dataset, scores in msa.items(): + table = grid.copy() + table["n_images"], table["msa_mean"], table["msa_std"] = 5, scores, 0.0 + tables[dataset] = table + ranked = rs.rank_shared(tables, reference={"sigma": 0.5, "boundary_magnitude_max": None}) + ranked = ranked.set_index(["sigma", "boundary_magnitude_max"]) + # The reference row: no change, not passing. + assert ranked.loc[(0.5, "none"), "balanced_gain"] == pytest.approx(0.0) + assert not ranked.loc[(0.5, "none"), "passed"] + # sigma 0.5 with the filter improves every dataset by at least 10 %: passes. + assert ranked.loc[(0.5, 0.4), "passed"] and ranked.loc[(0.5, 0.4), "n_up"] == 3 + assert ranked.loc[(0.5, 0.4), "rel_c"] == pytest.approx(0.10) + # sigma 1.0 with the filter halves dataset c: fails the loss limit despite the balanced gain. + assert not ranked.loc[(1.0, 0.4), "passed"] + assert ranked["mean_relative_optimum"].max() <= 1.0 + with pytest.raises(ValueError, match="matches 0 rows"): + rs.rank_shared(tables, reference={"sigma": 2.0, "boundary_magnitude_max": None}) + + +def test_sweep_tables_union_keeps_mechanism_families(monkeypatch, tmp_path): + import report_ais_sweep as rs + + def fake_load(grid_path, *_args, **_kwargs): + table = pd.DataFrame({"n_iter": [800], "n_images": [2], "msa_mean": [0.5], "msa_std": [0.1]}) + if grid_path.stem == "ridge": + table["contact_weight"] = 1.0 + return {"a": table, "b": table.copy()} + + monkeypatch.setattr(rs, "load_sweep_tables", fake_load) + tables = rs.load_sweep_tables_many( + [Path("base.json"), Path("ridge.json")], ["primary"], tmp_path, tmp_path, tmp_path, + "hvit_t", "boundary", + ) + assert set(tables["a"]["mechanism_family"]) == {"base", "ridge"} + assert set(tables["a"]["contact_weight"].astype(str)) == {"none", "1.0"} + ranked = rs.rank_shared(tables) + assert len(ranked) == 2 and set(ranked["mechanism_family"]) == {"base", "ridge"} + + +def test_sweep_table_keeps_embedded_mechanism_family(monkeypatch, tmp_path): + import report_ais_sweep as rs + + table = pd.DataFrame({ + "n_iter": [800, 800], "mechanism_family": ["base", "ridge"], + "n_images": [2, 2], "msa_mean": [0.5, 0.6], "msa_std": [0.1, 0.1], + }) + monkeypatch.setattr(rs, "load_sweep_tables", lambda *_args, **_kwargs: {"a": table}) + tables = rs.load_sweep_tables_many( + [Path("boundary.json")], ["primary"], tmp_path, tmp_path, tmp_path, "hvit_t", "boundary", + ) + assert set(tables["a"]["mechanism_family"]) == {"base", "ridge"} + + +def test_sweep_plateau_selection_prefers_robust_cheaper_candidate(): + import report_ais_sweep as rs + + ranked = pd.DataFrame([ + {"balanced": 0.6000, "min_relative_optimum": 0.96, "n_iter": 1600, + "contact_weight": 2.0, "contact_mask_threshold": 0.5, "foreground_threshold": 0.4}, + {"balanced": 0.5995, "min_relative_optimum": 0.98, "n_iter": 800, + "contact_weight": "none", "contact_mask_threshold": "none", "foreground_threshold": 0.45}, + {"balanced": 0.5900, "min_relative_optimum": 1.00, "n_iter": 400, + "contact_weight": "none", "contact_mask_threshold": "none", "foreground_threshold": 0.5}, + ]) + selected = rs.select_plateau(ranked, tolerance=0.001) + assert selected["foreground_threshold"] == 0.45 + config = rs.selected_config(selected, "baseline-dice-optimum") + assert config["params_2d"]["sparse"] == {"foreground_threshold": 0.45, "n_iter": 800} + assert rs.select_plateau(ranked.drop(columns=["contact_weight"]), tolerance=0.001)[ + "foreground_threshold" + ] == 0.45 + selected["boundary_magnitude_max"] = np.inf + assert rs.selected_config(selected, "off")["params_2d"]["sparse"]["boundary_magnitude_max"] == "off" + + +def test_polish_grid_refines_boundary_coordinates_and_edges(): + import prepare_ais_reoptimization_polish as polish + + ranking = pd.DataFrame([ + {"mechanism_family": "ridge", "balanced": 0.50, "foreground_threshold": 0.4, + "foreground_weight": 0.75, "min_size": 50, "boundary_magnitude_max": 0.4, + "n_iter": 1200, "contact_weight": 1.0}, + {"mechanism_family": "ridge", "balanced": 0.502, "foreground_threshold": 0.4, + "foreground_weight": 0.75, "min_size": 50, "boundary_magnitude_max": 0.4, + "n_iter": 1600, "contact_weight": 1.0}, + ]) + combinations = polish.polish_combinations(ranking, top_per_family=1) + assert any(combo.get("contact_weight") == 3.0 for combo in combinations) + assert any(combo.get("n_iter") == 2400 for combo in combinations) + assert any(combo.get("boundary_magnitude_max") == "off" for combo in combinations) + + +def test_polish_grid_restores_numeric_optional_parameters_from_csv_strings(): + import prepare_ais_reoptimization_polish as polish + + ranking = pd.DataFrame([{ + "mechanism_family": "combined", "balanced": 0.5, "foreground_threshold": 0.45, + "foreground_weight": 0.5, "min_size": 50, "boundary_magnitude_max": 0.4, + "density_threshold": 20.0, "n_iter": 1200, "sigma": 0.5, "dt": 0.5, + "contact_weight": "1.0", "contact_mask_threshold": "0.5", + }]) + combinations = polish.polish_combinations(ranking, top_per_family=1) + assert combinations + assert all( + not isinstance(combo.get(key), str) + for combo in combinations + for key in ("contact_weight", "contact_mask_threshold") + if key in combo + ) + + +def test_polish_cli_reports_safe_shard_count(tmp_path, capsys): + import prepare_ais_reoptimization_polish as polish + + ranking = pd.DataFrame([{ + "mechanism_family": "base", "balanced": 0.5, "foreground_threshold": 0.4, + "foreground_weight": 0.75, "min_size": 50, "boundary_magnitude_max": 0.4, + "n_iter": 800, "sigma": 0.5, "dt": 0.5, + }]) + ranking_path, output_path = tmp_path / "ranking.csv", tmp_path / "polish.json" + ranking.to_csv(ranking_path, index=False) + assert polish.main(["--ranking", str(ranking_path), "--output", str(output_path)]) == 0 + output = capsys.readouterr().out + combinations = json.loads(output_path.read_text())["combinations"] + resolved = [ais.resolve_postprocessing({"sparse": combo}, "hvit_t")["sparse"] for combo in combinations] + expected = len({tuple(combo[key] for key in ais.SWEEP_CACHE_KEYS["sparse"]) for combo in resolved}) + assert f"{expected} flow-cache groups" in output + assert f"no more than {expected} sweep shards" in output + + +@pytest.fixture(scope="module") +def contact_prediction(geodesic_prediction): + """The fixture's field plus a fifth channel with the ground-truth contact lines.""" + from micro_sam.v2.transforms.labels import touching_boundaries + + prediction, labels = geodesic_prediction + contact = touching_boundaries(labels).astype("float32")[None] + return np.concatenate([prediction, contact], axis=0), labels + + +def test_resolve_postprocessing_accepts_the_contact_keywords(): + params = ais.resolve_postprocessing({"contact_weight": 1.0, "contact_mask_threshold": 0.5}, "hvit_t")["sparse"] + assert params["contact_weight"] == 1.0 and params["contact_mask_threshold"] == 0.5 + assert "contact_weight" not in ais.resolve_postprocessing({}, "hvit_t")["sparse"] + + +@pytest.mark.parametrize("overrides", [{}, {"contact_weight": 1.0}, {"contact_mask_threshold": 0.5}]) +def test_sparse_pipeline_matches_library_with_a_contact_channel(contact_prediction, overrides): + from micro_sam.v2.postprocessing import flow_instance_segmentation + + prediction, labels = contact_prediction + params = ais.resolve_postprocessing({"min_size": 20, "foreground_weight": 1.0, **overrides}, "hvit_t")["sparse"] + expected = flow_instance_segmentation( + prediction[0], prediction[1:4], model_type="hvit_t", n_threads=2, contact=prediction[4], **params, + ) + intermediates = ais.sparse_pipeline(prediction, params, None, 2) + assert np.array_equal(intermediates["segmentation"], expected) + mine = ais.segment_prediction(prediction, params, dense=False, spacing=None, model_type="hvit_t", n_threads=2) + assert np.array_equal(mine, expected) + diagnostics = ais.seed_diagnostics(intermediates, labels, expected) + assert 0.5 < diagnostics["fg_area_ratio"] < 2.0 + assert "fg_area_ratio" in ais.METRIC_COLUMNS + + +@pytest.mark.parametrize( + "overrides", + [ + {}, + {"contact_weight": 1.0}, + {"contact_mask_threshold": 0.5}, + {"contact_weight": 1.0, "contact_mask_threshold": 0.5}, + ], +) +def test_cached_sparse_scorer_matches_library_with_boundary_channel(contact_prediction, overrides, monkeypatch): + import parameter_search + from micro_sam.v2.postprocessing import flow_instance_segmentation + + prediction, labels = contact_prediction + params = ais.resolve_postprocessing( + { + "min_size": 20, + "n_iter": 200, + "density_threshold": 5.0, + "foreground_weight": 1.0, + **overrides, + }, + "hvit_t", + )["sparse"] + expected = flow_instance_segmentation( + prediction[0], prediction[1:4], contact=prediction[4], model_type="hvit_t", n_threads=2, **params, + ) + monkeypatch.setattr( + parameter_search, + "compute_metrics", + lambda segmentation, *_args, **_kwargs: {"segmentation": segmentation.copy()}, + ) + result = parameter_search.score_image_sparse_cached(prediction, labels, [params], n_threads=2)[0] + assert np.array_equal(result["segmentation"], expected) + + +def test_cached_sparse_scorer_rejects_boundary_parameters_without_channel(geodesic_prediction): + import parameter_search + + prediction, labels = geodesic_prediction + params = ais.resolve_postprocessing({"contact_weight": 1.0}, "hvit_t")["sparse"] + with pytest.raises(ValueError, match="need prediction channel 4"): + parameter_search.score_image_sparse_cached(prediction, labels, [params], n_threads=2) diff --git a/test/test_apg_3d_runner.py b/test/test_apg_3d_runner.py index 2f8654f91..2ac58d7df 100644 --- a/test/test_apg_3d_runner.py +++ b/test/test_apg_3d_runner.py @@ -3,7 +3,6 @@ from pathlib import Path import pytest -import numpy as np OPTIMIZATION_ROOT = Path(__file__).parents[1] / "finetuning/v2/evaluation/optimization" @@ -37,32 +36,6 @@ def test_volume_params_apply_overrides_and_reject_unknown_keys(tmp_path): def test_run_identity_is_stable(): - first = runner.run_identity("cfg", {"a": 1}, "checkpoint", "manifest", "trial-1") - assert first == runner.run_identity("cfg", {"a": 1}, "checkpoint", "manifest", "trial-1") - assert first != runner.run_identity("cfg", {"a": 2}, "checkpoint", "manifest", "trial-1") - - -@pytest.mark.parametrize("field", ["checkpoint_id", "manifest_checksum", "trial_id"]) -def test_run_identity_separates_experiments(tmp_path, field): - identity = {"checkpoint_id": "checkpoint", "manifest_checksum": "manifest", "trial_id": "trial-1"} - first = runner.run_dir(tmp_path, "primary", "cfg", {}, **identity) - first.mkdir(parents=True) - identity[field] = "different" - second = runner.run_dir(tmp_path, "primary", "cfg", {}, **identity) - assert first != second - assert runner.sibling_run_dirs(second) == [] - - -@pytest.mark.parametrize("missed_id", [0, 1, 2]) -def test_object_counts_include_matched_severed_objects(missed_id): - labels = np.zeros((8, 8, 8), dtype="uint32") - labels[2:5, 2:4, 2:4] = 1 - labels[5:7, 5:7, 6:] = 2 - segmentation = labels.copy() - segmentation[segmentation == missed_id] = 0 - counts = runner.object_counts(labels, segmentation) - assert counts == { - "gt_objects": 2, "severed_objects": 1, "merged": 2 - int(missed_id != 0), - "non_severed_matches": int(missed_id != 1), "unmatched": int(missed_id != 0), - "genuine_misses": int(missed_id == 1), - } + first = runner.run_identity("cfg", {"a": 1}) + assert first == runner.run_identity("cfg", {"a": 1}) + assert first != runner.run_identity("cfg", {"a": 2}) diff --git a/test/test_apg_manifest_subsets.py b/test/test_apg_manifest_subsets.py index 050109517..1e40c0cdc 100644 --- a/test/test_apg_manifest_subsets.py +++ b/test/test_apg_manifest_subsets.py @@ -38,3 +38,45 @@ def test_sample_counts_and_subsets(): assert set(benchmark.TRAINING_EXTRA_DATASETS).isdisjoint(benchmark.DATASETS_2D) with pytest.raises(ValueError): benchmark._sample_counts_2d("unknown") + + +def test_ood_extended_selection_is_stratified_and_uses_test_data(monkeypatch, tmp_path): + def candidates(dataset, _root, split="val", validate_raw=False, skip_read_errors=False): + assert split == "test" and validate_raw and skip_read_errors + if dataset == "bitdepth_nucseg": + strata = benchmark.OOD_EXTENDED_STRATUM_COUNTS[dataset] + template = f"{dataset}/data/{{stratum}}/images/img{{index}}.tif" + elif dataset == "cellbindb": + strata = {key: value + 2 for key, value in benchmark.OOD_EXTENDED_STRATUM_COUNTS[dataset].items()} + template = f"{dataset}/Other/{{stratum}}/sample{{index}}/img.tif" + elif dataset == "vicar": + strata = {key: value + 2 for key, value in benchmark.OOD_EXTENDED_STRATUM_COUNTS[dataset].items()} + template = f"{dataset}/labelled/{{stratum}}/img{{index}}.tif" + else: + return _fake_candidates(dataset, benchmark.SAMPLE_COUNTS_2D_OOD_EXTENDED[dataset]) + result = [] + for stratum, count in strata.items(): + for index in range(count): + sample = _fake_candidates(dataset, 1)[0] + sample["raw_path"] = template.format(stratum=stratum, index=index) + sample["label_path"] = sample["raw_path"].replace("img", "lab") + sample["object_count"] = index + 1 + result.append(sample) + return result + + monkeypatch.setattr(benchmark, "_scan_2d_dataset", candidates) + samples = benchmark._select_ood_extended_samples(tmp_path) + counts = {} + strata = {} + for sample in samples: + counts[sample["dataset"]] = counts.get(sample["dataset"], 0) + 1 + if "stratum" in sample: + key = (sample["dataset"], sample["stratum"]) + strata[key] = strata.get(key, 0) + 1 + assert counts == benchmark.SAMPLE_COUNTS_2D_OOD_EXTENDED + assert strata == { + (dataset, stratum): count + for dataset, expected in benchmark.OOD_EXTENDED_STRATUM_COUNTS.items() + for stratum, count in expected.items() + } + assert len({sample["sample_id"] for sample in samples}) == sum(counts.values()) == 180 diff --git a/test/test_models/test_unisam2.py b/test/test_models/test_unisam2.py index 76cf3cb1c..05b32724a 100644 --- a/test/test_models/test_unisam2.py +++ b/test/test_models/test_unisam2.py @@ -35,3 +35,26 @@ def test_unisam2_loads_a_narrow_state_dict_strictly(encoder): rebuilt.load_state_dict(state) assert all(torch.equal(value, torch.zeros_like(value)) for key, value in rebuilt.state_dict().items() if key.startswith(("out_conv", "base", "decoder"))) + + +def test_unisam2_fifth_channel_is_a_probability(encoder): + model = UniSAM2(encoder=encoder, output_channels=5, device="cpu", initial_features=32) + assert model.out_conv.out_channels == 5 and model.out_channels == 5 + with torch.no_grad(): + out = model(torch.rand(1, 3, 1, 256, 256)) + assert tuple(out.shape) == (1, 5, 1, 256, 256) + assert out[:, 0].min() >= 0 and out[:, 4].min() >= 0 and out[:, 4].max() <= 1 + assert out[:, 1:4].min() >= -1 and out[:, 1:4].max() <= 1 + + +def test_get_unisam2_model_reads_the_channel_count_off_the_checkpoint(encoder, tmp_path): + from micro_sam.v2.instance_segmentation import get_unisam2_model + + model = UniSAM2(encoder=encoder, output_channels=5, device="cpu", initial_features=32) + torch.save(model.state_dict(), tmp_path / "five.pt") + four = UniSAM2(encoder=encoder, output_channels=4, device="cpu", initial_features=32) + torch.save({"unetr_state": four.state_dict()}, tmp_path / "joint.pt") + + loaded = get_unisam2_model(tmp_path / "five.pt", device="cpu", encoder=encoder) + assert loaded.out_channels == 5 and loaded.out_conv.in_channels == 32 + assert get_unisam2_model(tmp_path / "joint.pt", device="cpu", encoder=encoder).out_channels == 4 diff --git a/test/test_submit_optimization_jobs.py b/test/test_submit_optimization_jobs.py index 185ad87a2..d32eb9047 100644 --- a/test/test_submit_optimization_jobs.py +++ b/test/test_submit_optimization_jobs.py @@ -1,5 +1,5 @@ import sys -import shlex +import subprocess from pathlib import Path import pytest @@ -51,7 +51,7 @@ def test_job_script_header_and_activation_order(tmp_path): assert all(i < first_command for i, line in enumerate(lines) if line.startswith("#SBATCH")) order = [ lines.index("set -eo pipefail"), lines.index("source ~/.bashrc"), lines.index("set -u"), - lines.index("micromamba activate super"), lines.index(f"cd {soj.REPOSITORY_ROOT}"), + lines.index("micromamba activate new-stack"), lines.index(f"cd {soj.REPOSITORY_ROOT}"), lines.index("export PYTHONUNBUFFERED=1"), ] assert order == sorted(order) @@ -60,6 +60,21 @@ def test_job_script_header_and_activation_order(tmp_path): assert "$SLURM_RESTART_COUNT " not in script and "$SLURM_RESTART_COUNT\"" not in script +def test_cpu_test_preset_does_not_request_a_gpu(tmp_path): + resources = soj.PRESETS["cpu-test"] + script = soj.render_job_script("cpu", tmp_path, 100, resources, tasks_per_job=48) + assert "#SBATCH -p standard96s:test" in script + assert "#SBATCH -t 00:59:00" in script + assert "#SBATCH -c 192" in script and "#SBATCH --mem=500G" in script + assert "#SBATCH --array=0-2%8" in script + assert not any(line.startswith("#SBATCH -G") for line in script.splitlines()) + assert "first_task=$((SLURM_ARRAY_TASK_ID * 48))" in script + assert 'for pid in "${pids[@]}"' in script + assert subprocess.run(["bash", "-n"], input=script, text=True, check=False).returncode == 0 + with pytest.raises(ValueError, match="positive"): + soj.render_job_script("bad", tmp_path, 2, resources, tasks_per_job=0) + + def test_preset_overrides_and_optional_lines(tmp_path): tasks = _tasks_file(tmp_path, [("a", "echo a")]) jobs_root = tmp_path / "jobs" @@ -126,6 +141,28 @@ def test_status_maps_sacct_rows_to_tags(tmp_path, monkeypatch, capsys): assert soj._expand_array_ids("15465049") == [] +def test_status_maps_packed_array_rows_to_all_child_tasks(tmp_path, monkeypatch, capsys): + tasks = _tasks_file(tmp_path, [(name, "echo") for name in ("a", "b", "c", "d")]) + jobs_root = tmp_path / "jobs" + soj.main([ + "submit", "--name", "packed", "--preset", "cpu-test", "--tasks-file", str(tasks), + "--tasks-per-job", "2", "--dry-run", "--jobs-root", str(jobs_root), + ]) + job_dir = _only_job_dir(jobs_root) + (job_dir / "job_id.txt").write_text("15465049\n") + rows = [ + {"JobID": "15465049_0", "State": "COMPLETED", "ExitCode": "0:0", "Elapsed": "00:10:00", "Restarts": "0", + "NodeList": "c0201"}, + {"JobID": "15465049_1", "State": "FAILED", "ExitCode": "1:0", "Elapsed": "00:10:00", "Restarts": "0", + "NodeList": "c0202"}, + ] + monkeypatch.setattr(soj, "_sacct_rows", lambda job_id: rows) + assert soj.status(job_dir) == 1 + lines = capsys.readouterr().out.splitlines() + assert sum("COMPLETED" in line for line in lines) == 2 + assert sum("FAILED" in line for line in lines) == 2 + + def test_benchmark_builder(tmp_path): config = tmp_path / "apg_my_config.json" config.write_text("{}") @@ -134,10 +171,7 @@ def test_benchmark_builder(tmp_path): assert len(tags) == len(set(tags)) == 4 for _, command in tasks: assert "--trial-id" in command and "--ndim 2" in command and "--subset holdout" in command - arguments = [shlex.split(command) for _, command in tasks] - assert any( - args[args.index("--config") + 1] == str(config.resolve()) for args in arguments if "--config" in args - ) + assert any(f"--config {config.resolve()}" in command for _, command in tasks) assert any("my_config" in tag for tag in tags) serial = campaign.benchmark_tasks([config], ["trial-1"], serialize=True, bracket=True) assert len(serial) == 1 diff --git a/test/test_v2_automatic_prompt_generation.py b/test/test_v2_automatic_prompt_generation.py index bf5ba7a26..35ec9cb53 100644 --- a/test/test_v2_automatic_prompt_generation.py +++ b/test/test_v2_automatic_prompt_generation.py @@ -1643,297 +1643,3 @@ def fake_stitch_segmentation(*, shape, **kwargs): assert calls["shape"] == (8, 12) assert segmentation.shape == (8, 12) - - -def _fake_tiled_embeddings(shape, tile_shape, halo, video): - """Tiled embeddings in the layout of `precompute_image_embeddings`, for a volume of `shape` (z, y, x) if - `video`, else for an image of `shape` (y, x). Every stored slice holds 100 * tile_id + z.""" - import zarr - - from bioimage_cpp.utils import Blocking - - from micro_sam.util import _create_dataset_without_data - from micro_sam.v2.batched_inference import _create_feature_dataset, _create_feature_levels - - root = zarr.group() - features = root.require_group("features") - features.attrs.update(shape=list(shape), tile_shape=list(tile_shape), halo=list(halo)) - tiling = Blocking([0, 0], list(shape[-2:]), list(tile_shape)) - for tile_id in range(tiling.number_of_blocks): - name = str(tile_id) - outer = tiling.get_block_with_halo(tile_id, list(halo)).outer_block - - def value(z): - return np.full((1, 2, 2, 2), 100 * tile_id + z, dtype="float32") - - if video: - n_slices = shape[0] - dataset = _create_feature_dataset(features, name, n_slices, value(0)) - levels = _create_feature_levels(root.require_group("fpn").require_group(name), n_slices, [value(0)] * 2) - for z in range(n_slices): - dataset[z] = value(z) - for level in levels: - level[z] = value(z) - pos_enc = _create_feature_levels(root.require_group("pos_enc").require_group(name), 1, [value(0)]) - pos_enc[0][0] = value(0) - else: - dataset = _create_dataset_without_data( - features, name, shape=(1, 2, 2, 2), dtype="float32", chunks=(1, 2, 2, 2), - ) - dataset[:] = value(0) - high_res = root.require_group("high_res_feats").require_group(name) - _create_dataset_without_data(high_res, "0", shape=(1, 2, 4, 4), dtype="float32", chunks=(1, 2, 4, 4))[:] = 0 - dataset.attrs["input_size"] = 8 - dataset.attrs["original_size"] = [int(e - b) for b, e in zip(outer.begin, outer.end)] - - embeddings = {"features": features, "input_size": None, "original_size": None} - for key in ("fpn", "pos_enc", "high_res_feats"): - if key in root: - embeddings[key] = root[key] - return embeddings - - -def _run_tiled_apg_with_embeddings(image, ndim, tile_shape, halo, image_embeddings, i=None): - """Run the tiled generator with the real stitching and a stand-in worker; return what each block got.""" - calls = [] - - class RecordingGenerator: - _pruning_protected_margin = None - - def initialize(self, block, **kwargs): - calls.append((np.asarray(block), kwargs)) - - def generate(self, **params): - return np.zeros(calls[-1][0].shape[:ndim], dtype="uint32") - - def clear_state(self): - pass - - segmenter = TiledAutomaticPromptGenerator(torch.nn.Identity(), _fake_apg_predictor()) - segmenter._pool = [RecordingGenerator()] - segmenter.initialize(image, ndim=ndim, tile_shape=tile_shape, halo=halo, image_embeddings=image_embeddings, i=i) - segmenter.generate() - return calls - - -def _tile_crop(shape, tile_shape, halo, tile_id): - from bioimage_cpp.utils import Blocking - - outer = Blocking([0, 0], list(shape), list(tile_shape)).get_block_with_halo(tile_id, list(halo)).outer_block - return tuple(slice(b, e) for b, e in zip(outer.begin, outer.end)) - - -@pytest.mark.skipif( - automatic_prompt_generation.bp is None, reason="Tiled stitching requires the optional 'bioimage_py'." -) -def test_tiled_apg_blocks_of_a_volume_read_their_tile_slices_from_the_embeddings(): - shape, in_plane_tile, in_plane_halo = (6, 8, 12), (8, 8), (2, 2) - volume = np.arange(np.prod(shape), dtype="float32").reshape(shape) - embeddings = _fake_tiled_embeddings(shape, in_plane_tile, in_plane_halo, video=True) - - calls = _run_tiled_apg_with_embeddings(volume, 3, (4, *in_plane_tile), (1, *in_plane_halo), embeddings) - - assert len(calls) == 4 # 2 z blocks x 2 in-plane tiles - blocks = set() - for block, kwargs in calls: - block_embeddings = kwargs["image_embeddings"] - assert kwargs["i"] is None - assert kwargs["normalization_bounds"] is None # nothing is encoded, so nothing is normalized - # The stored values give the tile id and the slice of each block slice. - values = np.asarray(block_embeddings["features"])[:, 0, 0, 0, 0].astype(int) - tile_id, z_start = values[0] // 100, values[0] % 100 - np.testing.assert_array_equal(values, 100 * tile_id + np.arange(z_start, z_start + block.shape[0])) - for level in block_embeddings["fpn"]: - np.testing.assert_array_equal(np.asarray(level)[:, 0, 0, 0, 0].astype(int), values) - assert int(np.asarray(block_embeddings["pos_enc"][0]).flat[0]) == 100 * tile_id - # The block must be the crop of the tile that the encoder saw. - crop = _tile_crop(shape[1:], in_plane_tile, in_plane_halo, tile_id) - np.testing.assert_array_equal(block, volume[(slice(z_start, z_start + block.shape[0]), *crop)]) - assert list(block_embeddings["original_size"]) == list(block.shape[1:]) - blocks.add((tile_id, z_start)) - assert blocks == {(0, 0), (1, 0), (0, 3), (1, 3)} - - -@pytest.mark.skipif( - automatic_prompt_generation.bp is None, reason="Tiled stitching requires the optional 'bioimage_py'." -) -def test_tiled_apg_blocks_of_a_volume_slice_read_that_slice_from_the_embeddings(): - shape, tile_shape, halo = (3, 8, 12), (8, 8), (2, 2) - volume = np.arange(np.prod(shape), dtype="float32").reshape(shape) - embeddings = _fake_tiled_embeddings(shape, tile_shape, halo, video=True) - - calls = _run_tiled_apg_with_embeddings(volume[1], 2, tile_shape, halo, embeddings, i=1) - - assert len(calls) == 2 - for block, kwargs in calls: - assert kwargs["i"] == 0 # the block's embeddings hold only the one slice - values = np.asarray(kwargs["image_embeddings"]["features"])[:, 0, 0, 0, 0].astype(int) - assert len(values) == 1 and values[0] % 100 == 1 - crop = _tile_crop(shape[1:], tile_shape, halo, values[0] // 100) - np.testing.assert_array_equal(block, volume[1][crop]) - - -@pytest.mark.skipif( - automatic_prompt_generation.bp is None, reason="Tiled stitching requires the optional 'bioimage_py'." -) -def test_tiled_apg_blocks_of_an_image_read_their_tile_from_the_embeddings(): - shape, tile_shape, halo = (8, 12), (8, 8), (2, 2) - image = np.arange(np.prod(shape), dtype="float32").reshape(shape) - embeddings = _fake_tiled_embeddings(shape, tile_shape, halo, video=False) - - calls = _run_tiled_apg_with_embeddings(image, 2, tile_shape, halo, embeddings) - - assert len(calls) == 2 - for block, kwargs in calls: - block_embeddings = kwargs["image_embeddings"] - assert kwargs["i"] is None and "high_res_feats" in block_embeddings - tile_id = int(np.asarray(block_embeddings["features"]).flat[0]) // 100 - np.testing.assert_array_equal(block, image[_tile_crop(shape, tile_shape, halo, tile_id)]) - - -def test_tiled_apg_without_embeddings_encodes_every_block(): - segmenter = TiledAutomaticPromptGenerator(torch.nn.Identity(), _fake_apg_predictor()) - segmenter.initialize(np.zeros((8, 12)), ndim=2, tile_shape=(8, 8), halo=(2, 2)) - assert segmenter._block_embedding_kwargs(0) == {} - - -@pytest.mark.parametrize("tile_shape, halo, image_shape, i, match", [ - ((4, 8), (2, 2), (8, 12), None, "in-plane tiling"), # another tile shape than the embeddings - ((8, 8), (1, 1), (8, 12), None, "in-plane tiling"), # another halo - ((8, 8), (2, 2), (8, 16), None, "are not for the input"), # another image - ((8, 8), (2, 2), (8, 12), 3, "are not for the input"), # a slice index the volume does not have -]) -def test_tiled_apg_rejects_embeddings_it_cannot_reuse(tile_shape, halo, image_shape, i, match): - embeddings = _fake_tiled_embeddings((3, 8, 12), (8, 8), (2, 2), video=True) - segmenter = TiledAutomaticPromptGenerator(torch.nn.Identity(), _fake_apg_predictor()) - with pytest.raises(ValueError, match=match): - segmenter.initialize( - np.zeros(image_shape), ndim=2, tile_shape=tile_shape, halo=halo, image_embeddings=embeddings, i=i, - ) - - -def test_tiled_apg_rejects_image_embeddings_for_a_volume(): - embeddings = _fake_tiled_embeddings((8, 12), (8, 8), (2, 2), video=False) - segmenter = TiledAutomaticPromptGenerator(torch.nn.Identity(), _fake_apg_predictor()) - with pytest.raises(ValueError, match="are not for the input"): - segmenter.initialize( - np.zeros((3, 8, 12)), ndim=3, tile_shape=(4, 8, 8), halo=(1, 2, 2), image_embeddings=embeddings, - ) - - -@pytest.mark.parametrize("n_prompts", [0, 5]) -def test_apg_proposal_progress_counts_completed_batches(monkeypatch, n_prompts): - segmenter = object.__new__(AutomaticPromptGenerator) - segmenter._is_initialized = True - segmenter._model_type = "hvit_t_cells" - segmenter._prediction = np.zeros((4, 16, 16), dtype="float32") - predictor = _RecordingPredictor((16, 16)) - segmenter._predictor = predictor - prompts = None if n_prompts == 0 else { - "points": np.full((n_prompts, 1, 2), 6, dtype="float32"), - "point_labels": np.ones((n_prompts, 1), dtype="int32"), - } - monkeypatch.setattr(automatic_prompt_generation, "derive_point_prompts", lambda *a, **k: prompts) - stages, updates = [], [] - proposals = segmenter.propose( - batch_size=2, pbar_init=lambda total, desc: stages.append((total, desc)), - pbar_update=lambda n: updates.append((n, len(predictor.calls))), - ) - assert len(proposals) == n_prompts - assert stages[0] == (1, "APG: deriving prompts") - assert updates[0] == (1, 0) - if n_prompts: - assert stages[1] == (3, "APG: prompting batches") - assert updates[1:] == [(1, 1), (1, 2), (1, 3)] - else: - assert len(stages) == len(updates) == 1 - - -@pytest.mark.parametrize("execution", ["thread", "process"]) -def test_tiled_apg_progress_is_live_and_on_the_calling_thread(monkeypatch, execution): - import threading - from concurrent.futures import ThreadPoolExecutor - - caller = threading.get_ident() - received = threading.Event() - stages, updates = [], [] - segmenter = TiledAutomaticPromptGenerator(torch.nn.Identity(), _fake_apg_predictor()) - segmenter.initialize(np.zeros((8, 8, 3), dtype="uint8"), tile_shape=(4, 4), halo=(1, 1)) - segmenter._execution = execution - - def dispatch(params, *args): - assert "pbar_init" not in params and "pbar_update" not in params - return lambda block, block_id: np.zeros((4, 4), dtype="uint32"), 2 - - monkeypatch.setattr(segmenter, f"_{execution}_dispatch", dispatch) - - def stitch(*, segmentation_function, shape, **kwargs): - assert not torch.is_grad_enabled() - with ThreadPoolExecutor(max_workers=2) as pool: - for block_id in range(4): - pool.submit(segmentation_function, np.zeros((4, 4, 3)), block_id).result() - # The callback must run while stitching is still active, not after it returns. - assert received.wait(timeout=5) - return np.zeros(shape, dtype="uint32") - - monkeypatch.setattr(automatic_prompt_generation, "bp", types.SimpleNamespace( - segmentation=types.SimpleNamespace(stitch_segmentation=stitch), - )) - - def update(n): - assert threading.get_ident() == caller - updates.append(n) - received.set() - - result = segmenter.generate(pbar_init=lambda n, desc: stages.append((n, desc)), pbar_update=update) - assert result.shape == (8, 8) - assert stages == [(4, "APG: segmenting tiles")] - assert updates == [1, 1, 1, 1] - - -def test_apg_parallel_progress_propagates_errors(): - updates = [] - - def fail(update): - update(1) - raise RuntimeError("segmentation failed") - - with pytest.raises(RuntimeError, match="segmentation failed"): - automatic_prompt_generation._run_with_progress(fail, updates.append) - assert updates == [1] - - -def test_volume_apg_progress_reports_scoring_and_propagation(monkeypatch): - import threading - - shape = (32, 32) - mask = _mask(shape, slice(4, 12), slice(4, 12)) - predictor = _VolumePredictor([([mask, mask], [0.9, 0.8]), ([mask], [0.9])]) - segmenter, _ = _volume_generator(monkeypatch, (3, *shape), predictor) - segmenter._model_type = "hvit_t_cells" - monkeypatch.setattr(automatic_prompt_generation, "derive_volume_prompts", lambda *a, **k: _two_anchor_prompts()) - propagator = _RecordingPropagator() - propagator.predictor_devices = [(predictor, "cpu")] - segmenter._propagator = propagator - segmenter._scoring_predictor_pool = [predictor] - stages, counts = [], [] - caller = threading.get_ident() - - def initialize(total, description): - assert threading.get_ident() == caller - stages.append((total, description)) - counts.append(0) - - def update(n): - assert threading.get_ident() == caller - counts[-1] += n - - result = segmenter.generate( - refinement=None, propagation_waves=1, pbar_init=initialize, pbar_update=update, - ) - assert result.shape == (3, *shape) - assert [description for _, description in stages] == [ - "APG: deriving volume prompts", "APG: scoring anchor slices", - "APG: propagation passes (wave 1/1)", "APG: merging volume masks", - ] - assert counts == [total for total, _ in stages] == [1, 2, 2, 1] diff --git a/test/test_v2_automatic_segmentation.py b/test/test_v2_automatic_segmentation.py index 46b3ec374..9d3787aaa 100644 --- a/test/test_v2_automatic_segmentation.py +++ b/test/test_v2_automatic_segmentation.py @@ -1145,25 +1145,188 @@ def test_decoder_output_is_moved_to_cpu_before_the_float_cast(): assert calls == ["detach", "cpu", "float"] -def test_decoder_width_mismatch_names_torch_em(): - """An outdated torch-em builds a fixed-width decoder; say so instead of dumping size mismatches.""" - from micro_sam.v2.instance_segmentation import CONFIGURABLE_DECODER_WIDTH_VERSION, _check_decoder_width - - # Only 'out_conv.weight.shape[1]' is read, so a bare namespace stands in for the built model. - model = types.SimpleNamespace(out_conv=types.SimpleNamespace(weight=torch.zeros(4, 64, 1, 1, 1))) - - _check_decoder_width(model, 64) # Matching width: no error. - - with pytest.raises(RuntimeError) as excinfo: - _check_decoder_width(model, 32) - message = str(excinfo.value) - assert "torch-em" in message - assert CONFIGURABLE_DECODER_WIDTH_VERSION in message - assert "64" in message and "32" in message - - -def test_decoder_width_check_skips_models_without_out_conv(): - """The check is a diagnostic, so a module that has no 'out_conv' passes through it untouched.""" - from micro_sam.v2.instance_segmentation import _check_decoder_width - - _check_decoder_width(types.SimpleNamespace(), 32) +def _geodesic_field_with_false_region(): + """Two real objects in the geodesic hybrid field plus a false foreground blob carrying the background fill.""" + from micro_sam.v2.transforms.labels import GeodesicHybridDistanceTransform + + labels = np.zeros((96, 128), dtype="uint32") + yy, xx = np.indices(labels.shape) + labels[((yy - 30) / 18) ** 2 + ((xx - 32) / 16) ** 2 <= 1] = 1 + labels[((yy - 60) / 16) ** 2 + ((xx - 90) / 20) ** 2 <= 1] = 2 + target = GeodesicHybridDistanceTransform(foreground=True)(labels).astype("float32") + false_blob = ((yy - 22) / 9) ** 2 + ((xx - 100) / 12) ** 2 <= 1 + target[0][false_blob] = 1.0 # confident foreground ... + target[1:, false_blob] = 1.0 # ... with the fill value the decoder emits in the background + return target, labels, false_blob + + +def test_drop_instances_without_boundary_dip_removes_false_regions_only(): + from micro_sam.v2.postprocessing import drop_instances_without_boundary_dip, flow_instance_segmentation + + prediction, labels, false_blob = _geodesic_field_with_false_region() + params = dict(model_type="hvit_t", min_size=20, n_iter=200, dt=0.5, density_threshold=5.0, n_threads=1) + # The hvit_t default filter is on, so the unfiltered reference disables it explicitly. + unfiltered = flow_instance_segmentation( + prediction[0], prediction[1:], boundary_magnitude_max=float("inf"), **params + ) + assert len(np.unique(unfiltered)) - 1 == 3, "expected two objects and the false region" + filtered = drop_instances_without_boundary_dip(unfiltered, prediction[1:][-2:], max_median=0.5) + assert len(np.unique(filtered)) - 1 == 2 + assert (filtered[false_blob] == 0).all() + for index in (1, 2): + kept = np.unique(filtered[labels == index]) + assert len(kept[kept != 0]) == 1 + # Through the keyword, and through the hvit_t default (0.4), which drops the same false region here. + via_keyword = flow_instance_segmentation(prediction[0], prediction[1:], boundary_magnitude_max=0.5, **params) + assert np.array_equal(via_keyword, filtered) + via_default = flow_instance_segmentation(prediction[0], prediction[1:], **params) + assert np.array_equal(via_default, filtered) + + +def test_default_postprocessing_per_backbone_and_dimension(): + from micro_sam.v2.postprocessing import DEFAULT_POSTPROCESSING, default_postprocessing + + # The optimized hvit_t defaults: images get the filter, wider smoothing and a ground-truth-like size + # floor; volumes keep the registry smoothing and size floor and add the filter. The other backbones keep + # the registry values. + images = default_postprocessing("hvit_t", "sparse", ndim=2) + volumes = default_postprocessing("hvit_t", "sparse", ndim=3) + assert images["boundary_magnitude_max"] == 0.4 and images["sigma"] == 1.0 and images["min_size"] == 50 + assert volumes["boundary_magnitude_max"] == 0.4 and volumes["sigma"] == 0.5 and volumes["min_size"] == 100 + assert {k: v for k, v in volumes.items() if k not in ("min_size", "sigma")} == { + k: v for k, v in images.items() if k not in ("min_size", "sigma") + } + for backbone in ("hvit_s", "hvit_b", "hvit_l"): + assert default_postprocessing(backbone, "sparse")["boundary_magnitude_max"] is None + assert default_postprocessing(backbone, "sparse", ndim=3) == default_postprocessing(backbone, "sparse") + assert "sparse_volume" in DEFAULT_POSTPROCESSING["hvit_t"] + # A finetuned model built on the backbone resolves to the backbone's table. + assert default_postprocessing("hvit_t_cells", "sparse") == images + # The returned dict is a copy: mutating it must not change the table. + images["sigma"] = 99.0 + assert default_postprocessing("hvit_t", "sparse")["sigma"] == 1.0 + + +def test_lower_height_under_seeds_modes(): + from micro_sam.v2.postprocessing import lower_height_under_seeds + from bioimage_cpp.segmentation import watershed + + # Two touching squares; a single-pixel seed on a height spike at each centre. With the monotone flooding + # the seed on the higher spike floods last and loses its square; a floor restores both. + hmap = np.full((20, 40), 0.3, dtype="float32") + hmap[:, 19:21] = 0.45 # the contact ridge + seeds = np.zeros((20, 40), dtype="uint64") + seeds[10, 10], seeds[10, 30] = 1, 2 + hmap[10, 10], hmap[10, 30] = 0.5, 0.6 + mask = np.ones((20, 40), dtype=bool) + broken = watershed(hmap, markers=seeds, mask=mask) + assert (broken == 2).sum() <= 1, "the test needs the monotone-flooding failure to be present" + for mode in ("zero", "ring"): + lowered = lower_height_under_seeds(hmap, seeds, mode) + assert lowered.dtype == np.float32 and lowered.flags["C_CONTIGUOUS"] + fixed = watershed(lowered, markers=seeds, mask=mask) + assert abs(int((fixed == 1).sum()) - int((fixed == 2).sum())) <= 40 + assert lower_height_under_seeds(hmap, seeds, "zero")[10, 10] == 0.0 + ring = lower_height_under_seeds(hmap, seeds, "ring") + assert ring[10, 10] == pytest.approx(0.3) and ring[10, 30] == pytest.approx(0.3) + assert ring[5, 5] == pytest.approx(0.3) and ring[10, 19] == pytest.approx(0.45) + assert lower_height_under_seeds(hmap, seeds, "none") is hmap + with pytest.raises(ValueError, match="Unknown seed floor"): + lower_height_under_seeds(hmap, seeds, "deep") + + +def test_seed_floor_default_is_off_everywhere(): + from micro_sam.v2.postprocessing import DEFAULT_POSTPROCESSING, default_postprocessing + + for backbone in DEFAULT_POSTPROCESSING: + for ndim in (2, 3): + assert default_postprocessing(backbone, "sparse", ndim=ndim)["seed_floor"] == "none" + + +def _touching_ellipses(shape, centers, radii): + """Ellipses with consecutive ids; later ones do not overwrite earlier ones.""" + labels = np.zeros(shape, dtype="uint32") + grid = np.indices(shape) + for index, (center, radius) in enumerate(zip(centers, radii), start=1): + distance = sum(((g - c) / r) ** 2 for g, c, r in zip(grid, center, radius)) + labels[(distance <= 1) & (labels == 0)] = index + return labels + + +def _best_iou(labels, segmentation, label_id): + mask = labels == label_id + ious = [ + (mask & (segmentation == seg_id)).sum() / (mask | (segmentation == seg_id)).sum() + for seg_id in np.unique(segmentation) if seg_id != 0 + ] + return max(ious) if ious else 0.0 + + +@pytest.fixture(scope="module") +def big_small_contact_prediction(): + """A large and a small ellipse touching each other, with the contact line as a fifth channel. + + With a flat height map (foreground weight 1) the fronts of the two seeds meet halfway between the seeds, + so the small object's basin falls below the size floor and the big instance swallows it. The contact + channel puts the split back onto the true contact line. + """ + from micro_sam.v2.transforms.labels import GeodesicHybridDistanceTransform + + labels = _touching_ellipses((128, 160), [(64, 50), (64, 104)], [(40, 40), (16, 16)]) + prediction = GeodesicHybridDistanceTransform(contact=True)(labels).astype("float32") + return prediction, labels + + +def test_flow_segmentation_contact_ridge_and_mask_split_touching_objects(big_small_contact_prediction): + from micro_sam.v2.postprocessing import flow_instance_segmentation + + prediction, labels = big_small_contact_prediction + foreground, distances, contact = prediction[0], prediction[1:4], prediction[4] + common = dict(model_type="hvit_t", foreground_weight=1.0, boundary_magnitude_max=float("inf")) + + merged = flow_instance_segmentation(foreground, distances, **common) + assert len(np.unique(merged)) - 1 == 1 + assert _best_iou(labels, merged, 2) < 0.2 + + # An unused contact map changes nothing. + assert np.array_equal(flow_instance_segmentation(foreground, distances, contact=contact, **common), merged) + + ridge = flow_instance_segmentation(foreground, distances, contact=contact, contact_weight=1.0, **common) + masked = flow_instance_segmentation(foreground, distances, contact=contact, contact_mask_threshold=0.5, **common) + for segmentation in (ridge, masked): + assert len(np.unique(segmentation)) - 1 == 2 + assert _best_iou(labels, segmentation, 1) > 0.95 and _best_iou(labels, segmentation, 2) > 0.9 + # Every foreground pixel is assigned, also the excluded contact pixels of the mask mode. + assert np.array_equal(segmentation > 0, foreground > 0.5) + + +def test_flow_segmentation_rejects_wrong_channel_counts_and_orphan_contact_keywords(big_small_contact_prediction): + from micro_sam.v2.postprocessing import flow_instance_segmentation + + prediction, _ = big_small_contact_prediction + foreground, distances, contact = prediction[0], prediction[1:4], prediction[4] + # Three channels for a 2d prediction drop the z channel; the 2d channels alone work as well. + reference = flow_instance_segmentation(foreground, distances, model_type="hvit_t") + assert np.array_equal(flow_instance_segmentation(foreground, distances[1:], model_type="hvit_t"), reference) + with pytest.raises(ValueError, match="distance channels"): + flow_instance_segmentation(foreground, prediction[1:], model_type="hvit_t") + with pytest.raises(ValueError, match="contact"): + flow_instance_segmentation(foreground, distances, model_type="hvit_t", contact_weight=1.0) + with pytest.raises(ValueError, match="contact"): + flow_instance_segmentation(foreground, distances, model_type="hvit_t", contact_mask_threshold=0.5) + with pytest.raises(ValueError, match="shape"): + flow_instance_segmentation(foreground, distances, model_type="hvit_t", contact=contact[:-1], contact_weight=1.0) + + +def test_segment_from_predictions_forwards_the_contact_channel(big_small_contact_prediction): + from micro_sam.v2.instance_segmentation import _segment_from_predictions + + prediction, labels = big_small_contact_prediction + common = dict(model_type="hvit_t", foreground_weight=1.0, boundary_magnitude_max=float("inf")) + four = _segment_from_predictions(prediction[:4], mode="sparse", **common) + five = _segment_from_predictions(prediction, mode="sparse", **common) + assert np.array_equal(four, five) + ridge = _segment_from_predictions(prediction, mode="sparse", contact_weight=1.0, **common) + assert len(np.unique(ridge)) - 1 == 2 and _best_iou(labels, ridge, 2) > 0.9 + with pytest.raises(ValueError, match="contact"): + _segment_from_predictions(prediction[:4], mode="sparse", contact_weight=1.0, **common) diff --git a/test/test_v2_label_transforms.py b/test/test_v2_label_transforms.py new file mode 100644 index 000000000..1228533cc --- /dev/null +++ b/test/test_v2_label_transforms.py @@ -0,0 +1,101 @@ +"""Tests for the label transforms of the automatic branch, in particular the contact channel.""" + +import numpy as np +import pytest + +from micro_sam.v2.transforms.labels import ( + DirectedPerObjectBoundaryDistanceTransform, GeodesicHybridDistanceTransform, _JointGeodesicLabelTransform, + touching_boundaries, +) + + +def _two_squares(gap: int) -> np.ndarray: + """Two squares side by side, touching for gap=0 or separated by 'gap' background columns.""" + labels = np.zeros((40, 60), dtype="uint16") + labels[10:30, 10:30] = 1 + labels[10:30, 30 + gap:50 + gap] = 2 + return labels + + +def test_touching_boundaries_marks_both_sides_of_a_direct_contact(): + contact = touching_boundaries(_two_squares(gap=0), dilation=0) + rows, cols = np.nonzero(contact) + assert set(cols.tolist()) == {29, 30} + # The background pixel just beyond either end of the line sees both objects as well. + assert rows.min() == 9 and rows.max() == 30 + assert contact[10:30, 29].all() and contact[10:30, 30].all() + + +def test_touching_boundaries_marks_a_one_pixel_gap(): + contact = touching_boundaries(_two_squares(gap=1), dilation=0) + assert set(np.nonzero(contact)[1].tolist()) == {30} + # A two pixel gap is out of reach of the default radius. + assert not touching_boundaries(_two_squares(gap=2), dilation=0).any() + + +def test_touching_boundaries_ignores_isolated_objects_and_dilates(): + labels = _two_squares(gap=0) + labels[2:8, 52:58] = 3 + contact = touching_boundaries(labels) + assert not contact[2:8, 52:58].any() + # One dilation pass widens the two pixel line to four pixels. + assert set(np.nonzero(contact[20])[0].tolist()) == {28, 29, 30, 31} + assert not touching_boundaries(np.zeros((8, 8), dtype="uint8")).any() + + +def test_touching_boundaries_handles_3d_and_large_ids(): + labels = np.zeros((3, 20, 20), dtype="uint32") + labels[:, 5:10, 5:15] = 70000 + labels[:, 10:15, 5:15] = 3 + contact = touching_boundaries(labels, dilation=0) + assert contact.shape == labels.shape + assert set(np.nonzero(contact[1])[0].tolist()) == {9, 10} + + +@pytest.mark.parametrize( + "transform_class", [DirectedPerObjectBoundaryDistanceTransform, GeodesicHybridDistanceTransform], +) +def test_contact_channel_is_appended_last(transform_class): + labels = _two_squares(gap=0) + # Ellipsoidal ends so that the objects do not fill their bounding boxes. + labels[10:12, 10:12] = 0 + labels[28:30, 48:50] = 0 + plain = transform_class()(labels) + with_contact = transform_class(contact=True)(labels) + assert plain.shape == (4, 40, 60) + assert with_contact.shape == (5, 40, 60) + assert with_contact.dtype == np.float32 + np.testing.assert_array_equal(with_contact[:4], plain) + np.testing.assert_array_equal(with_contact[4] > 0, touching_boundaries(labels)) + assert set(np.unique(with_contact[4]).tolist()) == {0.0, 1.0} + + +def test_contact_channel_follows_the_instance_channel_layout_and_3d_input(): + labels = _two_squares(gap=0) + joint = _JointGeodesicLabelTransform(contact=True)(labels) + assert joint.shape == (6, 40, 60) + np.testing.assert_array_equal(joint[0] > 0, labels > 0) + np.testing.assert_array_equal(joint[5] > 0, touching_boundaries(labels)) + + volume = np.stack([labels, labels]) + target = GeodesicHybridDistanceTransform(contact=True)(volume) + assert target.shape == (5, 2, 40, 60) + np.testing.assert_array_equal(target[4] > 0, touching_boundaries(volume)) + + +def test_object_boundaries_mode_covers_every_object_edge(): + from micro_sam.v2.transforms.labels import object_boundaries + + labels = _two_squares(gap=0) + labels[2:8, 52:58] = 3 + full = object_boundaries(labels, dilation=0) + # Every object contributes its inner boundary, the isolated one included. + assert full[2, 52:58].all() and full[10, 10:30].all() and full[10:30, 29].all() + assert not full[15, 15:25].any() + contact = touching_boundaries(labels, dilation=0) + assert (contact & ~full).sum() <= contact.sum() // 2 # the contact line is (mostly) a subset of the boundaries + target = GeodesicHybridDistanceTransform(contact=True, contact_mode="all")(labels) + assert target.shape == (5, 40, 60) + np.testing.assert_array_equal(target[4] > 0, object_boundaries(labels)) + with pytest.raises(ValueError, match="contact_mode"): + GeodesicHybridDistanceTransform(contact=True, contact_mode="edges") diff --git a/test/test_v2_training.py b/test/test_v2_training.py index 95b83d9d1..dca5e308b 100644 --- a/test/test_v2_training.py +++ b/test/test_v2_training.py @@ -16,7 +16,7 @@ import torch.multiprocessing as mp from micro_sam.v2.transforms.raw import VideoAugment -from micro_sam.v2.loss.directed_distance_based import _masked_mse, DirectedDistanceLoss +from micro_sam.v2.loss.directed_distance_based import _masked_mse, boundary_band, DirectedDistanceLoss def _free_port(): @@ -638,3 +638,63 @@ def test_the_trainer_trains_in_bfloat16_without_a_scaler_on_ampere(monkeypatch): if __name__ == "__main__": unittest.main() + + +class TestDirectedDistanceLossVariants(unittest.TestCase): + """The contact and boundary-weighted terms extend the loss without touching its default behaviour.""" + + def _batch(self, n_channels=4): + torch.manual_seed(0) + prediction = torch.rand(2, n_channels, 1, 32, 32) + target = torch.zeros(2, n_channels, 1, 32, 32) + target[:, 0, :, 8:24, 8:24] = 1.0 + target[:, 1:4] = 1.0 + target[:, 2:4, :, 8:24, 8:24] = 0.3 + if n_channels == 5: + target[:, 4, :, 8:24, 15:17] = 1.0 + return prediction, target + + def test_default_loss_is_unchanged(self): + from torch_em.loss import DiceLoss + + prediction, target = self._batch() + fg_mask, z_mask = target[:, 0:1], torch.zeros_like(target[:, 0:1]) + expected = ( + DiceLoss()(prediction[:, 0:1], target[:, 0:1]) + + _masked_mse(prediction[:, 1:2], target[:, 1:2], z_mask) + + _masked_mse(prediction[:, 2:3], target[:, 2:3], fg_mask) + + _masked_mse(prediction[:, 3:4], target[:, 3:4], fg_mask) + ) + loss = DirectedDistanceLoss() + self.assertEqual(loss.n_channels, 4) + self.assertTrue(torch.equal(loss(prediction, target), expected)) + + def test_boundary_weight_adds_a_cross_entropy_term_in_a_band(self): + prediction, target = self._batch() + band = boundary_band(target[:, 0:1], radius=2) + # A 16 x 16 square: a two pixel ring inside (112 px) and outside (144 px) of its edge, per sample. + self.assertEqual(int(band.sum()), 2 * 256) + self.assertTrue(band[0, 0, 0, 16, 16] == 0 and band[0, 0, 0, 8, 16] == 1 and band[0, 0, 0, 6, 16] == 1) + plain = DirectedDistanceLoss()(prediction, target) + weighted = DirectedDistanceLoss(boundary_weight=4.0)(prediction, target) + self.assertGreater(weighted.item(), plain.item()) + # bfloat16 predictions must not produce an infinite log. + self.assertTrue(torch.isfinite(DirectedDistanceLoss(boundary_weight=4.0)(prediction.bfloat16(), target))) + + def test_contact_channel_is_trained_and_required(self): + prediction, target = self._batch(n_channels=5) + loss = DirectedDistanceLoss(contact=True) + self.assertEqual(loss.n_channels, 5) + value = loss(prediction, target) + self.assertTrue(torch.isfinite(value)) + without_contact = DirectedDistanceLoss()(prediction[:, :4], target[:, :4]) + self.assertGreater(value.item(), without_contact.item()) + # A perfect contact prediction adds (almost) nothing. + perfect = prediction.clone() + perfect[:, 4] = target[:, 4] + self.assertAlmostEqual(loss(perfect, target).item(), without_contact.item(), places=3) + with self.assertRaises(AssertionError): + DirectedDistanceLoss()(prediction, target) + with self.assertRaises(AssertionError): + loss(prediction[:, :4], target[:, :4]) + self.assertEqual(loss.init_kwargs["contact"], True)