DOCSEARCHbyAlgolia
Submit Docs

MobileSAM

repository·master·Indexed 27 days ago

https://github.com/chaoningzhang/mobilesam

Lightweight versions of the Segment Anything Model (SAM) optimized for mobile and edge applications. It includes MobileSAM and MobileSAMv2, which replace the heavy ViT-H encoder with a smaller Tiny-ViT encoder to achieve faster inference. The library supports prompt-guided mask prediction via SamPredictor, automatic mask generation via SamAutomaticMaskGenerator, and ONNX export. MobileSAMv2 integrates with the ultralytics package for object tracking using BoT-SORT and ByteTracker.

Tokens
3.1K
Snippets
14
Records
17
Agent score
90%

What's inside MobileSAM

  1. Overview of MobileSAM

    master
    MobileSAM is a lightweight version of the Segment Anything Model (SAM). It maintains the same pipeline as the original SAM but replaces the heavyweight ViT-H encoder (632M) with a much smaller Tiny-ViT (5M). This allows it to perform on par with the original SAM visually while being significantly faster, running approximately 12ms per image on a single GPU (8ms for the image encoder and 4ms for the mask decoder).

  3. Export MobileSAM to ONNX

    master

    MobileSAM supports ONNX export. Use the export_onnx_model.py script to convert your checkpoint.

    Recommended versions for testing:

    • onnx==1.12.0
    • onnxruntime==1.13.1
    python scripts/export_onnx_model.py --checkpoint ./weights/mobile_sam.pt --model-type vit_t --output ./mobile_sam.onnx

  4. Track objects using the Python interface

    master

    Use the model.track() method from the ultralytics YOLO model to perform object tracking on video streams or files.

    Supported trackers include:

    • botsort.yaml (BoT-SORT)
    • bytetrack.yaml (ByteTracker)

    When processing a sequence of frames manually (e.g., in a loop or from a folder of images), you must set persist=True. This ensures the model maintains object IDs across frames instead of treating each frame as a new tracking session.

    from ultralytics import YOLO

# Initialize model (detection or segmentation)
model = YOLO("yolov8n.pt")

# Basic tracking usage
model.track(
    source="video/streams",
    stream=True,
    tracker="botsort.yaml",  # or 'bytetrack.yaml'
    show=True,
)

# Accessing tracked object IDs
for result in model.track(source="video.mp4"):
    if result.boxes.id is not None:
        print(result.boxes.id.cpu().numpy().astype(int))

  5. Install MobileSAM

    master

    MobileSAM requires python>=3.8, pytorch>=1.7, and torchvision>=0.8. It is highly recommended to install PyTorch and TorchVision with CUDA support.

    You can install MobileSAM directly via pip or by cloning the repository.

    pip install git+https://github.com/ChaoningZhang/MobileSAM.git

# Or via local clone
git clone git@github.com:ChaoningZhang/MobileSAM.git
cd MobileSAM; pip install -e .

  9. Track objects manually in a video loop

    master

    When iterating through video frames using OpenCV, use model.track(frame, persist=True) to ensure object IDs remain consistent across the sequence.

    import cv2
from ultralytics import YOLO

cap = cv2.VideoCapture("video.mp4")
model = YOLO("yolov8n.pt")
while True:
    ret, frame = cap.read()
    if not ret:
        break
    # persist=True is critical for maintaining IDs across frames
    results = model.track(frame, persist=True)
    
    if results[0].boxes.id is not None:
        boxes = results[0].boxes.xyxy.cpu().numpy().astype(int)
        ids = results[0].boxes.id.cpu().numpy().astype(int)
        for box, id in zip(boxes, ids):
            cv2.rectangle(frame, (box[0], box[1]), (box[2], box[3]), (0, 255, 0), 2)
            cv2.putText(
                frame,
                f"Id {id}",
                (box[0], box[1]),
                cv2.FONT_HERSHEY_SIMPLEX,
                1,
                (0, 0, 255),
                2,
            )
    cv2.imshow("frame", frame)
    if cv2.waitKey(1) & 0xFF == ord("q"):
        break

  10. Use MobileSAM for Prompt-Guided Mask Prediction

    master

    You can use SamPredictor to predict masks based on specific input prompts (like points or boxes).

    1. Load the model using sam_model_registry with the vit_t model type.
    2. Initialize SamPredictor with the loaded model.
    3. Set the image using .set_image().
    4. Predict masks using .predict().
    from mobile_sam import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor
import torch

model_type = "vit_t"
sam_checkpoint = "./weights/mobile_sam.pt"

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

mobile_sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
mobile_sam.to(device=device)
mobile_sam.eval()

predictor = SamPredictor(mobile_sam)
predictor.set_image(<your_image>)
masks, _, _ = predictor.predict(<input_prompts>)

  11. Use MobileSAM for Automatic Mask Generation

    master

    To generate masks for an entire image automatically, use the SamAutomaticMaskGenerator class.

    from mobile_sam import sam_model_registry, SamAutomaticMaskGenerator

# Assuming mobile_sam is already loaded as shown in the predictor example
mask_generator = SamAutomaticMaskGenerator(mobile_sam)
masks = mask_generator.generate(<your_image>)

  12. Use YOLO models via Python API

    master

    You can interact with YOLO models in a Python environment using the ultralytics package. You can build a model from a configuration file (.yaml) to train from scratch, or load a pre-trained checkpoint (.pt) for transfer learning or inference. The YOLO class accepts the same arguments as the CLI.

    from ultralytics import YOLO

model = YOLO("model.yaml")  # build a YOLOv8n model from scratch
# YOLO("model.pt")  use pre-trained model if available
model.info()  # display model information
model.train(data="coco128.yaml", epochs=100)  # train the model