get_mask.py

import argparse
import os

import numpy as np
import torch
from PIL import Image


def parse_points(points_text):
    points = []
    for item in points_text.split(";"):
        item = item.strip()
        if not item:
            continue
        x, y = item.split(",")
        points.append([float(x), float(y)])
    if not points:
        raise ValueError("At least one positive point is required, for example: --points 480,1470")
    return points


def main():
    parser = argparse.ArgumentParser(description="Generate a reference subject mask with SAM-2.")
    parser.add_argument("--image", required=True, help="Reference image path.")
    parser.add_argument("--output", required=True, help="Output mask path.")
    parser.add_argument(
        "--points",
        required=True,
        help="Positive SAM points formatted as 'x,y;x,y'. Coordinates are in image pixels.",
    )
    parser.add_argument(
        "--sam2_model",
        default="facebook/sam2.1-hiera-large",
        help="SAM-2 model id or local directory.",
    )
    parser.add_argument("--mask_index", type=int, default=0, help="SAM-2 returns multiple masks. Select one index.")
    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    args = parser.parse_args()
    try:
        from transformers import Sam2Model, Sam2Processor
    except ImportError as exc:
        raise ImportError(
            "SAM-2 mask generation requires a transformers version that provides "
            "Sam2Model and Sam2Processor. Please upgrade transformers or pass a "
            "precomputed mask to the inference script."
        ) from exc

    image = Image.open(args.image).convert("RGB")
    points = parse_points(args.points)
    input_points = [[points]]
    input_labels = [[[1 for _ in points]]]

    processor = Sam2Processor.from_pretrained(args.sam2_model)
    model = Sam2Model.from_pretrained(args.sam2_model).to(args.device)
    inputs = processor(
        images=image,
        input_points=input_points,
        input_labels=input_labels,
        return_tensors="pt",
    ).to(args.device)

    with torch.no_grad():
        outputs = model(**inputs)

    masks = processor.post_process_masks(outputs.pred_masks.cpu(), inputs["original_sizes"])[0]
    if args.mask_index >= masks.shape[1]:
        raise ValueError(f"mask_index={args.mask_index} is out of range. SAM-2 returned {masks.shape[1]} masks.")

    mask = masks[0, args.mask_index].numpy()
    mask = (mask > 0).astype(np.uint8) * 255
    mask = Image.fromarray(mask).convert("RGB")

    os.makedirs(os.path.dirname(os.path.abspath(args.output)), exist_ok=True)
    mask.save(args.output)
    print(f"Saved mask to {args.output}")


if __name__ == "__main__":
    main()