自动标注工具可以集成SAM实现自动标注的功能。附源码和使用说明

主要有三个模型，一般用的B,L比较多一点

L模型是1.2G 兼顾精度和速度 大概75ms以内。

下载地址：https://huggingface.co/datasets/Gourieff/ReActor/blob/main/models/sams/sam_vit_l_0b3195.pth

我是直接暴露端口的，然后提供给标注平台使用。

启动之前记得导入环境变量：

export SAM_CHECKPOINT="/home/sam_vit_l_0b3195.pth"

import base64
import io
import os
from typing import List, Optional, Tuple

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from PIL import Image

def try_import_numpy():
    try:
        import numpy as np  # type: ignore
        return np
    except Exception:
        return None

def try_import_cv2():
    try:
        import cv2  # type: ignore
        return cv2
    except Exception:
        return None

def try_import_torch():
    try:
        import torch  # type: ignore
        return torch
    except Exception:
        return None

def try_import_segment_anything():
    try:
        from segment_anything import SamPredictor, sam_model_registry  # type: ignore
        return SamPredictor, sam_model_registry
    except Exception:
        return None, None

np = try_import_numpy()
cv2 = try_import_cv2()
torch = try_import_torch()
SamPredictor, sam_model_registry = try_import_segment_anything()

class SegmentRequest(BaseModel):
    image: str
    x: int
    y: int
    multi: bool = True

class SamService:
    def __init__(self) -> None:
        self.backend = "unavailable"
        self.device = "cpu"
        self.predictor = None
        self.reason = ""
        self.sam_model_type = os.getenv("SAM_MODEL_TYPE", "vit_l")
        self.sam_checkpoint = os.getenv(
            "SAM_CHECKPOINT",
            "/home/sam_vit_l_0b3195.pth"
        ).strip()
        self.max_results = int(os.getenv("SAM_MAX_RESULTS", "3"))
        self.polygon_epsilon_ratio = float(os.getenv("SAM_POLYGON_EPSILON_RATIO", "0.005"))
        self.min_area = int(os.getenv("SAM_MIN_AREA", "80"))
        self._init_backend()

    def _init_backend(self) -> None:
        if np is None or cv2 is None:
            self.backend = "unavailable"
            self.device = "missing_numpy_or_cv2"
            self.reason = "numpy/cv2 not installed"
            return

        if SamPredictor is None or sam_model_registry is None:
            self.backend = "unavailable"
            self.device = "missing_segment_anything"
            self.reason = "segment_anything is not installed"
            return

        if torch is None:
            self.backend = "unavailable"
            self.device = "missing_torch"
            self.reason = "torch is not installed"
            return

        if not self.sam_checkpoint:
            self.backend = "unavailable"
            self.device = "missing_checkpoint"
            self.reason = "SAM_CHECKPOINT is not set"
            return

        if not os.path.exists(self.sam_checkpoint):
            self.backend = "unavailable"
            self.device = "missing_checkpoint_file"
            self.reason = f"SAM checkpoint not found: {self.sam_checkpoint}"
            return

        try:
            self.device = "cuda" if torch.cuda.is_available() else "cpu"
            sam = sam_model_registry[self.sam_model_type](checkpoint=self.sam_checkpoint)
            sam.to(device=self.device)
            self.predictor = SamPredictor(sam)
            self.backend = "sam"
            self.reason = ""
        except Exception as exc:
            self.backend = "unavailable"
            self.device = "init_failed"
            self.predictor = None
            self.reason = f"failed to initialize SAM: {exc}"

    def health(self) -> dict:
        if self.backend == "unavailable":
            return {
                "ok": False,
                "backend": self.backend,
                "device": self.device,
                "message": self.reason or "SAM is unavailable",
                "model_type": self.sam_model_type,
                "checkpoint": self.sam_checkpoint or None
            }
        return {
            "ok": True,
            "backend": self.backend,
            "device": self.device,
            "model_type": self.sam_model_type if self.backend == "sam" else None,
            "checkpoint": self.sam_checkpoint if self.backend == "sam" else None
        }

    def segment(self, image_data_url: str, x: int, y: int, multi: bool = True) -> List[dict]:
        image_bgr = self._decode_image(image_data_url)
        height, width = image_bgr.shape[:2]
        if not (0 <= x < width and 0 <= y < height):
            raise HTTPException(status_code=400, detail="Click point is outside the image")

        if self.backend == "sam" and self.predictor is not None:
            return self._segment_with_sam(image_bgr, x, y, multi)

        raise HTTPException(status_code=503, detail=self.reason or "SAM service backend is unavailable")

    def _decode_image(self, image_data_url: str):
        if np is None or cv2 is None:
            raise HTTPException(status_code=503, detail="numpy/cv2 not installed")

        if "," not in image_data_url:
            raise HTTPException(status_code=400, detail="Invalid image data URL")

        try:
            _, encoded = image_data_url.split(",", 1)
            image_bytes = base64.b64decode(encoded)
            image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
            image_np = np.array(image)
            return cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
        except Exception as exc:
            raise HTTPException(status_code=400, detail=f"Failed to decode image: {exc}")

    def _segment_with_sam(self, image_bgr, x: int, y: int, multi: bool) -> List[dict]:
        image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
        self.predictor.set_image(image_rgb)

        point_coords = np.array([[x, y]])
        point_labels = np.array([1])
        masks, scores, _ = self.predictor.predict(
            point_coords=point_coords,
            point_labels=point_labels,
            multimask_output=bool(multi),
        )

        results = []
        for mask, score in zip(masks, scores):
            polygon = self._mask_to_polygon(mask.astype("uint8") * 255, x, y)
            if polygon is None:
                continue
            results.append({
                "polygon": polygon,
                "score": float(score),
            })

        results.sort(key=lambda item: item["score"], reverse=True)
        return results[: self.max_results]

    def _mask_to_polygon(self, binary_mask, click_x: int, click_y: int) -> Optional[List[List[int]]]:
        contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        if not contours:
            return None

        selected = None
        selected_area = 0.0
        for contour in contours:
            area = cv2.contourArea(contour)
            if area < self.min_area:
                continue

            inside = cv2.pointPolygonTest(contour, (float(click_x), float(click_y)), False) >= 0
            if inside and area > selected_area:
                selected = contour
                selected_area = area

        if selected is None:
            selected = max(contours, key=cv2.contourArea)
            if cv2.contourArea(selected) < self.min_area:
                return None

        perimeter = cv2.arcLength(selected, True)
        epsilon = max(1.0, perimeter * self.polygon_epsilon_ratio)
        approx = cv2.approxPolyDP(selected, epsilon, True)

        polygon = [[int(point[0][0]), int(point[0][1])] for point in approx]
        if len(polygon) < 3:
            return None
        return polygon

    def _dedupe_results(self, results: List[dict]) -> List[dict]:
        deduped = []
        seen = set()
        for item in results:
            polygon = item["polygon"]
            key = tuple((int(x), int(y)) for x, y in polygon)
            if key in seen:
                continue
            seen.add(key)
            deduped.append(item)
        return deduped

service = SamService()
app = FastAPI(title="SAM Helper Service")

@app.get("/health")
def health():
    return service.health()

@app.post("/segment")
def segment(request: SegmentRequest):
    results = service.segment(request.image, request.x, request.y, request.multi)
    return {
        "success": len(results) > 0,
        "backend": service.backend,
        "device": service.device,
        "results": results,
    }

if __name__ == "__main__":
    import uvicorn

    host = os.getenv("SAM_HOST", "127.0.0.1")
    port = int(os.getenv("SAM_PORT", "8002"))
    uvicorn.run(app, host=host, port=port)