Refactor automatic segmentation v2 API design (#1301)

* Refactor automatic segmentation v2 API design

* Improve code design by removing redundant code

* Clarfiy AMG helpers and drop checkpoint module patches

* Consolidate AIS stub encoders into one class

* Update training CLI entry point for joint training

Author: anwai98

finetuning/v2/evaluation/common.py

@@ -467,13 +467,13 @@ def load_unisam2_model(checkpoint_path, device):
 
 
 def predict_unisam2(model, raw, ndim, device):
-    from micro_sam.v2.automatic_segmentation import run_unisam2_inference
+    from micro_sam.v2.instance_segmentation import get_unisam2_segmentation_generator
     is_3d = (ndim == 3)
     tile_shape = (4, 384, 384) if is_3d else (384, 384)
     halo = (2, 64, 64) if is_3d else (64, 64)
-    return run_unisam2_inference(
-        model=model, raw=raw, ndim=ndim, device=device, tile_shape=tile_shape, halo=halo,
-    )
+    segmenter = get_unisam2_segmentation_generator(model, is_tiled=True, device=device)
+    segmenter.initialize(raw, ndim=ndim, tile_shape=tile_shape, halo=halo)
+    return segmenter.get_state()["prediction"]
 
 
 def postprocess_unisam2(out, dataset_name, backend="python"):

finetuning/v2/evaluation/grid_search_automatic_cells.py

@@ -122,7 +122,7 @@ def load_model(device, checkpoint_path=None, model_name=MODEL_NAME):
     Returns:
         The UniSAM2 model in eval mode.
     """
-    from micro_sam.v2.automatic_segmentation import get_unisam2_model
+    from micro_sam.v2.instance_segmentation import get_unisam2_model
 
     if checkpoint_path is not None:
         print(f"Loading UniSAM2 model from custom checkpoint '{checkpoint_path}'.")

micro_sam/_cli.py

@@ -492,15 +492,16 @@ def inference_segmentation(
 
     from .v2.util import DEFAULT_MODEL
     from .v1.automatic_segmentation import _get_inputs_from_paths
-    from .v2.automatic_segmentation import get_segmenter, automatic_instance_segmentation
+    from .v2.automatic_segmentation import get_predictor_and_segmenter, automatic_instance_segmentation
 
     model_type = model_type or DEFAULT_MODEL
     tile_shape = _parse_shape(tile_shape)
     halo = _parse_shape(halo)
     generate_kwargs = _parse_extra(ctx.args)
 
-    segmenter = get_segmenter(
-        model_type=model_type, checkpoint=checkpoint_path, device=device, is_tiled=tile_shape is not None,
+    predictor, segmenter = get_predictor_and_segmenter(
+        model_type=model_type, checkpoint=checkpoint_path, device=device,
+        segmentation_mode="ais", is_tiled=tile_shape is not None,
     )
 
     input_paths = _get_inputs_from_paths(list(input_path), pattern)
@@ -521,6 +522,7 @@ def inference_segmentation(
             embedding_fpath = os.path.join(embedding_folder, f"{os.path.splitext(os.path.basename(path))[0]}.zarr")
 
         segmentation = automatic_instance_segmentation(
+            predictor=predictor,
             segmenter=segmenter,
             input_path=path,
             output_path=output_fpath,
@@ -580,17 +582,19 @@ def inference_tracking(
     for sparse or '--beta' for dense) can be passed through and are forwarded.
     """
     from .v2.util import DEFAULT_MODEL
-    from .v2.automatic_segmentation import get_segmenter, automatic_tracking
+    from .v2.automatic_segmentation import get_predictor_and_segmenter, automatic_tracking
 
     model_type = model_type or DEFAULT_MODEL
     tile_shape = _parse_shape(tile_shape)
     halo = _parse_shape(halo)
     generate_kwargs = _parse_extra(ctx.args)
 
-    segmenter = get_segmenter(
-        model_type=model_type, checkpoint=checkpoint_path, device=device, is_tiled=tile_shape is not None,
+    predictor, segmenter = get_predictor_and_segmenter(
+        model_type=model_type, checkpoint=checkpoint_path, device=device,
+        segmentation_mode="ais", is_tiled=tile_shape is not None,
     )
     automatic_tracking(
+        predictor=predictor,
         segmenter=segmenter,
         input_path=input_path,
         output_path=output_path,
@@ -791,28 +795,48 @@ def precompute_embeddings(
 
 
 @cli.command("train")
-@click.option("-c", "--config", required=True, help="The filepath to the SAM2 training config file.")
-@click.option("--use_cluster", type=int, default=None, help="Whether to launch on a cluster: 0 local, 1 cluster.")
-@click.option("--partition", default=None, help="SLURM partition.")
-@click.option("--account", default=None, help="SLURM account.")
-@click.option("--qos", default=None, help="SLURM qos.")
-@click.option("--num_gpus", type=int, default=None, help="Number of GPUs per node.")
-@click.option("--num_nodes", type=int, default=None, help="Number of nodes.")
-def train(config, use_cluster, partition, account, qos, num_gpus, num_nodes):
-    """Training a custom `micro-sam2` model."""
-    from .v2.train import train_sam2, register_omegaconf_resolvers
-
-    register_omegaconf_resolvers()
-    train_sam2(
-        config=config,
-        use_cluster=bool(use_cluster) if use_cluster is not None else None,
-        partition=partition,
-        account=account,
-        qos=qos,
-        num_gpus=num_gpus,
-        num_nodes=num_nodes,
+@click.option("-i", "--input_path", required=True, help="Root of the generalist training data (dataset subfolders).")
+@click.option("-m", "--model_type", default="hvit_t", help="SAM2 variant: hvit_t, hvit_s, hvit_b, hvit_l.")
+@click.option("-s", "--save_root", default=None, help="Directory for checkpoints and logs. Default: current dir.")
+@click.option("--name", default=None, help="Checkpoint/log folder name. Auto-generated if unset.")
+@click.option("--dataset_choice", default="both", help="Data modality to train on: 'lm', 'em' or 'both'.")
+@click.option("--n_epochs", type=int, default=100, help="Number of training epochs.")
+@click.option("--n_iterations", type=int, default=None, help="Fixed iteration budget (overrides --n_epochs).")
+@click.option("--batch_size", type=int, default=1, help="Batch size per GPU for 3d groups.")
+@click.option("--batch_size_2d", type=int, default=8, help="Batch size per GPU for 2d groups.")
+@click.option("--z_slices", type=int, multiple=True, default=(8,), help="Z-slice counts for 3d groups.")
+@click.option("-c", "--checkpoint_path", default=None, help="SAM2 checkpoint to start from; default downloaded.")
+def train(
+    input_path, model_type, save_root, name, dataset_choice, n_epochs, n_iterations,
+    batch_size, batch_size_2d, z_slices, checkpoint_path,
+):
+    """Train a joint SAM2 + UniSAM2 `micro-sam2` model on the generalist datasets."""
+    import torch
+
+    n_gpus = torch.cuda.device_count()
+    if name is None:
+        name = f"joint_sam2_{model_type}_{'multi' if n_gpus > 1 else 'single'}_gpu"
+
+    # Joint interactive + automatic recipe from finetuning/v2/generalist/train_joint.py.
+    kwargs = dict(
+        name=name, model_type=model_type, input_path=input_path, save_root=save_root,
+        dataset_choice=dataset_choice, n_epochs=n_epochs, n_iterations=n_iterations,
+        batch_size=batch_size, batch_size_2d=batch_size_2d, z_slices=list(z_slices),
+        n_workers=8, checkpoint_path=checkpoint_path, lr=1e-5, max_num_objects=5,
+        prob_to_use_pt_input=1.0, prob_to_use_box_input=0.5, num_frames_to_correct=2,
+        rand_frames_to_correct=True, prob_to_sample_from_gt=0.1, add_all_frames_to_correct_as_cond=True,
+        num_correction_pt_per_frame=7, num_init_cond_frames=2, clip_grad_norm=None, largest_first=True,
+        bidirectional=True, use_focal_loss=True, focal_weight=1.0, use_object_score_loss=True,
+        average_over_frames=False,
     )
 
+    if n_gpus > 1:
+        from .v2.training import train_joint_sam2_multi_gpu
+        train_joint_sam2_multi_gpu(**kwargs)
+    else:
+        from .v2.training import train_joint_sam2
+        train_joint_sam2(**kwargs)
+
 
 @cli.command("info")
 @click.option(

micro_sam/precompute_state.py

@@ -449,7 +449,7 @@ def _cache_amg_state_v2(
 def _cache_amg_slice(segmenter, save_path, i, init_fn, embedding_signature=None):
     """Load slice `i`'s AMG state from `save_path` if present and matching, else init and save.
 
-    Used by `micro_sam.v2.instance_segmentation.automatic_3d_segmentation` to cache the per-slice
+    Used by `micro_sam.v2.instance_segmentation.amg_3d_segmentation` to cache the per-slice
     grid-prediction state of a volume. `init_fn(i)` runs the (expensive) `initialize` for the slice.
     """
     key = _autoseg_state_key(i)
@@ -553,7 +553,7 @@ def _cache_ais_state_v2(
     'save_path=None' to skip caching.
 
     Args:
-        decoder: The UniSAM2 model, loaded via `micro_sam.v2.automatic_segmentation.get_unisam2_model`.
+        decoder: The UniSAM2 model, loaded via `micro_sam.v2.instance_segmentation.get_unisam2_model`.
         raw: The image data.
         image_embeddings: The (optionally precomputed) image embeddings. When given only the decoder
             is run on them (no encoder pass).
@@ -576,7 +576,7 @@ def _cache_ais_state_v2(
     Returns:
         The AIS segmenter with the (cached or freshly computed) state set.
     """
-    from .v2.automatic_segmentation import get_unisam2_segmentation_generator
+    from .v2.instance_segmentation import get_unisam2_segmentation_generator
 
     if is_tiled is None:
         is_tiled = image_embeddings is not None and image_embeddings.get("input_size") is None
@@ -612,12 +612,12 @@ def _cache_ais_state_v2(
 def _resolve_unisam2_decoder(model_type, checkpoint_path, device):
     """Return a UniSAM2 decoder for the SAM2 model if one is available, else None (fall back to AMG).
 
-    Mirrors `micro_sam.v2.automatic_segmentation.get_segmenter`: a decoder from a custom
+    Mirrors `micro_sam.v2.instance_segmentation.get_decoder`: a decoder from a custom
     `checkpoint_path`, or the registered decoder of a finetuned model (e.g. 'hvit_t_cells'). Any
     failure (e.g. an interactive-only checkpoint without a decoder) returns None.
     """
     from .v2.util import FINETUNED_MODELS, has_registered_decoder, _download_finetuned_sam2_model
-    from .v2.automatic_segmentation import get_unisam2_model
+    from .v2.instance_segmentation import get_unisam2_model
 
     encoder = model_type[:6]
     if checkpoint_path is not None:

micro_sam/sam_annotator/_state.py

@@ -201,7 +201,7 @@ def initialize_predictor(
         # finetuned checkpoint passed via 'checkpoint_path', or a finetuned model from the download
         # console (e.g. 'hvit_t_cells') whose decoder is downloaded from the SAM2 model registry.
         if self.is_sam2 and prefer_decoder and self.decoder is None:
-            from micro_sam.v2.automatic_segmentation import get_unisam2_model
+            from micro_sam.v2.instance_segmentation import get_unisam2_model
             from micro_sam.v2.util import FINETUNED_MODELS, _download_finetuned_sam2_model
 
             # The decoder is built on the base SAM2 backbone, i.e. the first 6 characters of the

micro_sam/sam_annotator/_widgets.py

@@ -4322,7 +4322,7 @@ def _run_unisam2(self, state, run_raw, ndim, z, pbar_init=None, pbar_update=None
         return self._segmenter.generate(mode=self.mode, **self._postproc_kwargs())
 
     def _run_amg(self, state, run_raw, ndim, z, pbar_init=None, pbar_update=None):
-        from micro_sam.v2.instance_segmentation import get_amg_segmenter, automatic_3d_segmentation
+        from micro_sam.v2.instance_segmentation import get_amg_segmenter, amg_3d_segmentation
         from micro_sam.precompute_state import cache_autoseg_state
 
         # The SAM2 model: 'state.predictor' is the image predictor (2d) wrapping the model, or the
@@ -4342,7 +4342,7 @@ def _run_amg(self, state, run_raw, ndim, z, pbar_init=None, pbar_update=None):
             segmenter = get_amg_segmenter(model, is_tiled=tile_shape is not None, model_type=model_type, **amg_params)
             # Reuse the precomputed 3d embeddings per slice (tiled or not) so AMG does not re-encode,
             # and cache each slice's grid-prediction state in the embedding Zarr.
-            return automatic_3d_segmentation(
+            return amg_3d_segmentation(
                 run_raw, segmenter, tile_shape=tile_shape, halo=halo,
                 image_embeddings=state.image_embeddings, state_save_path=save_path,
                 pbar_init=pbar_init, pbar_update=pbar_update, **generate_kwargs,