micro_sam.bioimageio.model_export

  1import os
  2import shutil
  3import tempfile
  4from pathlib import Path
  5from typing import Optional, Union
  6
  7import xarray
  8import numpy as np
  9import matplotlib.pyplot as plt
 10
 11import torch
 12
 13import bioimageio.core
 14import bioimageio.spec.model.v0_5 as spec
 15from bioimageio.spec import save_bioimageio_package
 16from bioimageio.core.digest_spec import create_sample_for_model
 17
 18from .. import util
 19from ..prompt_generators import PointAndBoxPromptGenerator
 20from ..evaluation.model_comparison import _enhance_image, _overlay_outline, _overlay_box
 21from ..prompt_based_segmentation import _compute_logits_from_mask
 22from .predictor_adaptor import PredictorAdaptor
 23
 24
 25DEFAULTS = {
 26    "authors": [
 27        spec.Author(name="Anwai Archit", affiliation="University Goettingen", github_user="anwai98"),
 28        spec.Author(name="Constantin Pape", affiliation="University Goettingen", github_user="constantinpape"),
 29    ],
 30    "description": "Finetuned Segment Anything Model for Microscopy",
 31    "cite": [
 32        spec.CiteEntry(
 33            text="Archit et al. Segment Anything for Microscopy",
 34            doi=spec.Doi("10.1038/s41592-024-02580-4")
 35        ),
 36    ],
 37    "tags": ["segment-anything", "instance-segmentation"],
 38}
 39
 40# Reference: https://github.com/bioimage-io/spec-bioimage-io/commit/39d343681d427ec93cf69eef7597d9eb9678deb1#diff-0bbdaa8196fa31f945afabcf04a4295ff098f1f24400ef9e59b0f684d411905eL269  # noqa
 41# We had this parameter in bioimageio.spec. This has been removed. We just make a copy of the same parameter.
 42ARBITRARY_SIZE = spec.ParameterizedSize(min=1, step=1)
 43
 44
 45def _get_architecture_model_type(model_type):
 46    # Derived models use their base SAM architecture.
 47    return model_type[:5]
 48
 49
 50def _create_test_inputs_and_outputs(image, labels, model_type, checkpoint_path, tmp_dir, with_decoder):
 51
 52    predictor = PredictorAdaptor(model_type=_get_architecture_model_type(model_type))
 53    predictor.load_state_dict(torch.load(checkpoint_path, map_location="cpu", weights_only=True))
 54    # Match the evaluation mode that BioImageIO uses.
 55    predictor.eval()
 56
 57    input_ = util._to_image(image).transpose(2, 0, 1)[None]
 58    image_path = os.path.join(tmp_dir, "input.npy")
 59    np.save(image_path, input_)
 60
 61    if with_decoder:
 62        # Test automatic instance segmentation through prompt-free inference.
 63        with torch.no_grad():
 64            masks, scores, embeddings = predictor(image=torch.from_numpy(input_))
 65        if masks.shape[1] == 0:
 66            raise RuntimeError(
 67                "Automatic instance segmentation did not find any objects in the test image. "
 68                "The exported model would fail its test procedure, which requires at least one object "
 69                "in the mask output. Please use a test image in which the model finds instances."
 70            )
 71        inputs = {"image": image_path}
 72    else:
 73        # For now we just generate a single box prompt here, but we could also generate more input prompts.
 74        generator = PointAndBoxPromptGenerator(
 75            n_positive_points=1,
 76            n_negative_points=2,
 77            dilation_strength=2,
 78            get_point_prompts=True,
 79            get_box_prompts=True,
 80        )
 81        centers, bounding_boxes = util.get_centers_and_bounding_boxes(labels)
 82        masks = util.segmentation_to_one_hot(labels.astype("int64"), segmentation_ids=[1, 2])  # type: ignore
 83        point_prompts, point_labels, box_prompts, _ = generator(masks, [bounding_boxes[1], bounding_boxes[2]])
 84
 85        box_prompts = box_prompts.numpy()[None]
 86        point_prompts = point_prompts.numpy()[None]
 87        point_labels = point_labels.numpy()[None]
 88
 89        # Generate logits from the two
 90        mask_prompts = np.stack(
 91            [_compute_logits_from_mask(labels == 1), _compute_logits_from_mask(labels == 2)]
 92        )[None]
 93
 94        with torch.no_grad():
 95            masks, scores, embeddings = predictor(
 96                image=torch.from_numpy(input_),
 97                embeddings=None,
 98                box_prompts=torch.from_numpy(box_prompts),
 99                point_prompts=torch.from_numpy(point_prompts),
100                point_labels=torch.from_numpy(point_labels),
101                mask_prompts=torch.from_numpy(mask_prompts),
102            )
103
104        box_prompt_path = os.path.join(tmp_dir, "box_prompts.npy")
105        point_prompt_path = os.path.join(tmp_dir, "point_prompts.npy")
106        point_label_path = os.path.join(tmp_dir, "point_labels.npy")
107        mask_prompt_path = os.path.join(tmp_dir, "mask_prompts.npy")
108        np.save(box_prompt_path, box_prompts.astype("int64"))
109        np.save(point_prompt_path, point_prompts)
110        np.save(point_label_path, point_labels)
111        np.save(mask_prompt_path, mask_prompts)
112
113        inputs = {
114            "image": image_path,
115            "box_prompts": box_prompt_path,
116            "point_prompts": point_prompt_path,
117            "point_labels": point_label_path,
118            "mask_prompts": mask_prompt_path,
119        }
120
121    mask_path = os.path.join(tmp_dir, "mask.npy")
122    score_path = os.path.join(tmp_dir, "scores.npy")
123    embed_path = os.path.join(tmp_dir, "embeddings.npy")
124    np.save(mask_path, masks.numpy())
125    np.save(score_path, scores.numpy())
126    np.save(embed_path, embeddings.numpy())
127
128    outputs = {"mask": mask_path, "score": score_path, "embeddings": embed_path}
129    return inputs, outputs
130
131
132def _write_documentation(doc, model_type, tmp_dir):
133    tmp_doc_path = os.path.join(tmp_dir, "documentation.md")
134
135    if doc is None:
136        with open(tmp_doc_path, "w") as f:
137            f.write("# Segment Anything for Microscopy\n")
138            f.write("We extend Segment Anything, a vision foundation model for image segmentation ")
139            f.write("by training specialized models for microscopy data.\n")
140        return tmp_doc_path
141
142    elif os.path.exists(doc):
143        return doc
144
145    else:
146        with open(tmp_doc_path, "w") as f:
147            f.write(doc)
148        return tmp_doc_path
149
150
151def _get_checkpoint(model_type, checkpoint_path, tmp_dir):
152    # If we don't have a checkpoint we get the corresponding model from the registry.
153    if checkpoint_path is None:
154        model_registry = util.models()
155        checkpoint_path = model_registry.fetch(model_type)
156        decoder_name = f"{model_type}_decoder"
157        decoder_path = model_registry.fetch(decoder_name) if decoder_name in model_registry.registry else None
158        return checkpoint_path, decoder_path
159
160    # Otherwise we have to load the checkpoint to see if it is the state dict of an encoder,
161    # or the checkpoint for a custom SAM model.
162    state, model_state = util._load_checkpoint(checkpoint_path)
163
164    if "model_state" in state:  # This is a finetuning checkpoint -> we have to resave the state.
165        new_checkpoint_path = os.path.join(tmp_dir, f"{model_type}.pt")
166        torch.save(model_state, new_checkpoint_path)
167
168        # We may also have an instance segmentation decoder in that case.
169        # If we have it we also resave this one and return it.
170        if "decoder_state" in state:
171            decoder_path = os.path.join(tmp_dir, f"{model_type}_decoder.pt")
172            decoder_state = state["decoder_state"]
173            torch.save(decoder_state, decoder_path)
174        else:
175            decoder_path = None
176
177        return new_checkpoint_path, decoder_path
178
179    else:  # This is a SAM encoder state -> we don't have to resave.
180        return checkpoint_path, None
181
182
183def _get_weight_checkpoint(model_type, checkpoint_path, decoder_path, tmp_dir):
184    """Combine the SAM and decoder states for the exported architecture."""
185    if decoder_path is None:
186        return checkpoint_path
187
188    weight_dir = os.path.join(tmp_dir, "combined_weights")
189    os.makedirs(weight_dir, exist_ok=True)
190    weight_path = os.path.join(weight_dir, f"{model_type}.pt")
191
192    model_state = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
193    decoder_state = torch.load(decoder_path, map_location="cpu", weights_only=True)
194    torch.save({"model_state": model_state, "decoder_state": decoder_state}, weight_path)
195    return weight_path
196
197
198def _get_decoder_attachment(model_type, decoder_path, tmp_dir):
199    # Registry decoder files do not have an extension.
200    attachment_name = f"{_get_architecture_model_type(model_type)}_decoder.pt"
201    if os.path.basename(decoder_path) == attachment_name:
202        return decoder_path
203    attachment_path = os.path.join(tmp_dir, attachment_name)
204    shutil.copyfile(decoder_path, attachment_path)
205    return attachment_path
206
207
208# TODO: Update this with our latest yaml file updates.
209def _write_dependencies(dependency_file, require_mobile_sam, require_micro_sam):
210    if require_micro_sam:
211        content = """name: sam
212channels:
213    - conda-forge
214dependencies:
215    - micro_sam"""
216    else:
217        content = """name: sam
218channels:
219    - pytorch
220    - conda-forge
221dependencies:
222    - segment-anything"""
223    if require_mobile_sam:
224        content += """
225    - timm
226    - pip:
227        - git+https://github.com/ChaoningZhang/MobileSAM.git"""
228    with open(dependency_file, "w") as f:
229        f.write(content)
230
231
232def _generate_covers(input_paths, result_paths, tmp_dir, with_decoder):
233    image = np.load(input_paths["image"]).squeeze()
234    mask = np.load(result_paths["mask"])
235
236    # create the image overlay
237    if image.ndim == 2:
238        overlay = np.stack([image, image, image]).transpose((1, 2, 0))
239    elif image.shape[0] == 3:
240        overlay = image.transpose((1, 2, 0))
241    else:
242        overlay = image
243    overlay = _enhance_image(overlay.astype("float32"))
244
245    if with_decoder:
246        # Outline all automatic instances.
247        overlay = _overlay_outline(overlay, mask[0, :, 0].max(axis=0), outline_dilation=2)
248    else:
249        # overlay the mask as outline
250        overlay = _overlay_outline(overlay, mask[0, 0, 0], outline_dilation=2)
251
252        # overlay the bounding box prompt
253        prompts = np.load(input_paths["box_prompts"])
254        prompt = prompts[0, 0][[1, 0, 3, 2]]
255        prompt = np.array([prompt[:2], prompt[2:]])
256        overlay = _overlay_box(overlay, prompt, outline_dilation=4)
257
258    # write  the cover image
259    fig, ax = plt.subplots(1)
260    ax.axis("off")
261    ax.imshow(overlay.astype("uint8"))
262    cover_path = os.path.join(tmp_dir, "cover.jpeg")
263    plt.savefig(cover_path, bbox_inches="tight")
264    plt.close()
265
266    covers = [cover_path]
267    return covers
268
269
270def _check_model(model_description, input_paths, result_paths, with_decoder):
271    image = xarray.DataArray(np.load(input_paths["image"]), dims=("batch", "channel", "y", "x"))
272
273    if with_decoder:
274        # Check automatic instance segmentation against its reference outputs.
275        mask = np.load(result_paths["mask"])
276        scores = np.load(result_paths["score"])
277        embeddings = np.load(result_paths["embeddings"])
278
279        with bioimageio.core.create_prediction_pipeline(model_description, devices=["cpu"]) as pp:
280            sample = create_sample_for_model(model=model_description, inputs={"image": image})
281            prediction = pp.predict_sample_without_blocking(sample)
282            assert np.allclose(mask, prediction.members["masks"].data)
283            assert np.allclose(scores, prediction.members["scores"].data)
284            assert np.allclose(embeddings, prediction.members["embeddings"].data)
285        return
286
287    # Load inputs.
288    embeddings = xarray.DataArray(np.load(result_paths["embeddings"]), dims=("batch", "channel", "y", "x"))
289    box_prompts = xarray.DataArray(np.load(input_paths["box_prompts"]), dims=("batch", "object", "channel"))
290    point_prompts = xarray.DataArray(
291        np.load(input_paths["point_prompts"]), dims=("batch", "object", "point", "channel")
292    )
293    point_labels = xarray.DataArray(np.load(input_paths["point_labels"]), dims=("batch", "object", "point"))
294    mask_prompts = xarray.DataArray(np.load(input_paths["mask_prompts"]), dims=("batch", "object", "channel", "y", "x"))
295
296    # Load outputs.
297    mask = np.load(result_paths["mask"])
298
299    # Match the device used to generate the reference outputs.
300    with bioimageio.core.create_prediction_pipeline(model_description, devices=["cpu"]) as pp:
301
302        # Check with all prompts. We only check the result for this setting,
303        # because this was used to generate the test data.
304        sample = create_sample_for_model(
305            model=model_description,
306            inputs={
307                "image": image,
308                "box_prompts": box_prompts,
309                "point_prompts": point_prompts,
310                "point_labels": point_labels,
311                "mask_prompts": mask_prompts,
312                "embeddings": embeddings,
313            },
314        )
315        prediction = pp.predict_sample_without_blocking(sample)
316
317        predicted_mask = prediction.members["masks"].data
318        assert predicted_mask.shape == mask.shape
319        assert np.allclose(mask, predicted_mask)
320
321        # Run the checks with partial prompts.
322        prompt_kwargs = [
323            # With boxes.
324            {"box_prompts": box_prompts},
325            # With point prompts.
326            {"point_prompts": point_prompts, "point_labels": point_labels},
327            # With masks.
328            {"mask_prompts": mask_prompts},
329            # With boxes and points.
330            {"box_prompts": box_prompts, "point_prompts": point_prompts, "point_labels": point_labels},
331            # With boxes and masks.
332            {"box_prompts": box_prompts, "mask_prompts": mask_prompts},
333            # With points and masks.
334            {"mask_prompts": mask_prompts, "point_prompts": point_prompts, "point_labels": point_labels},
335        ]
336
337        for kwargs in prompt_kwargs:
338            sample = create_sample_for_model(
339                model=model_description, inputs={"image": image, "embeddings": embeddings, **kwargs},
340            )
341            prediction = pp.predict_sample_without_blocking(sample)
342            predicted_mask = prediction.members["masks"].data
343            assert predicted_mask.shape == mask.shape
344
345    # Use a fresh pipeline to verify image-only inference.
346    with bioimageio.core.create_prediction_pipeline(model_description, devices=["cpu"]) as pp:
347        sample = create_sample_for_model(model=model_description, inputs={"image": image})
348        prediction = pp.predict_sample_without_blocking(sample)
349        predicted_mask = prediction.members["masks"].data
350        # AIS may return no objects.
351        assert predicted_mask.ndim == 5
352        assert predicted_mask.shape[-2:] == mask.shape[-2:]
353
354
355def _regenerate_reference_outputs(model_description, input_paths, result_paths):
356    """Regenerate reference outputs through the packaged model.
357
358    This avoids numerical differences between direct calls and BioImageIO calls.
359    """
360    from bioimageio.core.digest_spec import get_test_input_sample
361
362    sample = get_test_input_sample(model_description)
363    # Keep reference generation deterministic across CI runners.
364    with bioimageio.core.create_prediction_pipeline(model_description, devices=["cpu"]) as pp:
365        pp.apply_preprocessing(sample)
366        prediction = pp.predict_sample_without_blocking(
367            sample,
368            skip_preprocessing=True,
369            skip_postprocessing=True,
370            skip_input_padding=True,
371            skip_output_cropping=True,
372        )
373        pp.apply_postprocessing(prediction)
374    np.save(result_paths["mask"], np.asarray(prediction.members["masks"].data).astype("uint8"))
375    np.save(result_paths["score"], np.asarray(prediction.members["scores"].data).astype("float32"))
376    np.save(result_paths["embeddings"], np.asarray(prediction.members["embeddings"].data).astype("float32"))
377
378
379def _build_model_description(
380    name, input_paths, result_paths, weight_descriptions, doc_path, covers, extra_kwargs, with_decoder, **kwargs
381):
382    # Rebuild after reference generation to update the test tensor hashes.
383    if with_decoder:
384        # Prevent BioImageIO size probes from cropping instances.
385        test_image_shape = np.load(input_paths["image"], mmap_mode="r").shape
386        image_y_size = spec.ParameterizedSize(min=test_image_shape[-2], step=1)
387        image_x_size = spec.ParameterizedSize(min=test_image_shape[-1], step=1)
388    else:
389        image_y_size = image_x_size = ARBITRARY_SIZE
390
391    input_descriptions = [
392        # First input: the image data.
393        spec.InputTensorDescr(
394            id=spec.TensorId("image"),
395            axes=[
396                spec.BatchAxis(size=1),
397                # NOTE: to support 1 and 3 channels we can add another preprocessing.
398                # Best solution: Have a pre-processing for this! (1C -> RGB)
399                spec.ChannelAxis(channel_names=[spec.Identifier(cname) for cname in "RGB"]),
400                spec.SpaceInputAxis(id=spec.AxisId("y"), size=image_y_size),
401                spec.SpaceInputAxis(id=spec.AxisId("x"), size=image_x_size),
402            ],
403            test_tensor=spec.FileDescr(source=input_paths["image"]),
404            data=spec.IntervalOrRatioDataDescr(type="uint8")
405        ),
406    ]
407
408    # Decoder exports keep only the image input to test automatic instance segmentation.
409    if not with_decoder:
410        input_descriptions += [
411            # Second input: the box prompts (optional)
412            spec.InputTensorDescr(
413                id=spec.TensorId("box_prompts"),
414                optional=True,
415                axes=[
416                    spec.BatchAxis(size=1),
417                    spec.IndexInputAxis(
418                        id=spec.AxisId("object"),
419                        size=ARBITRARY_SIZE
420                    ),
421                    spec.ChannelAxis(channel_names=[spec.Identifier(bname) for bname in "hwxy"]),
422                ],
423                test_tensor=spec.FileDescr(source=input_paths["box_prompts"]),
424                data=spec.IntervalOrRatioDataDescr(type="int64")
425            ),
426
427            # Third input: the point prompt coordinates (optional)
428            spec.InputTensorDescr(
429                id=spec.TensorId("point_prompts"),
430                optional=True,
431                axes=[
432                    spec.BatchAxis(size=1),
433                    spec.IndexInputAxis(
434                        id=spec.AxisId("object"),
435                        size=ARBITRARY_SIZE
436                    ),
437                    spec.IndexInputAxis(
438                        id=spec.AxisId("point"),
439                        size=ARBITRARY_SIZE
440                    ),
441                    spec.ChannelAxis(channel_names=[spec.Identifier(bname) for bname in "xy"]),
442                ],
443                test_tensor=spec.FileDescr(source=input_paths["point_prompts"]),
444                data=spec.IntervalOrRatioDataDescr(type="int64")
445            ),
446
447            # Fourth input: the point prompt labels (optional)
448            spec.InputTensorDescr(
449                id=spec.TensorId("point_labels"),
450                optional=True,
451                axes=[
452                    spec.BatchAxis(size=1),
453                    spec.IndexInputAxis(
454                        id=spec.AxisId("object"),
455                        size=ARBITRARY_SIZE
456                    ),
457                    spec.IndexInputAxis(
458                        id=spec.AxisId("point"),
459                        size=ARBITRARY_SIZE
460                    ),
461                ],
462                test_tensor=spec.FileDescr(source=input_paths["point_labels"]),
463                data=spec.IntervalOrRatioDataDescr(type="int64")
464            ),
465
466            # Fifth input: the mask prompts (optional)
467            spec.InputTensorDescr(
468                id=spec.TensorId("mask_prompts"),
469                optional=True,
470                axes=[
471                    spec.BatchAxis(size=1),
472                    spec.IndexInputAxis(
473                        id=spec.AxisId("object"),
474                        size=ARBITRARY_SIZE
475                    ),
476                    spec.ChannelAxis(channel_names=["channel"]),
477                    spec.SpaceInputAxis(id=spec.AxisId("y"), size=256),
478                    spec.SpaceInputAxis(id=spec.AxisId("x"), size=256),
479                ],
480                test_tensor=spec.FileDescr(source=input_paths["mask_prompts"]),
481                data=spec.IntervalOrRatioDataDescr(type="float32")
482            ),
483
484            # Sixth input: the image embeddings (optional)
485            spec.InputTensorDescr(
486                id=spec.TensorId("embeddings"),
487                optional=True,
488                axes=[
489                    spec.BatchAxis(size=1),
490                    # NOTE: we currently have to specify all the channel names
491                    # (It would be nice to also support size)
492                    spec.ChannelAxis(channel_names=[spec.Identifier(f"c{i}") for i in range(256)]),
493                    spec.SpaceInputAxis(id=spec.AxisId("y"), size=64),
494                    spec.SpaceInputAxis(id=spec.AxisId("x"), size=64),
495                ],
496                test_tensor=spec.FileDescr(source=result_paths["embeddings"]),
497                data=spec.IntervalOrRatioDataDescr(type="float32")
498            ),
499        ]
500
501    output_descriptions = [
502        # First output: The mask predictions.
503        spec.OutputTensorDescr(
504            id=spec.TensorId("masks"),
505            axes=[
506                spec.BatchAxis(size=1),
507                # NOTE: we use the data dependent size here to avoid dependency on optional inputs
508                spec.IndexOutputAxis(
509                    id=spec.AxisId("object"), size=spec.DataDependentSize(),
510                ),
511                # NOTE: this could be a 3 once we use multi-masking
512                spec.ChannelAxis(channel_names=[spec.Identifier("mask")]),
513                spec.SpaceOutputAxis(
514                    id=spec.AxisId("y"),
515                    size=spec.SizeReference(
516                        tensor_id=spec.TensorId("image"), axis_id=spec.AxisId("y"),
517                    )
518                ),
519                spec.SpaceOutputAxis(
520                    id=spec.AxisId("x"),
521                    size=spec.SizeReference(
522                        tensor_id=spec.TensorId("image"), axis_id=spec.AxisId("x"),
523                    )
524                )
525            ],
526            data=spec.IntervalOrRatioDataDescr(type="uint8"),
527            test_tensor=spec.FileDescr(source=result_paths["mask"])
528        ),
529
530        # The score predictions
531        spec.OutputTensorDescr(
532            id=spec.TensorId("scores"),
533            axes=[
534                spec.BatchAxis(size=1),
535                # NOTE: we use the data dependent size here to avoid dependency on optional inputs
536                spec.IndexOutputAxis(
537                    id=spec.AxisId("object"), size=spec.DataDependentSize(),
538                ),
539                # NOTE: this could be a 3 once we use multi-masking
540                spec.ChannelAxis(channel_names=[spec.Identifier("mask")]),
541            ],
542            data=spec.IntervalOrRatioDataDescr(type="float32"),
543            test_tensor=spec.FileDescr(source=result_paths["score"])
544        ),
545
546        # The image embeddings
547        spec.OutputTensorDescr(
548            id=spec.TensorId("embeddings"),
549            axes=[
550                spec.BatchAxis(size=1),
551                spec.ChannelAxis(channel_names=[spec.Identifier(f"c{i}") for i in range(256)]),
552                spec.SpaceOutputAxis(id=spec.AxisId("y"), size=64),
553                spec.SpaceOutputAxis(id=spec.AxisId("x"), size=64),
554            ],
555            data=spec.IntervalOrRatioDataDescr(type="float32"),
556            test_tensor=spec.FileDescr(source=result_paths["embeddings"])
557        )
558    ]
559
560    return spec.ModelDescr(
561        name=name,
562        inputs=input_descriptions,
563        outputs=output_descriptions,
564        weights=weight_descriptions,
565        description=kwargs.get("description", DEFAULTS["description"]),
566        authors=kwargs.get("authors", DEFAULTS["authors"]),
567        cite=kwargs.get("cite", DEFAULTS["cite"]),
568        license=spec.LicenseId("CC-BY-4.0"),
569        documentation=spec.FileDescr(source=Path(doc_path)),
570        git_repo=spec.HttpUrl("https://github.com/computational-cell-analytics/micro-sam"),
571        tags=kwargs.get("tags", DEFAULTS["tags"]),
572        covers=[spec.FileDescr(source=Path(cover)) for cover in covers],
573        **extra_kwargs,
574        config=spec.Config(
575            bioimageio=spec.BioimageioConfig(
576                # Allow minor numerical differences in raw embeddings across devices.
577                reproducibility_tolerance=[
578                    spec.ReproducibilityTolerance(
579                        mismatched_elements_per_million=1000,
580                        output_ids=[spec.TensorId("embeddings")],
581                    )
582                ],
583            ),
584        ),
585    )
586
587
588def export_sam_model(
589    image: np.ndarray,
590    label_image: np.ndarray,
591    model_type: str,
592    name: str,
593    output_path: Union[str, os.PathLike],
594    checkpoint_path: Optional[Union[str, os.PathLike]] = None,
595    **kwargs
596) -> None:
597    """Export SAM model to BioImage.IO model format.
598
599    The exported model can be uploaded to [bioimage.io](https://bioimage.io/#/) and
600    be used in tools that support the BioImage.IO model format.
601
602    Models with a decoder expose only the image input.
603    Their test outputs use automatic instance segmentation.
604    Use micro_sam directly for prompted segmentation with these models.
605    Models without a decoder expose the interactive prompt inputs.
606
607    Args:
608        image: The image for generating test data.
609        label_image: The segmentation corresponding to `image`.
610            It is used to derive prompt inputs for the model.
611        model_type: The type of the SAM model.
612        name: The name of the exported model.
613        output_path: Where the exported model is saved.
614        checkpoint_path: Optional checkpoint for loading the SAM model.
615    """
616    with tempfile.TemporaryDirectory() as tmp_dir:
617        checkpoint_path, decoder_path = _get_checkpoint(model_type, checkpoint_path, tmp_dir)
618        with_decoder = decoder_path is not None
619        weight_path = _get_weight_checkpoint(model_type, checkpoint_path, decoder_path, tmp_dir)
620        input_paths, result_paths = _create_test_inputs_and_outputs(
621            image, label_image, model_type, weight_path, tmp_dir, with_decoder=with_decoder,
622        )
623        architecture_path = os.path.join(os.path.split(__file__)[0], "predictor_adaptor.py")
624        architecture = spec.ArchitectureFromFileDescr(
625            source=Path(architecture_path),
626            callable="PredictorAdaptor",
627            kwargs={"model_type": _get_architecture_model_type(model_type)}
628        )
629
630        dependency_file = os.path.join(tmp_dir, "environment.yaml")
631        _write_dependencies(
632            dependency_file, require_mobile_sam=model_type.startswith("vit_t"), require_micro_sam=with_decoder,
633        )
634
635        weight_descriptions = spec.WeightsDescr(
636            pytorch_state_dict=spec.PytorchStateDictWeightsDescr(
637                source=Path(weight_path),
638                architecture=architecture,
639                pytorch_version=spec.Version(torch.__version__),
640                dependencies=spec.FileDescr(source=dependency_file),
641            )
642        )
643
644        doc_path = _write_documentation(kwargs.get("documentation", None), model_type, tmp_dir)
645
646        covers = kwargs.pop("covers", None)
647        if covers is None:
648            covers = _generate_covers(input_paths, result_paths, tmp_dir, with_decoder=with_decoder)
649        else:
650            assert all(os.path.exists(cov) for cov in covers)
651
652        # the uploader information is only added if explicitly passed
653        extra_kwargs = {}
654        if "id" in kwargs:
655            extra_kwargs["id"] = kwargs["id"]
656        if "id_emoji" in kwargs:
657            extra_kwargs["id_emoji"] = kwargs["id_emoji"]
658        if "uploader" in kwargs:
659            extra_kwargs["uploader"] = kwargs["uploader"]
660        if "version" in kwargs:
661            extra_kwargs["version"] = kwargs["version"]
662
663        if with_decoder:
664            # Keep the decoder available as a standalone registry attachment.
665            attachment_path = _get_decoder_attachment(model_type, decoder_path, tmp_dir)
666            extra_kwargs["attachments"] = [spec.FileDescr(source=attachment_path)]
667
668        model_description = _build_model_description(
669            name, input_paths, result_paths, weight_descriptions, doc_path, covers, extra_kwargs,
670            with_decoder=with_decoder, **kwargs,
671        )
672
673        if with_decoder:
674            # Rebuild the description after the reference output hashes change.
675            _regenerate_reference_outputs(model_description, input_paths, result_paths)
676            model_description = _build_model_description(
677                name, input_paths, result_paths, weight_descriptions, doc_path, covers, extra_kwargs,
678                with_decoder=with_decoder, **kwargs,
679            )
680
681        _check_model(model_description, input_paths, result_paths, with_decoder=with_decoder)
682
683        save_bioimageio_package(model_description, output_path=output_path)