tarekziade HF Staff
Publish MIC usage, slim fixture and reviewed Transformers expectations
66f4e01 verified | # SAM2: promptable image segmentation with Transformers | |
| This page is the Transformers contract of the SAM2 MIC and also a usage example. [doc-builder](https://github.com/huggingface/doc-builder/blob/main/docs/runnable-code-blocks.md) renders the `runnable` blocks below as documentation, and its pytest plugin runs them as written against `facebook/sam2.1-hiera-tiny`. The MIC runner runs the same blocks against each slim fixture and each checkpoint in [`integration.json`](../../integration.json). Lines marked `# doc-builder: hide` only matter there: they swap in the pinned checkpoint or variant and print the result record. | |
| ## Segment an object from a point prompt | |
| ```py runnable:sam2_image | |
| import numpy as np | |
| import torch | |
| from PIL import Image | |
| from transformers import Sam2Model, Sam2Processor, set_seed | |
| import json # doc-builder: hide | |
| import os # doc-builder: hide | |
| model_id = "facebook/sam2.1-hiera-tiny" | |
| model_id = os.environ.get("MIC_MODEL", model_id) # doc-builder: hide | |
| variant = json.loads(os.environ.get("MIC_MODEL_KWARGS", "{}")) # doc-builder: hide | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| set_seed(0) # seeds torch, NumPy, and Python | |
| processor = Sam2Processor.from_pretrained(model_id) | |
| model = Sam2Model.from_pretrained(model_id).to(device) | |
| if variant: # doc-builder: hide | |
| del model | |
| model = Sam2Model.from_pretrained(model_id, **variant).to(device) | |
| # Two synthetic 480x640 images: a bright rectangle on a dim noisy background. | |
| # `boxes` holds each rectangle as (x0, y0, x1, y1) in pixels. | |
| rng = np.random.default_rng(42) | |
| boxes = [(160, 120, 400, 300), (240, 180, 560, 420)] # both cover the image centre | |
| images, object_masks = [], [] | |
| for x0, y0, x1, y1 in boxes: | |
| pixels = rng.integers(0, 80, (480, 640, 3), dtype=np.uint8) | |
| pixels[y0:y1, x0:x1] = (230, 180, 40) | |
| images.append(Image.fromarray(pixels)) | |
| mask = np.zeros((480, 640), dtype=np.float32) | |
| mask[y0:y1, x0:x1] = 1.0 | |
| object_masks.append(mask) | |
| # The same positive point on both images, at the image centre: only the pixels | |
| # decide which rectangle is segmented. | |
| # Nesting: [image][object][point][x, y] and [image][object][label]. | |
| input_points = [[[[320, 240]]] for _ in boxes] | |
| input_labels = [[[1]] for _ in boxes] | |
| inputs = processor(images=images, input_points=input_points, input_labels=input_labels, return_tensors="pt") | |
| # Moves every tensor to the device; only floating tensors (pixels, points) take the model's dtype | |
| inputs = inputs.to(device, dtype=model.dtype) | |
| model.eval() | |
| with torch.no_grad(): | |
| outputs = model(**inputs, multimask_output=False) | |
| # pred_masks: low-resolution mask logits, (images, objects, masks, 256, 256) | |
| # iou_scores: the model's estimate of each mask's quality, (images, objects, masks) | |
| masks = processor.post_process_masks(outputs.pred_masks.float().cpu(), inputs["original_sizes"]) | |
| print(f"pred_masks: {list(outputs.pred_masks.shape)}, upscaled to {list(masks[0].shape)}") | |
| print(f"iou_scores: {outputs.iou_scores.flatten().tolist()}") | |
| # Compact probes: a 4x4 grid of mask logits per image, and the IoU estimates | |
| mask_probe = outputs.pred_masks[:, 0, 0, ::64, ::64].float().flatten().tolist() | |
| iou_probe = outputs.iou_scores.float().flatten().tolist() | |
| ``` | |
| ## One fine-tuning step | |
| `Sam2Model` returns predictions but no loss, so a fine-tuning loop computes its own. The loss below is binary cross-entropy between the low-resolution mask logits and each object's ground-truth mask, resized to the logits' 256x256 grid. It is computed in float32 so that a bfloat16 model gets a full-precision loss. | |
| ```py runnable:sam2_image:2 | |
| import torch.nn.functional as F | |
| # Ground truth at the mask logits' resolution, (images, objects, masks, 256, 256) | |
| target = torch.from_numpy(np.stack(object_masks))[:, None] | |
| target = F.interpolate(target, size=outputs.pred_masks.shape[-2:], mode="nearest") | |
| target = target[:, None].to(device) | |
| model.train() | |
| optimizer = torch.optim.SGD(model.parameters(), lr=0.01) | |
| outputs = model(**inputs, multimask_output=False) | |
| loss = F.binary_cross_entropy_with_logits(outputs.pred_masks.float(), target) | |
| loss.backward() | |
| before = {n: p.detach().cpu().clone() for n, p in model.named_parameters() if p.requires_grad} # doc-builder: hide | |
| optimizer.step() | |
| print(f"loss: {loss.item():.4f}") | |
| if os.environ.get("MIC_RUNNER"): # doc-builder: hide | |
| named = dict(model.named_parameters()) | |
| checked = ( | |
| "vision_encoder.backbone.patch_embed.projection.weight", | |
| "vision_encoder.neck.convs.0.weight", | |
| "prompt_encoder.point_embed.weight", | |
| "mask_decoder.upscale_conv1.weight", | |
| ) | |
| result = { | |
| "runtime": { | |
| "device": model.device.type, | |
| "dtype": str(model.dtype).removeprefix("torch."), | |
| "attn_implementation": model.config._attn_implementation, | |
| }, | |
| "pred_masks": mask_probe, | |
| "iou_scores": iou_probe, | |
| "training": { | |
| "loss": loss.item(), | |
| "gradients": { | |
| n: bool( | |
| named[n].grad is not None | |
| and torch.count_nonzero(named[n].grad) > 0 | |
| and torch.isfinite(named[n].grad).all() | |
| ) | |
| for n in checked | |
| }, | |
| "updated": any(not torch.equal(p.detach().cpu(), before[n]) for n, p in model.named_parameters() if n in before), | |
| }, | |
| } | |
| print(json.dumps(result)) | |
| ``` | |