"Initial commit"

docker-compose.yaml

+services:
+  frontend:
+    image: sam2/frontend
+    build:
+      context: ./demo/frontend
+      dockerfile: frontend.Dockerfile
+    ports:
+      - 7262:80
+
+  backend:
+    image: sam2/backend
+    build:
+      context: .
+      dockerfile: backend.Dockerfile
+    ports:
+      - 7263:5000
+    volumes:
+      - ./demo/data/:/data/:rw
+    environment:
+      - SERVER_ENVIRONMENT=DEV
+      - GUNICORN_WORKERS=1
+      # Inference API needs to have at least 2 threads to handle an incoming
+      # parallel cancel propagation request
+      - GUNICORN_THREADS=2
+      - GUNICORN_PORT=5000
+      - API_URL=http://localhost:7263
+      - DEFAULT_VIDEO_PATH=gallery/05_default_juggle.mp4
+      # # ffmpeg/video encode settings
+      - FFMPEG_NUM_THREADS=1
+      - VIDEO_ENCODE_CODEC=libx264
+      - VIDEO_ENCODE_CRF=23
+      - VIDEO_ENCODE_FPS=24
+      - VIDEO_ENCODE_MAX_WIDTH=1280
+      - VIDEO_ENCODE_MAX_HEIGHT=720
+      - VIDEO_ENCODE_VERBOSE=False
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: 1
+              capabilities: [gpu]

image_predictor_example.py

