Troubles using checkpoint file to headlessly run inference on unseen data (unpickling errors?)

[evaknichols]

Hi Anwai and Constantin! Hope you are both doing well this winter season. I have made some headway in finetuning a decent model, but am running into issues deploying it on other datasets. Below I describe my overall task and how I am trying to get there with microSAM.

Big Picture: from whole embryo LSFM data, predict the region where the somite structures are likely to be (semantic segmentation). The somite regions are sparse, about 10% of the total embryo (but my training data is dense; only planes that have masks are passed). Then, from that foreground, separate masks into instances. I tested code using notebooks but do the full training in a python scripts submitted to our cluster. The version of microSAM I'm using is 1.6.2.

Current progress: I have successfully fine-tuned a Micro-SAM vit_b_lm model for somite segmentation on 2D slices from volumetric mouse embryo data with great IOU scores (around 0.7-0.8). The results look fantastic (below), but are purely semantic (0=background and 1=foreground). I should be able to separate these into instances (which should be achievable by microSAM prompting, right?). I'd like to use this checkpoint to run inference on all unseen data (10 embryos) to extract foreground/semantic masks as a .tiff. Then that can be used as input for the prompting to get instance masks (I think...).

The first step is to run inference. I use to have notebook code working with an earlier version of microSAM, but now I get unpickling errors (below). I've brainstormed with ChatGPT but it is leading me out of the microSAM API. Its consensus is that the checkpoints file is not sufficient to run inference outside of the training environment. I feel like this cannot be since my purpose is straightforward. So, I'm probably missing something critical. How do I use a checkpoints file from training to run inference outside of the training environment? ChatGPT said that current functions in the API document are incompatible with my intended purpose, which makes me question if I'm using microSAM in a way that is "off label use". I could use a bit of help deciding a path forward with the checkpoint that I have, and if you think my approach is infeasible to get instance masks.

Thanks for reading!

The Code (ran as python script, but worked in a notebook in an earlier version of microSAM)

# imports
import zarr
import numpy as np
import torch
import matplotlib.pyplot as plt
import os 

def run_automatic_instance_segmentation(
    image: np.ndarray,
    checkpoint_path: Union[os.PathLike ,str],
    model_type: str = "vit_b_lm",
    device: Optional[Union[str, torch.device]] = None,
    tile_shape: Optional[Tuple[int, int]] = None,
    halo: Optional[Tuple[int, int]] = None,
):
    """Automatic Instance Segmentation (AIS) by training an additional instance decoder in SAM.

    NOTE: AIS is supported only for `µsam` models.

    Args:
        image: The input image.
        checkpoint_path: The path to stored checkpoints.
        model_type: The choice of the `µsam` model.
        device: The device to run the model inference.
        tile_shape: The tile shape for tiling-based segmentation.
        halo: The overlap shape on each side per tile for stitching the segmented tiles.

    Returns:
        The instance segmentation.
    """
    # Step 1: Get the 'predictor' and 'segmenter' to perform automatic instance segmentation.
    predictor, segmenter = get_predictor_and_segmenter(
        model_type=model_type,  # choice of the Segment Anything model
        checkpoint=checkpoint_path,  # overwrite to pass your own finetuned model.
        device=device,  # the device to run the model inference.
        is_tiled=(tile_shape is not None),  # whether to run automatic segmentation.
    )

    # Step 2: Get the instance segmentation for the given image.
    prediction = automatic_instance_segmentation(
        predictor=predictor,  # the predictor for the Segment Anything model.
        segmenter=segmenter,  # the segmenter class responsible for generating predictions.
        input_path=image,  # the filepath to image or the input array for automatic segmentation.
        ndim=2,  # the number of input dimensions.
        tile_shape=tile_shape,  # the tile shape for tiling-based prediction.
        halo=halo,  # the overlap shape for tiling-based prediction.
    )

    return prediction

# Path definition
raw_path = "/net/beliveau/vol2/instrument/E9.5_300/Zoom_300/dataset_fused_rechunked_blocks.zarr"
raw_key = "ch2/s3"

