MobileSAM
lastmod 2025-02-11created 2024-06-2615 edits
This is the official code for MobileSAM project that makes SAM lightweight for mobile applications and beyond!
MobileSAM vs FastSAM
- FastSAM
- Trained on only 2% of SA-1B, SAM's training data, and reportedly runs inference about 50x faster at runtime on an NVIDIA GeForce RTX 3090.
- The parameter count did not shrink dramatically, but the actual speed seems much faster. (SAM(Vit-B)-136M, FastSAM-68M)
- MobileSAM
- A model from a paper published by Kyung Hee University that replaces only the image encoder in the SAM architecture with TinyViT.
- TinyViT(5M) is a model that makes ViT lightweight using the distillation proposed in the paper 'Fast Pretraining Distillation for Small Vision Transformers', and it has about 120x fewer parameters than ViT-H(632M), SAM's image encoder. FastSAM has 68M parameters in total while MobileSAM has 9.66M, which shows just how few parameters it has.
Install
Installing on CUDA 12.2: (proceed after setting up the venv environment with the opy script)
git clone --depth 1 https://github.com/ChaoningZhang/MobileSAM
cd MobileSAM
cp {나의 opy python 단독 스크립트} ./python
./python -m pip install torch==2.3.0 torchvision==0.18.0 torchaudio==2.3.0 --index-url https://download.pytorch.org/whl/cu121
./python -m pip install timm
./python opencv-python
Example
# -*- coding: utf-8 -*-
from mobile_sam import sam_model_registry, SamPredictor
import cv2
import os
import numpy as np
from pathlib import Path
from glob import glob
from datetime import datetime
model_type = "vit_t"
sam_checkpoint = "./weights/mobile_sam.pt"
sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
sam.to(device="cuda")
sam.eval()
predictor = SamPredictor(sam)
# files = glob("mpv-shot*.jpg")
# files = glob("*_cam1.bmp")
# files = glob("*.png")
files = glob("./input/*.jpg")
files.sort()
output_dir = Path("output")
if not output_dir.is_dir():
output_dir.mkdir()
for filename in files:
file_prefix, ext = os.path.splitext(filename)
input_point = np.array([[1673, 224]]) # Cam1 (4K)
# input_point = np.array([[836, 112]]) # Cam1 (FHD)
# input_point = np.array([[768, 224]]) # Cam1 (4k-cut)
input_label = np.array([1])
img = cv2.imread(filename)
begin = datetime.now()
predictor.set_image(img)
masks, scores, logits = predictor.predict(
point_coords=input_point,
point_labels=input_label,
multimask_output=True,
)
duration = (datetime.now() - begin).total_seconds()
print(f"Predict duration: {duration:.02f}s")
for i in range(masks.shape[0]):
mask_filename = file_prefix + f"-{i}.png"
mask_filepath = str(output_dir / mask_filename)
cv2.imwrite(mask_filepath, masks[0] * 255)
print(f"Save mask image: {mask_filename}")
ONNX Export
python scripts/export_onnx_model.py --checkpoint ./weights/mobile_sam.pt --model-type vit_t --output ./mobile_sam.onnx
The ONNX model has a different input signature than SamPredictor.predict.
All of the following inputs must be provided. Note the special cases for both point and mask inputs. All inputs are np.float32.
-
image_embeddings- The image embedding fromPredictor.get_image_embedding(). Has a batch index of length 1. -
point_coords- Coordinates of sparse input prompts, covering both point inputs and box inputs. Boxes are encoded with two points: one for the top-left corner and one for the bottom-right corner. Coordinates must already be transformed to long side 1024. Has a batch index of length 1. -
point_labels- Labels for the sparse input prompts. 0 is a negative input point, 1 is a positive input point, 2 is the top-left box corner, 3 is the bottom-right box corner, and -1 is a padding point. If there is no box input, a single padding point with label -1 and coordinates (0.0, 0.0) must be concatenated. -
Mask_input- The mask input to the model, with shape1x1x256x256. It must be provided even if there is no mask input. In that case it can be all zeros. -
has_mask_input- An indicator for the mask input. 1 indicates a mask input, 0 indicates no mask input. -
orig_im_size- The size of the input image in (H,W) format, before any transformation.
Also, the ONNX model does not threshold the output mask logits. To obtain a binary mask, the threshold is sam.mask_threshold (equal to 0.0).
session.get_inputs():
Expression: session.get_inputs()
- 0: <onnxruntime.capi.onnxruntime_pybind11_state.NodeArg object at 0x7f064d2d4170>
+ special variables:
- name: 'image_embeddings'
+ shape: [1, 256, 64, 64]
- type: 'tensor(float)'
- 1: <onnxruntime.capi.onnxruntime_pybind11_state.NodeArg object at 0x7f064d2d41b0>
+ special variables:
- name: 'point_coords'
+ shape: [1, 'num_points', 2]
- type: 'tensor(float)'
- 2: <onnxruntime.capi.onnxruntime_pybind11_state.NodeArg object at 0x7f064d2d41f0>
+ special variables:
- name: 'point_labels'
+ shape: [1, 'num_points']
- type: 'tensor(float)'
- 3: <onnxruntime.capi.onnxruntime_pybind11_state.NodeArg object at 0x7f064d2d4230>
+ special variables:
- name: 'mask_input'
+ shape: [1, 1, 256, 256]
- type: 'tensor(float)'
- 4: <onnxruntime.capi.onnxruntime_pybind11_state.NodeArg object at 0x7f064d2d4270>
+ special variables:
- name: 'has_mask_input'
+ shape: [1]
- type: 'tensor(float)'
- 5: <onnxruntime.capi.onnxruntime_pybind11_state.NodeArg object at 0x7f064d2d42b0>
+ special variables:
- name: 'orig_im_size'
+ shape: [2]
- type: 'tensor(float)'
session.get_outputs():
Expression: session.get_outputs()
- 0: <onnxruntime.capi.onnxruntime_pybind11_state.NodeArg object at 0x7f064d2d42f0>
+ special variables:
- name: 'masks'
+ shape: ['Resizemasks_dim_0', 'Resizemasks_dim_1', 'Resizemasks_dim_2', 'Resizemasks_dim_3']
- type: 'tensor(float)'
- 1: <onnxruntime.capi.onnxruntime_pybind11_state.NodeArg object at 0x7f064d2d4330>
+ special variables:
- name: 'iou_predictions'
+ shape: ['Gemmiou_predictions_dim_0', 4]
- type: 'tensor(float)'
- 2: <onnxruntime.capi.onnxruntime_pybind11_state.NodeArg object at 0x7f064d2d4370>
+ special variables:
- name: 'low_res_masks'
+ shape: ['Reshapelow_res_masks_dim_0', 'Reshapelow_res_masks_dim_1', 'Reshapelow_res_masks_dim_2', 'Reshapelow_res_masks_dim_3']
- type: 'tensor(float)'
ONNX Example
# -*- coding: utf-8 -*-
from mobile_sam import sam_model_registry, SamPredictor
import cv2
import os
import sys
import numpy as np
from pathlib import Path
from glob import glob
from datetime import datetime
import torch
import numpy as np
import cv2
import matplotlib.pyplot as plt
from mobile_sam import sam_model_registry, SamPredictor
from mobile_sam.utils.onnx import SamOnnxModel
import onnxruntime
from onnxruntime.quantization import QuantType
from onnxruntime.quantization.quantize import quantize_dynamic
files = glob("./input/*.jpg")
if not files:
print("Not found any files", file=sys.stderr)
exit(1)
files.sort()
output_dir = Path("output")
if not output_dir.is_dir():
output_dir.mkdir()
onnx_model_path = "mobile_sam.onnx"
session = onnxruntime.InferenceSession(onnx_model_path, providers=['CUDAExecutionProvider'])
iobinding = session.io_binding()
model_type = "vit_t"
sam_checkpoint = "./weights/mobile_sam.pt"
sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
sam.to(device="cuda")
sam.eval()
predictor = SamPredictor(sam)
for filename in files:
file_basename = os.path.basename(filename)
file_prefix, ext = os.path.splitext(file_basename)
input_point = np.array([[1673, 224]]) # Cam1 (4K)
# input_point = np.array([[836, 112]]) # Cam1 (FHD)
# input_point = np.array([[768, 224]]) # Cam1 (4k-cut)
input_label = np.array([1])
image = cv2.imread(filename)
begin = datetime.now()
predictor.set_image(image)
image_embedding = predictor.get_image_embedding().cpu().numpy()
print(f"image_embedding.shape: {image_embedding.shape}")
onnx_coord = np.concatenate([input_point, np.array([[0.0, 0.0]])], axis=0)[None, :, :]
onnx_label = np.concatenate([input_label, np.array([-1])], axis=0)[None, :].astype(np.float32)
onnx_coord = predictor.transform.apply_coords(onnx_coord, image.shape[:2]).astype(np.float32)
onnx_mask_input = np.zeros((1, 1, 256, 256), dtype=np.float32)
onnx_has_mask_input = np.zeros(1, dtype=np.float32)
ort_inputs = {
"image_embeddings": image_embedding,
"point_coords": onnx_coord,
"point_labels": onnx_label,
"mask_input": onnx_mask_input,
"has_mask_input": onnx_has_mask_input,
"orig_im_size": np.array(image.shape[:2], dtype=np.float32)
}
masks, _, low_res_logits = session.run(None, ort_inputs)
masks = masks > predictor.model.mask_threshold
# masks, scores, logits = predictor.predict(
# point_coords=input_point,
# point_labels=input_label,
# multimask_output=True,
# )
duration = (datetime.now() - begin).total_seconds()
print(f"Predict duration: {duration:.02f}s")
for i in range(masks.shape[0]):
mask_filename = file_prefix + f"-{i}.png"
mask_filepath = str(output_dir / mask_filename)
cv2.imwrite(mask_filepath, masks[0] * 255)
print(f"Save mask image: {mask_filepath}")
Internal test results
- NVIDIA GeForce RTX 3070 Ti
- 3840x2160 (4K Original) image - about 0.12s (about 120ms)
- 1600x2160 (ROI Crop) image - about 0.07s (about 70ms)
- NVIDIA GeForce RTX 4090
- 3840x2160 (4K Original) image - 0.0391s ~ 0.0458s (about 45ms)
- 1660x1926 (ROI Crop) image - 0.026s ~ 0.031s (about 31ms)