+import os
+import numpy as np
+import torch
+import matplotlib.pyplot as plt
+from PIL import Image
+from sam2.build_sam import build_sam2
+from sam2.sam2_image_predictor import SAM2ImagePredictor
+
+# Environment setup (assuming dependencies are installed)
+os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
+
+# Select device
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+    torch.autocast("cuda", dtype=torch.bfloat16).__enter__()
+    if torch.cuda.get_device_properties(0).major >= 8:
+        torch.backends.cuda.matmul.allow_tf32 = True
+        torch.backends.cudnn.allow_tf32 = True
+elif torch.backends.mps.is_available():
+    device = torch.device("mps")
+else:
+    device = torch.device("cpu")
+print(f"using device: {device}")
+
+# Helper functions (modified to save instead of show)
+def save_mask_on_image(image, mask, save_path, alpha=0.5, borders=True):
+    fig, ax = plt.subplots(figsize=(10, 10))
+    ax.imshow(image)
+
+    color = np.array([30/255, 144/255, 255/255, alpha])
+    h, w = mask.shape[-2:]
+    mask = mask.astype(np.uint8)
+    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
+
+    if borders:
+        import cv2
+        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
+        contours = [cv2.approxPolyDP(contour, epsilon=0.01, closed=True) for contour in contours]
+        mask_image = cv2.drawContours(mask_image, contours, -1, (1, 1, 1, 0.5), thickness=2)
+
+    ax.imshow(mask_image)
+    ax.axis('off')
+    plt.savefig(save_path, bbox_inches='tight', pad_inches=0)
+    plt.close(fig)
+
+def save_points_on_image(image, coords, labels, save_path):
+    fig, ax = plt.subplots(figsize=(10, 10))
+    ax.imshow(image)
+    pos_points = coords[labels == 1]
+    neg_points = coords[labels == 0]
+    marker_size = 375
+    ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)
+    ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)
+    ax.axis('off')
+    plt.savefig(save_path, bbox_inches='tight', pad_inches=0)
+    plt.close(fig)
+
+def save_box_on_image(image, box, save_path):
+    fig, ax = plt.subplots(figsize=(10, 10))
+    ax.imshow(image)
+    x0, y0 = box[0], box[1]
+    w, h = box[2] - box[0], box[3] - box[1]
+    ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0, 0, 0, 0), lw=2))
+    ax.axis('off')
+    plt.savefig(save_path, bbox_inches='tight', pad_inches=0)
+    plt.close(fig)
+
+
+# Load model
+sam2_checkpoint = "./checkpoints/sam2.1_hiera_large.pt"
+model_cfg = "configs/sam2.1/sam2.1_hiera_l.yaml"
+sam2 = build_sam2(model_cfg, sam2_checkpoint, device=device)
+predictor = SAM2ImagePredictor(sam2)
+
+# Example 1: Truck image with point prompt
+image = Image.open('./notebooks/images/truck.jpg').convert("RGB")
+image_np = np.array(image)
+
+with torch.inference_mode(), torch.autocast(device.type, dtype=torch.bfloat16):
+    predictor.set_image(image_np)
+    input_point = np.array([[825, 619]])
+    input_label = np.array([1])
+    masks, scores, logits = predictor.predict(
+        point_coords=input_point,
+        point_labels=input_label,
+        multimask_output=True,
+    )
+
+# Save results
+os.makedirs("results", exist_ok=True)
+for i, (mask, score) in enumerate(zip(masks, scores)):
+    save_path = f"results/truck_mask_{i}_score_{score:.3f}.png"
+    save_mask_on_image(image_np, mask, save_path)
+    np.save(f"results/truck_mask_{i}.npy", mask)
+
+save_points_on_image(image_np, input_point, input_label, "results/truck_points.png")
+
+# Example 2: Truck with box prompt
+with torch.inference_mode(), torch.autocast(device.type, dtype=torch.bfloat16):
+    input_box = np.array([425, 600, 700, 875])
+    masks, scores, _ = predictor.predict(
+        point_coords=None,
+        point_labels=None,
+        box=input_box,
+        multimask_output=False,
+    )
+
+for i, (mask, score) in enumerate(zip(masks, scores)):
+    save_path = f"results/truck_box_mask_{i}_score_{score:.3f}.png"
+    save_mask_on_image(image_np, mask, save_path)
+    np.save(f"results/truck_box_mask_{i}.npy", mask)
+
+save_box_on_image(image_np, input_box, "results/truck_box.png")
+
+# Example 3: Truck with multiple points
+with torch.inference_mode(), torch.autocast(device.type, dtype=torch.bfloat16):
+    input_points = np.array([[450, 600], [775, 800]])
+    input_labels = np.array([1, 1])
+    masks, scores, _ = predictor.predict(
+        point_coords=input_points,
+        point_labels=input_labels,
+        multimask_output=False,
+    )
+
+for i, (mask, score) in enumerate(zip(masks, scores)):
+    save_path = f"results/truck_multi_points_mask_{i}_score_{score:.3f}.png"
+    save_mask_on_image(image_np, mask, save_path)
+    np.save(f"results/truck_multi_points_mask_{i}.npy", mask)
+
+save_points_on_image(image_np, input_points, input_labels, "results/truck_multi_points.png")
+
+# Example 4: Groceries image with batch prompts
+image2 = Image.open('./notebooks/images/groceries.jpg').convert("RGB")
+image2_np = np.array(image2)
+
+with torch.inference_mode(), torch.autocast(device.type, dtype=torch.bfloat16):
+    predictor.set_image(image2_np)
+    input_points2 = np.array([[350, 450], [400, 450]])
+    input_labels2 = np.array([1, 1])
+    masks2, scores2, _ = predictor.predict(
+        point_coords=input_points2,
+        point_labels=input_labels2,
+        multimask_output=False,
+    )
+
+for i, (mask, score) in enumerate(zip(masks2, scores2)):
+    save_path = f"results/groceries_mask_{i}_score_{score:.3f}.png"
+    save_mask_on_image(image2_np, mask, save_path)
+    np.save(f"results/groceries_mask_{i}.npy", mask)
+
+save_points_on_image(image2_np, input_points2, input_labels2, "results/groceries_points.png")
+
+print("All results saved in 'results' folder.")
\ No newline at end of file

model.properties

+# 模型唯一标识
+modelCode=1879
+# 模型名称
+modelName=sam2_pytorch
+# 模型描述
+modelDescription=Segment Anything Model 2 (SAM 2)是一个用于解决图像和视频中可提示视觉分割问题的基础模型。
+# 运行过程
+processType=推理
+# 算法类别
+appCategory=图像分割
+# 框架类型
+frameType=pytorch
+# 加速卡类型
+accelerateType=BW1000