checkpoint_path = ("/net/beliveau/vol1/project/EKN_whole_embryo_structural_analysis/Somites/uSAM_finetuning/checkpoints/somite_usam_finetune_old_code_test_5kiter_1e-5lr_vit_b_LM/best.pt")

raw = zarr.open(raw_path, mode="r")[raw_key]
print(f"Loaded raw volume: {raw.shape}, dtype={raw.dtype}")

# defining padding function 
def center_pad(image, target_shape=(1024, 1024)):
    x, y = image.shape
    tx, ty = target_shape
    pad_x = max((tx - x) // 2, 0)
    pad_y = max((ty - y) // 2, 0)
    pad_width = ((pad_x, tx - x - pad_x), (pad_y, ty - y - pad_y))
    return np.pad(image, pad_width, mode="constant", constant_values=0)

# running inference on all the slices in the z-stack
device = "cuda" if torch.cuda.is_available() else "cpu"
n_slices = raw.shape[0]
pred_stack = []

print(f"Running inference on {n_slices} slices...")

for z in range(n_slices):
    raw_slice = raw[z, :, :].astype(np.float32)
    padded = center_pad(raw_slice)
    if padded.max() > 0:
        normed = (padded - padded.min()) / (padded.max() - padded.min())
    else:
        normed = np.zeros_like(padded)

    img_3ch = np.repeat(normed[:, :, None], 3, axis=2).astype(np.float32)

    pred_mask_bin = run_automatic_instance_segmentation(
        image=img_3ch,
        checkpoint_path=checkpoint_path,
        model_type="vit_b",
        device=device
    )
    pred_mask_bin = pred_mask_bin.astype(bool)
    pred_stack.append(pred_mask_bin)

pred_stack = np.stack(pred_stack)
print("Inference complete!")
print(f"Predicted stack shape: {pred_stack.shape}")

Error:

---------------------------------------------------------------------------
Traceback (most recent call last):
  File "/net/beliveau/vol1/project/EKN_whole_embryo_structural_analysis/Somites/uSAM_finetuning/microSAM_inference_to_instance_mask_script.py", line 54, in <module>
    sam, _ = get_sam_model(
             ~~~~~~~~~~~~~^
        model_type="vit_b_lm",
        ^^^^^^^^^^^^^^^^^^^^^^
    ...<2 lines>...
        return_sam=True
        ^^^^^^^^^^^^^^^
    )
    ^
  File "/net/beliveau/vol1/home/eknich/.conda/envs/sam/lib/python3.13/site-packages/micro_sam/util.py", line 425, in get_sam_model
    sam = custom_models.sam_model_registry[abbreviated_model_type](**model_kwargs)
  File "/net/beliveau/vol1/home/eknich/.conda/envs/sam/lib/python3.13/site-packages/micro_sam/models/build_sam.py", line 68, in build_sam_vit_b
    return _build_sam(
        encoder_embed_dim=768,
    ...<5 lines>...
        image_size=image_size,
    )
  File "/net/beliveau/vol1/home/eknich/.conda/envs/sam/lib/python3.13/site-packages/micro_sam/models/build_sam.py", line 139, in _build_sam
    state_dict = torch.load(f)
  File "/net/beliveau/vol1/home/eknich/.conda/envs/sam/lib/python3.13/site-packages/torch/serialization.py", line 1470, in load
    raise pickle.UnpicklingError(_get_wo_message(str(e))) from None
_pickle.UnpicklingError: Weights only load failed. This file can still be loaded, to do so you have two options, �[1mdo those steps only if you trust the source of the checkpoint�[0m. 
	(1) In PyTorch 2.6, we changed the default value of the `weights_only` argument in `torch.load` from `False` to `True`. Re-running `torch.load` with `weights_only` set to `False` will likely succeed, but it can result in arbitrary code execution. Do it only if you got the file from a trusted source.
	(2) Alternatively, to load with `weights_only=True` please check the recommended steps in the following error message.
	WeightsUnpickler error: Unsupported global: GLOBAL __main__.SomiteSegmentationDataset was not an allowed global by default. Please use `torch.serialization.add_safe_globals([SomiteSegmentationDataset])` or the `torch.serialization.safe_globals([SomiteSegmentationDataset])` context manager to allowlist this global if you trust this class/function.

Check the documentation of torch.load to learn more about types accepted by default with weights_only https://pytorch.org/docs/stable/generated/torch.load.html.

[constantinpape]

Hi Eva!

Nice to hear from you and to see the great update on your progress. Sorry for the bit late reply -- things were a bit busy before the holidays ;).

Regarding your overall approach

I should be able to separate these into instances (which should be achievable by microSAM prompting, right?). I'd like to use this checkpoint to run inference on all unseen data (10 embryos) to extract foreground/semantic masks as a .tiff. Then that can be used as input for the prompting to get instance masks (I think...).

This makes sense overall, but it may not be so trivial to derive prompts from the semantic masks for individual somites. If this turns out to be difficult, you would need to repeat the training with instance segmentation ground-truth, i.e. 1 label per individual somite structures. You can generate these with the interactive annotator if needed. Then, you can use both automatic as well as interactive segmentation with the newly trained model.

Regarding the current error

The problem is that the code runs this function which it should actually not run into. I am not 100% sure what happened here, can you check the following:

1. Can you check what happens if you just load the model? Does the code below run through? If no: what's the error?

from micro_sam.util import get_sam_model
checkpoint_path = "..."  # same path as before
model, state = get_sam_model(model_type="vit_b", checkpoint_path=checkpoint_path, return_state=True)

2. Did you train with any options for parameter-efficient fine-tuning? (If you don't know, you can also share your training code).

[evaknichols]

Hi Constantin,

No worries at all, I didn't expect to hear back before the holidays! Thanks so much. I have new information for you that could provide some clarity. It looks like I'm getting instance masks now and confirm that inference works so long as it's run after finetuning directly in the same training environment. So apart from training improvements (like adding more data and augmentations) think I no longer need to do the whole prompting-from-masks thing any longer, and only focus on trying to get get inference to work outside of the training environment. I could definitely use a mini code check to see if I'm handling the finetuning and inference correctly for this overall goal? My training code is below, from finetuning to inference, all-in-one python script that works as-is. I'm not sure if I am using the params you are thinking of.

First, running this chunk (in a notebook outside of the training script):

from micro_sam.util import get_sam_model
checkpoint_path = "..."  # same path as before
model, state = get_sam_model(model_type="vit_b", checkpoint_path=checkpoint_path, return_state=True)

Produced this output, which is definitely familiar:

---------------------------------------------------------------------------
UnpicklingError                           Traceback (most recent call last)
Cell In[1], line 3
      1 from micro_sam.util import get_sam_model
      2 checkpoint_path = "/net/beliveau/vol1/project/EKN_whole_embryo_structural_analysis/Somites/uSAM_finetuning/checkpoints/somite_usam_finetune_old_code_test_5kiter_1e-5lr_vit_b_LM_inf_after_train_2/best.pt"  # same path as before
----> 3 model, state = get_sam_model(model_type="vit_b", checkpoint_path=checkpoint_path, return_state=True)

File ~/.conda/envs/sam/lib/python3.13/site-packages/micro_sam/util.py:399, in get_sam_model(model_type, device, checkpoint_path, return_sam, return_state, peft_kwargs, flexible_load_checkpoint, progress_bar_factory, **model_kwargs)
    393 if abbreviated_model_type == "vit_t" and not VIT_T_SUPPORT:
    394     raise RuntimeError(
    395         "'mobile_sam' is required for the vit-tiny. "
    396         "You can install it via 'pip install git+https://github.com/ChaoningZhang/MobileSAM.git'"
    397     )
--> 399 state, model_state = _load_checkpoint(checkpoint_path)
    401 if _provided_checkpoint_path:
    402     # To get the model weights, we prioritize having the correct 'checkpoint_path' over 'model_type'
    403     # It is done to avoid strange parameter mismatch issues while incompatible model type and weights combination.
    404     from micro_sam.models.build_sam import _validate_model_type

File ~/.conda/envs/sam/lib/python3.13/site-packages/micro_sam/util.py:275, in _load_checkpoint(checkpoint_path)
    272 custom_pickle = pickle
    273 custom_pickle.Unpickler = _CustomUnpickler
--> 275 state = torch.load(checkpoint_path, map_location="cpu", pickle_module=custom_pickle)
    276 if "model_state" in state:
    277     # Copy the model weights from torch_em's training format.
    278     model_state = state["model_state"]

File ~/.conda/envs/sam/lib/python3.13/site-packages/torch/serialization.py:1471, in load(f, map_location, pickle_module, weights_only, mmap, **pickle_load_args)
   1469             except pickle.UnpicklingError as e:
   1470                 raise pickle.UnpicklingError(_get_wo_message(str(e))) from None
-> 1471         return _load(
   1472             opened_zipfile,
   1473             map_location,
   1474             pickle_module,
   1475             overall_storage=overall_storage,
   1476             **pickle_load_args,
   1477         )
   1478 if mmap:
   1479     f_name = "" if not isinstance(f, str) else f"{f}, "

File ~/.conda/envs/sam/lib/python3.13/site-packages/torch/serialization.py:1964, in _load(zip_file, map_location, pickle_module, pickle_file, overall_storage, **pickle_load_args)
   1962 global _serialization_tls
   1963 _serialization_tls.map_location = map_location
-> 1964 result = unpickler.load()
   1965 _serialization_tls.map_location = None
   1967 torch._utils._validate_loaded_sparse_tensors()

UnpicklingError: NEWOBJ class argument must be a type, not NoneType

Training code .py

# Imports 
import os
import warnings
warnings.filterwarnings("ignore")
from glob import glob
from IPython.display import FileLink
from typing import Union, Tuple, Optional
import numpy as np
import random
import zarr
import tifffile
from PIL import Image
import imageio.v3 as imageio
from matplotlib import pyplot as plt
from skimage.measure import label as connected_components
import torch
from torch_em.util.debug import check_loader
from torch_em.data import MinInstanceSampler
from torch_em.util.util import get_random_colors
from torchvision import transforms
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
import micro_sam.training as sam_training
from micro_sam.training import train_sam
from micro_sam.training.util import get_trainable_sam_model
from micro_sam.sample_data import fetch_tracking_example_data, fetch_tracking_segmentation_data
from micro_sam.automatic_segmentation import get_predictor_and_segmenter, automatic_instance_segmentation

# -----------------------------
# Paths and keys
# -----------------------------
raw_path = "/net/beliveau/vol2/instrument/E9.5_306/Zoom_306/dataset_fused_rechunked_blocks.zarr"
raw_key = "ch2/s3"
label_path = "/net/beliveau/vol1/project/EKN_whole_embryo_structural_analysis/Somites/E9.5_306/306_all_somites_no_PSM.tif"

# Open Zarr dataset
raw = zarr.open(raw_path, mode="r")[raw_key]

# ---------- augmentation pipelines ----------
# geometric transforms (must be same for image & mask)
train_geo = transforms.Compose([
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.RandomVerticalFlip(p=0.5),
    transforms.RandomRotation(degrees=(-90, 90)),  # geometric only
])

# intensity transforms (apply only to image)
train_intensity = transforms.Compose([
    transforms.RandomAutocontrast(p=0.5),
    transforms.GaussianBlur(kernel_size=(3, 3), sigma=(0.1, 3.0)),
])

# validation: no random augmentations
val_geo = transforms.Compose([])          # no geometry for val
val_intensity = transforms.Compose([])    # no intensity for val

# ---------- helper: center pad (TBA: use median values instead of 0) ----------
def center_pad(image, target_shape=(1024, 1024)):
    x, y = image.shape
    tx, ty = target_shape
    pad_x = max((tx - x) // 2, 0)
    pad_y = max((ty - y) // 2, 0)
    pad_width = ((pad_x, tx - x - pad_x), (pad_y, ty - y - pad_y))
    return np.pad(image, pad_width, mode="constant", constant_values=0)

# ---------- Dataset with separate geo/intensity transforms ----------
class SomiteSegmentationDataset(Dataset):
    def __init__(
        self,
        raw_path,
        raw_key,
        mask_path,
        split="train",
        split_ratios=(0.7, 0.2, 0.1),
        seed=42,
        transform_geo=None,
        transform_intensity=None,
    ):
        super().__init__()
        self.raw_zarr = zarr.open(raw_path, mode="r")
        self.raw = self.raw_zarr[raw_key]
        self.mask = tifffile.imread(mask_path)
        self.transform_geo = transform_geo
        self.transform_intensity = transform_intensity

        # find non-empty slices
        non_empty_indices = [i for i in range(self.mask.shape[0]) if np.any(self.mask[i] > 0)]
        random.seed(seed); random.shuffle(non_empty_indices)
        n = len(non_empty_indices)
        n_train = int(split_ratios[0] * n)
        n_val = int(split_ratios[1] * n)

        if split == "train":
            self.indices = non_empty_indices[:n_train]
        elif split == "val":
            self.indices = non_empty_indices[n_train:n_train + n_val]
        elif split == "test":
            self.indices = non_empty_indices[n_train + n_val:]
        else:
            raise ValueError(f"Invalid split '{split}'")
        print(f"Loaded {len(self.indices)} slices for {split} split.")

    def __len__(self):
        return len(self.indices)

    def __getitem__(self, idx):
        slice_idx = self.indices[idx]
        raw_slice = self.raw[slice_idx].astype(np.float32)   # e.g. uint16 -> float32
        mask_slice = self.mask[slice_idx].astype(np.uint8)

        # center-pad
        raw_padded = center_pad(raw_slice)
        mask_padded = center_pad(mask_slice)

        # normalize raw to 0..255 (uint8) for PIL/torchvision behavior
        if raw_padded.max() > raw_padded.min():
            raw_norm = (raw_padded - raw_padded.min()) / (raw_padded.max() - raw_padded.min())
        else:
            raw_norm = np.zeros_like(raw_padded, dtype=np.float32)
        raw_uint8 = (raw_norm * 255).astype(np.uint8)

        # mask already binary: convert to 0/255 for PIL
        mask_bin_uint8 = (mask_padded > 0).astype(np.uint8) * 255

        raw_img = Image.fromarray(raw_uint8)       # single-channel grayscale PIL
        mask_img = Image.fromarray(mask_bin_uint8)

        # Apply geometric transforms identically:
        if self.transform_geo is not None:
            # Use a deterministic seed to make geo transforms repeatable for pair
            # torchvision transforms are not guaranteed to be fully seeded by random.seed,
            # but this approach works well enough for typical RandomRotation/Flip.
            seed = np.random.randint(2**31 - 1)
            random.seed(seed); torch.manual_seed(seed)
            raw_img = self.transform_geo(raw_img)
            random.seed(seed); torch.manual_seed(seed)
            mask_img = self.transform_geo(mask_img)

        # Apply intensity transforms only to the raw image
        if self.transform_intensity is not None:
            raw_img = self.transform_intensity(raw_img)

        # Convert to tensors (C,H,W), values in [0,1]
        to_tensor = transforms.ToTensor()
        raw_tensor = to_tensor(raw_img).float()    # shape (1,H,W)
        mask_tensor = to_tensor(mask_img).float()  # shape (1,H,W), values 0 or 1
        
        # Scale raw_tensor to [0, 255] for Micro-SAM
        raw_tensor = raw_tensor * 255.0

        # Build 4-channel mask for SAM decoder
        H, W = mask_tensor.shape[1], mask_tensor.shape[2]
        mask_4ch = torch.zeros((4, H, W), dtype=torch.float32)
        mask_4ch[0] = 1.0 - mask_tensor[0]   # background
        mask_4ch[1] = mask_tensor[0]         # somite mask
        # channels 2 and 3 left as zeros

        return raw_tensor, mask_4ch
    
# --------------- instantiate datasets ---------------

# with augmentation
# train_dataset = SomiteSegmentationDataset(
#     raw_path, raw_key, label_path, split="train",
#     transform_geo=train_geo, transform_intensity=train_intensity
# )
# val_dataset = SomiteSegmentationDataset(
#     raw_path, raw_key, label_path, split="val",
#     transform_geo=val_geo, transform_intensity=val_intensity
# )

# no augmentation version
train_dataset = SomiteSegmentationDataset(
    raw_path, raw_key, label_path, split="train")

val_dataset = SomiteSegmentationDataset(
    raw_path, raw_key, label_path, split="val")

# ----------------- Training start -------------------

# Simple wrapper to give DataLoader a .shuffle attribute for micro-SAM compatibility
class LoaderWrapper:
    def __init__(self, loader, shuffle=False):
        self.loader = loader
        self.shuffle = shuffle  # this is what JointSamTrainer checks for

    def __iter__(self):
        return iter(self.loader)

    def __len__(self):
        return len(self.loader)
    
# Paths
checkpoint_root = "/net/beliveau/vol1/project/EKN_whole_embryo_structural_analysis/Somites/uSAM_finetuning/checkpoints"
run_name = "somite_usam_finetune_old_code_test_5kiter_1e-5lr_vit_b_LM_inf_after_train_2" # Be sure to rename this between runs :) 

run_dir = os.path.join(checkpoint_root, run_name)

# Choose model backbone
model_type = "vit_b_lm"  # valid options: "vit_t", "vit_b", "vit_l", "vit_h", etc.; instance seg. use "vit_b_lm"

# Use LoaderWrapper to add .shuffle attribute
train_loader = LoaderWrapper(
    DataLoader(train_dataset, batch_size=2, shuffle=True, num_workers=0),
    shuffle=True
)

val_loader = LoaderWrapper(
    DataLoader(val_dataset, batch_size=2, shuffle=False, num_workers=0),
    shuffle=False
)

# If your loaders are LoaderWrapper instances:
train_loader = train_loader.loader if hasattr(train_loader, "loader") else train_loader
val_loader = val_loader.loader if hasattr(val_loader, "loader") else val_loader

train_loader.shuffle = True
val_loader.shuffle = False

# Training hyperparameters
n_iterations = 5000 # original 5000
learning_rate = 1e-5 # original 1e-5
device = "cuda" if torch.cuda.is_available() else "cpu"

# # Start training
trainer = train_sam(
    name=run_name, 
    model_type=model_type,
    checkpoint=None, # usually None for FT
    train_loader=train_loader,
    val_loader=val_loader,
    n_iterations=n_iterations,
    lr=learning_rate,
    device=device,
)

# ------------------------------
# Inference on unseen embryo data
# ------------------------------

# ----------------------------------------------
# Paths
# ----------------------------------------------
raw_path = "/net/beliveau/vol2/instrument/E9.5_300/Zoom_300/dataset_fused_rechunked_blocks.zarr"
raw_key = "ch2/s3"

checkpoint_path = ("/net/beliveau/vol1/project/EKN_whole_embryo_structural_analysis/Somites/uSAM_finetuning/checkpoints/somite_usam_finetune_old_code_test_5kiter_1e-5lr_vit_b_LM_inf_after_train_2/best.pt")

fg_out_path = (
    "/net/beliveau/vol2/instrument/E9.5_300/Zoom_300/"
    "somite_prediction/predicted_somite_foreground.tif"
)

device = "cuda" if torch.cuda.is_available() else "cpu"

# ----------------------------------------------
# Load raw data
# ----------------------------------------------
raw = zarr.open(raw_path, mode="r")[raw_key]
print(f"Loaded raw volume: {raw.shape}")

# ----------------------------------------------
# Helper: center pad to 1024 × 1024
# ----------------------------------------------
def center_pad(image, target_shape=(1024, 1024)):
    x, y = image.shape
    tx, ty = target_shape
    pad_x = max((tx - x) // 2, 0)
    pad_y = max((ty - y) // 2, 0)
    return np.pad(
        image,
        ((pad_x, tx - x - pad_x), (pad_y, ty - y - pad_y)),
        mode="constant",
        constant_values=0,
    )

# ------------------------------
# Defining AIS Function
# ------------------------------

def run_automatic_instance_segmentation(
    image: np.ndarray,
    checkpoint_path: Union[os.PathLike ,str],
    model_type: str = model_type,
    device: Optional[Union[str, torch.device]] = None,
    tile_shape: Optional[Tuple[int, int]] = None,
    halo: Optional[Tuple[int, int]] = None,
):
    """Automatic Instance Segmentation (AIS) by training an additional instance decoder in SAM.

    NOTE: AIS is supported only for `µsam` models.

    Args:
        image: The input image.
        checkpoint_path: The path to stored checkpoints.
        model_type: The choice of the `µsam` model.
        device: The device to run the model inference.
        tile_shape: The tile shape for tiling-based segmentation.
        halo: The overlap shape on each side per tile for stitching the segmented tiles.

    Returns:
        The instance segmentation.
    """
    # Step 1: Get the 'predictor' and 'segmenter' to perform automatic instance segmentation.
    predictor, segmenter = get_predictor_and_segmenter(
        model_type=model_type,  # choice of the Segment Anything model
        checkpoint=checkpoint_path,  # overwrite to pass your own finetuned model.
        device=device,  # the device to run the model inference.
        is_tiled=(tile_shape is not None),  # whether to run automatic segmentation.
    )

    # Step 2: Get the instance segmentation for the given image.
    prediction = automatic_instance_segmentation(
        predictor=predictor,  # the predictor for the Segment Anything model.
        segmenter=segmenter,  # the segmenter class responsible for generating predictions.
        input_path=image,  # the filepath to image or the input array for automatic segmentation.
        ndim=2,  # the number of input dimensions.
        tile_shape=tile_shape,  # the tile shape for tiling-based prediction.
        halo=halo,  # the overlap shape for tiling-based prediction.
    )

    return prediction

# ============================================================
# Running prediction and saving results as .TIFF stack
# ============================================================

pred_stack = []

print("Running foreground inference...")

for z in range(raw.shape[0]):
    raw_slice = raw[z].astype(np.float32)
    padded = center_pad(raw_slice)

    if padded.max() > padded.min():
        normed = (padded - padded.min()) / (padded.max() - padded.min())
    else:
        normed = np.zeros_like(padded)

    img_rgb = np.repeat(normed[:, :, None], 3, axis=2)

    # AIS for fg mask
    fg_mask = run_automatic_instance_segmentation(
        image=img_rgb,
        checkpoint_path=checkpoint_path,
        model_type="vit_b_lm",
        device=device,
    )

    pred_stack.append(fg_mask.astype(np.uint8))

pred_stack = np.stack(pred_stack)

tifffile.imwrite(fg_out_path, pred_stack, bigtiff=True)
print(f"Saved foreground prediction to {fg_out_path}")

[constantinpape]

First, running this chunk (in a notebook outside of the training script):

I see! The issue is clear here: the problem is that a custom class (probably the SomiteSegmentationDataset) is pickled.
Then it can't be unpiclked unless you are in the training environment because this class is not found. So you could just import SomiteSegmentationDataset before loading the model and it should work. You can also use export_custom_sam_model in order to export the trained model without the pickled objects, so it can be used in any other environment. IMPORTANT: set with_segmentation_decoder=True to retain the extra decoder for instance segmentation.

My training code is below, from finetuning to inference, all-in-one python script that works as-is. I'm not sure if I am using the params you are thinking of.

You are not using the params I am thinking of -- the issue is due to SomiteSegmentationDataset, see above.

I did not check all of the code in detail, but there is one important point I notice right away: You are training the segmentation decoder very differently from what we usually do: typically, using default_sam_dataset the decoder is trained to predict foreground probabilities and the distances to object center and boundary, which are then used as inputs to a watershed for instance segmentation (within the automatic_instance_segmentation function). Instead, in your dataset you are setting the targets for the segmentation decoder to foreground, background and to zeros. Hence, automatic_instance_segmentation will not really work as expected.

You have the following options to fix this:

1. Use default_sam_dataset instead of your custom datasets. It should be possible to fit the data you have into it -- and we are looking into making this part easier with our new release in case you run into issues here.
2. Use the transformation PerObjectDistanceTransform that creates the outputs for the SegmentationDecoder inside of your SomiteSegmentationDataset.

[evaknichols]

Hi Constantin,

Thank you very much! I see now how I strayed too far out of alignment w/micro-SAM's intended use. In re-tracing my steps, I remembered having trouble with dismantling some of the sampling behavior w/default_sam_dataset params, since I was passing the full image. I didn't want there to be sub-sampling within the image due to the sparse labeling. Most of the image, except very few, will not have a valid image:mask pairs. But now I believe I fixed it (below) and now sadly I'm getting awful performance (IOU less than 0.2, random regions being segmented).

For clarity, all of the labeled training data I provide already are instance masks with non-zero ID's, generated with the micro-SAM GUI annotator. I have not been passing a binary foreground mask for the labels.

Perhaps my task could be too complex for micro-SAM as I am using it below? My task is closer to a medical segmentation task (like from CT imaging--few correct targets in a large 3D dataset), rather than many targets in one image. So maybe the prompting mechanic is failing me here somehow? What if I can focus the attention to the right area first?

I know you said it might not be trivial, but could providing a semantic foreground mask actually work (along with my labeled instance masks and raw iamge data)? I'll try in the meantime to produce a semantice foreground mask, maybe w/a simple U-Net. Then could this be used to focus attention for better prompting w/AIS? The structures are quite distinct so it should work well in principle. I just don't know yet how I could pass that information in using micro-SAM.

Otherwise I welcome any insight or ideas on how I should move forward with my goal....I'm hopeful to can figure it out soon so I can finish analysis for our paper. Perhaps in the worst case if a U-Net works I can just watershed the masks myself.

Thanks in advance for any advice you might have! I look forward to citing micro-SAM in any case, as it was really helpful for producing my labeled training data.

default_sam_dataset construction:

train_dataset = default_sam_dataset(
    raw_paths=raw_paths_train,          
    raw_key=None,
    label_paths=label_paths_train,
    label_key=None,
    patch_shape=(1024, 1024),           # passing full image, previously padded by hand to avoid anisotropic stretching
    with_segmentation_decoder=True,     # required for for AIS
    sampler=None,                       # setting to None to disable internal sampling 
    is_train=True,
    raw_transform=normalize,
    min_size=1,                         # setting to 1 to disable internal sampling 
    max_sampling_attempts=1,            # should be 1; each passed image should have label pair
)

Running:

from micro_sam.util import get_sam_model
checkpoint_path = "..."  # same path as before
model, state = get_sam_model(model_type="vit_b", checkpoint_path=checkpoint_path, return_state=True)

Produces this output (truncated); still works for me though:

UserWarning: SegmentationDataset could not be deserialized because of missing [/tmp/6011301.1.shendure-long.q/tmp_tnd8kou/lbl_0416.tif](https://vscode-remote+tunnel-002beva-005ftunnel.vscode-resource.vscode-cdn.net/tmp/6011301.1.shendure-long.q/tmp_tnd8kou/lbl_0416.tif), None.
The dataset is deserialized in order to allow loading trained models from a checkpoint.
But it cannot be used for further training and will throw an error.
  warnings.warn(msg)

[constantinpape]

Hi Eva,
From your description: I am sure micro_sam can be used for the task. It probably just needs a few minor tweaks to the default parameters. I think the easiest would be if you can share some example data with me -- best would be one volume with annotations (feel free to upload and send it via mail if you don't want to share it publicly).
I am on vacation until the 5th and can't promise that I will have time to check it out beforehand, but I can take a look ASAP once I am back..