sam2_api_server.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
鼠标点击分割服务
提供基于点击的交互式图像分割功能，支持前景/背景点击，并可与画笔工具结合使用
"""

import os
import io
import base64
import time
import logging
from pathlib import Path
from typing import Optional, List, Dict, Any

import torch
import numpy as np
import cv2
from PIL import Image
from fastapi import FastAPI, HTTPException
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field

# SAM2导入
import sys
sam2_path = Path(__file__).parent / "sam2"
sys.path.insert(0, str(sam2_path))

from sam2.build_sam import build_sam2
from sam2.sam2_image_predictor import SAM2ImagePredictor

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# 创建FastAPI应用
app = FastAPI(
    title="Interactive Segmentation API",
    description="交互式图像分割服务，支持鼠标点击和画笔修复",
    version="1.0.0"
)

# CORS配置
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 全局模型管理
class SAM2ModelManager:
    """SAM2模型管理器"""
    _instance = None
    _initialized = False

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
        return cls._instance

    def __init__(self):
        if not self._initialized:
            self.predictor = None
            self.device = None
            self.current_image = None
            self.current_image_embedding = None
            self.current_image_hash = None  # 新增：图像哈希缓存
            self.service_id = "sam2_interactive_segmentation"
            self._initialized = True

    def load_model(self, checkpoint_path: str, model_cfg: str):
        """加载SAM2模型"""
        if self.predictor is not None:
            logger.info("模型已加载，跳过重复加载")
            return True

        logger.info("=" * 60)
        logger.info("开始加载SAM2模型...")
        logger.info("=" * 60)

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

        try:
            sam2_model = build_sam2(model_cfg, checkpoint_path)
            self.predictor = SAM2ImagePredictor(sam2_model)

            logger.info("=" * 60)
            logger.info("✓ 模型加载成功！")
            logger.info(f"  - Device: {self.device}")
            logger.info(f"  - Model: {model_cfg}")
            logger.info(f"  - Checkpoint: {checkpoint_path}")
            logger.info("=" * 60)
            return True

        except Exception as e:
            logger.error(f"模型加载失败: {e}")
            import traceback
            traceback.print_exc()
            return False

    def is_ready(self):
        """检查模型是否就绪"""
        return self.predictor is not None

    def set_image(self, image: np.ndarray, image_hash: str = None, force: bool = False):
        """设置当前图像并计算embedding（带缓存）

        Args:
            image: 图像数组（用于推理，可能是缩放后的）
            image_hash: 原始图像的哈希（用于缓存判断）
            force: 是否强制重新计算
        """
        # 如果没有提供哈希，自己计算
        if image_hash is None:
            import hashlib
            image_hash = hashlib.md5(image.tobytes()).hexdigest()

        # 如果是同一张图，跳过重新计算
        if not force and self.current_image_hash == image_hash:
            logger.info(f"图像未变化，使用缓存的embedding（耗时<1ms）")
            return

        with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
            self.predictor.set_image(image)
            self.current_image = image
            self.current_image_hash = image_hash
            logger.info(f"图像已设置，尺寸: {image.shape}，哈希: {image_hash[:8]}...")

    def predict(self, point_coords: np.ndarray = None, point_labels: np.ndarray = None,
                box: np.ndarray = None, multimask_output: bool = True):
        """执行预测

        Args:
            point_coords: 点击坐标 [[x, y], ...]
            point_labels: 点击标签 [1, 0, ...]
            box: 边界框 [x1, y1, x2, y2]
            multimask_output: 是否输出多个mask
        """
        if self.predictor is None:
            raise RuntimeError("模型未加载")

        with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
            masks, scores, logits = self.predictor.predict(
                point_coords=point_coords,
                point_labels=point_labels,
                box=box,
                multimask_output=multimask_output,
            )
            return masks, scores, logits

# 全局模型管理器
model_manager = SAM2ModelManager()

# API请求模型
class Point(BaseModel):
    """点击点"""
    x: float = Field(..., description="X坐标")
    y: float = Field(..., description="Y坐标")
    label: int = Field(..., description="标签：1=前景，0=背景")

class SetImageRequest(BaseModel):
    """设置图像请求"""
    image: str = Field(..., description="图像（base64编码）")

class PredictRequest(BaseModel):
    """预测请求"""
    image: str = Field(..., description="图像（base64编码）")
    points: List[Point] = Field(default=[], description="点击点列表（可选）")
    box: List[float] = Field(default=None, description="边界框 [x1, y1, x2, y2]（可选）")
    multimask_output: bool = Field(True, description="是否输出多个mask")
    return_logits: bool = Field(False, description="是否返回logits（用于后续细化）")
    inference_size: int = Field(1024, description="推理尺寸（默认1024，加速推理）", ge=256, le=1024)
    filter_small_regions: bool = Field(False, description="是否过滤小碎片")
    dilate_kernel_size: int = Field(15, description="膨胀核大小（默认15，用于扩展mask边缘）", ge=0, le=50)

class RefineRequest(BaseModel):
    """细化mask请求"""
    points: List[Point] = Field(..., description="新增的点击点列表")
    previous_logits: str = Field(None, description="上一次的logits（base64编码的numpy数组）")
    multimask_output: bool = Field(False, description="是否输出多个mask")

class MergeMasksRequest(BaseModel):
    """合并mask请求"""
    mask1: str = Field(..., description="第一个mask（base64编码）")
    mask2: str = Field(..., description="第二个mask（base64编码）")
    operation: str = Field("union", description="操作类型：union(并集), intersect(交集), diff(差集)")

class ServiceInfo(BaseModel):
    """服务信息"""
    service_id: str
    name: str
    description: str
    category: str
    version: str
    model_type: str

# 工具函数
def base64_to_image(base64_str: str) -> Optional[np.ndarray]:
    """Base64字符串转numpy数组（RGB格式）"""
    try:
        if base64_str.startswith('data:image'):
            base64_str = base64_str.split(',', 1)[1]

        image_data = base64.b64decode(base64_str)
        image = Image.open(io.BytesIO(image_data)).convert('RGB')
        return np.array(image)
    except Exception as e:
        logger.error(f"Base64转图像失败: {e}")
        return None

def image_to_base64(image: np.ndarray, format: str = 'PNG') -> str:
    """numpy数组转Base64字符串"""
    try:
        pil_image = Image.fromarray(image.astype(np.uint8))
        buffer = io.BytesIO()
        pil_image.save(buffer, format=format)
        img_str = base64.b64encode(buffer.getvalue()).decode()
        return f"data:image/{format.lower()};base64,{img_str}"
    except Exception as e:
        logger.error(f"图像转Base64失败: {e}")
        return None

def mask_to_base64(mask: np.ndarray) -> str:
    """mask转Base64字符串（单通道）"""
    try:
        mask_uint8 = (mask * 255).astype(np.uint8)
        pil_mask = Image.fromarray(mask_uint8, mode='L')
        buffer = io.BytesIO()
        pil_mask.save(buffer, format='PNG')
        mask_str = base64.b64encode(buffer.getvalue()).decode()
        return f"data:image/png;base64,{mask_str}"
    except Exception as e:
        logger.error(f"Mask转Base64失败: {e}")
        return None

def base64_to_mask(base64_str: str) -> Optional[np.ndarray]:
    """Base64字符串转mask数组"""
    try:
        if base64_str.startswith('data:image'):
            base64_str = base64_str.split(',', 1)[1]

        mask_data = base64.b64decode(base64_str)
        mask_image = Image.open(io.BytesIO(mask_data)).convert('L')
        mask = np.array(mask_image) > 128
        return mask
    except Exception as e:
        logger.error(f"Base64转Mask失败: {e}")
        return None

def numpy_to_base64(arr: np.ndarray) -> str:
    """numpy数组转base64字符串（用于传输logits）"""
    buffer = io.BytesIO()
    np.save(buffer, arr)
    buffer.seek(0)
    return base64.b64encode(buffer.read()).decode()

def base64_to_numpy(base64_str: str) -> np.ndarray:
    """base64字符串转numpy数组"""
    buffer = io.BytesIO(base64.b64decode(base64_str))
    return np.load(buffer)

def apply_mask_to_image(image: np.ndarray, mask: np.ndarray, color=(0, 255, 0), alpha=0.3):
    """将mask叠加到图像上"""
    overlay = image.copy()
    overlay[mask > 0] = overlay[mask > 0] * (1 - alpha) + np.array(color) * alpha
    return overlay.astype(np.uint8)

def resize_image_keep_aspect(image: np.ndarray, target_size: int) -> tuple:
    """等比例缩放图像到目标尺寸，返回(缩放后图像, 原始尺寸)"""
    h, w = image.shape[:2]
    scale = target_size / max(h, w)

    if scale >= 1.0:
        # 图像已经够小，不需要缩放
        return image, (h, w)

    new_h, new_w = int(h * scale), int(w * scale)
    resized = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
    return resized, (h, w)

def resize_mask_to_original(mask: np.ndarray, original_size: tuple) -> np.ndarray:
    """将mask缩放回原始尺寸"""
    h, w = original_size
    if mask.shape[:2] == (h, w):
        return mask

    # 使用最近邻插值保持mask的二值特性
    resized = cv2.resize(mask.astype(np.uint8), (w, h), interpolation=cv2.INTER_NEAREST)
    return resized.astype(bool)

def filter_small_regions(mask: np.ndarray, min_area_ratio: float = 0.01) -> np.ndarray:
    """过滤掉小碎片，只保留主体mask

    Args:
        mask: 输入mask
        min_area_ratio: 最小区域面积占比（相对于最大连通域）

    Returns:
        过滤后的mask
    """
    # 查找所有连通域
    num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(
        mask.astype(np.uint8), connectivity=8
    )

    if num_labels <= 1:
        return mask

    # 获取每个连通域的面积（跳过背景label=0）
    areas = stats[1:, cv2.CC_STAT_AREA]

    if len(areas) == 0:
        return mask

    # 找到最大连通域
    max_area = areas.max()
    min_area = max_area * min_area_ratio

    # 创建新mask，只保留足够大的连通域
    filtered_mask = np.zeros_like(mask, dtype=bool)

    for label_idx in range(1, num_labels):
        area = stats[label_idx, cv2.CC_STAT_AREA]
        if area >= min_area:
            filtered_mask[labels == label_idx] = True

    return filtered_mask

def dilate_mask(mask: np.ndarray, kernel_size: int = 15) -> np.ndarray:
    """膨胀mask，扩展边缘区域

    Args:
        mask: 输入mask（bool或uint8）
        kernel_size: 膨胀核大小（默认15）

    Returns:
        膨胀后的mask
    """
    if kernel_size <= 0:
        return mask

    # 转换为uint8
    mask_uint8 = mask.astype(np.uint8) * 255

    # 创建膨胀核
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size))

    # 执行膨胀操作
    dilated = cv2.dilate(mask_uint8, kernel, iterations=1)

    # 转换回bool
    return dilated > 128

# API端点
@app.get("/")
async def root():
    """根路径"""
    return {
        "service": "SAM2 Interactive Segmentation API",
        "service_id": model_manager.service_id,
        "version": "1.0.0",
        "status": "running",
        "model_loaded": model_manager.is_ready()
    }

@app.get("/api/health")
async def health_check():
    """健康检查"""
    return {
        "status": "healthy" if model_manager.is_ready() else "initializing",
        "model_loaded": model_manager.is_ready(),
        "device": str(model_manager.device) if model_manager.device else "unknown",
        "service_id": model_manager.service_id
    }

@app.get("/api/service/info", response_model=ServiceInfo)
async def get_service_info():
    """获取服务信息"""
    return ServiceInfo(
        service_id="sam2_interactive_segmentation",
        name="SAM2智能分割",
        description="基于SAM2的交互式图像分割，通过点击生成精确mask",
        category="image_segmentation",
        version="1.0.0",
        model_type="SAM2.1 Hiera Tiny"
    )

@app.post("/api/predict")
async def predict(request: PredictRequest):
    """
    点击/框选预测接口 - 返回最佳mask

    输入：
    - image: 原始图像（base64）
    - points: 点击点列表（可选），每个点包含 x, y, label（1=前景，0=背景）
    - box: 边界框（可选），格式 [x1, y1, x2, y2]
    - multimask_output: 是否输出多个mask候选（默认True）
    - return_logits: 是否返回logits用于后续细化（默认False）
    - dilate_kernel_size: 膨胀核大小（默认15，用于扩展mask边缘，便于消除）

    输出：
    - mask: 最佳mask（base64编码，仅data部分无前缀，已膨胀）
    - score: 置信度分数
    - preview: 预览图（mask叠加在原图上）

    使用示例：
    1. 点击模式：提供points
    2. 框选模式：提供box [x1, y1, x2, y2]
    3. 组合模式：同时提供points和box（更精确）

    注意：返回的mask默认已膨胀15像素，可通过dilate_kernel_size调整或设为0关闭
    """
    if not model_manager.is_ready():
        raise HTTPException(status_code=503, detail="模型未加载，请稍后重试")

    start_time = time.time()

    try:
        # 1. 解码图像
        logger.info("解码输入图像...")
        original_image = base64_to_image(request.image)
        if original_image is None:
            raise HTTPException(status_code=400, detail="无效的图像数据")

        original_h, original_w = original_image.shape[:2]
        logger.info(f"原始图像尺寸: {original_w}x{original_h}")

        # 2. 缩放图像到推理尺寸（先缩放！）
        resized_image, original_size = resize_image_keep_aspect(
            original_image,
            request.inference_size
        )
        resized_h, resized_w = resized_image.shape[:2]

        if (resized_h, resized_w) != (original_h, original_w):
            logger.info(f"图像已缩放: {original_w}x{original_h} → {resized_w}x{resized_h}")
            scale_x = resized_w / original_w
            scale_y = resized_h / original_h
        else:
            logger.info("图像无需缩放")
            scale_x = scale_y = 1.0

        # 3. 计算哈希（用缩放后的图！）
        import hashlib
        resized_image_hash = hashlib.md5(resized_image.tobytes()).hexdigest()

        # 4. 设置图像（传入缩放后的图像 + 缩放后图像的哈希）
        model_manager.set_image(resized_image, image_hash=resized_image_hash)

        # 5. 准备prompt数据（点击或框）
        if not request.points and not request.box:
            raise HTTPException(status_code=400, detail="需要提供points或box")

        # 准备点击点
        point_coords = None
        point_labels = None
        if request.points and len(request.points) > 0:
            point_coords = np.array([
                [p.x * scale_x, p.y * scale_y] for p in request.points
            ], dtype=np.float32)
            point_labels = np.array([p.label for p in request.points], dtype=np.int32)
            logger.info(f"点击点: {len(request.points)}个, 坐标: {point_coords.tolist()}")

        # 准备边界框
        box = None
        if request.box and len(request.box) == 4:
            x1, y1, x2, y2 = request.box
            box = np.array([
                x1 * scale_x, y1 * scale_y,
                x2 * scale_x, y2 * scale_y
            ], dtype=np.float32)
            logger.info(f"边界框: {box.tolist()}")

        # 6. 执行预测
        masks, scores, logits = model_manager.predict(
            point_coords=point_coords,
            point_labels=point_labels,
            box=box,
            multimask_output=request.multimask_output
        )

        logger.info(f"生成了 {len(masks)} 个mask，scores: {scores}")

        # 6. 选择最佳mask
        best_mask = masks[0]
        best_score = float(scores[0])

        # 7. 过滤小碎片（可选）
        if request.filter_small_regions:
            original_mask_area = best_mask.sum()
            best_mask = filter_small_regions(best_mask, min_area_ratio=0.05)
            filtered_area = best_mask.sum()

            if filtered_area < original_mask_area:
                logger.info(f"已过滤小碎片: {original_mask_area} → {filtered_area} pixels")

        # 8. 膨胀mask（在缩放尺寸上操作，默认kernel_size=15）
        if request.dilate_kernel_size > 0:
            original_mask_area = best_mask.sum()
            best_mask = dilate_mask(best_mask, kernel_size=request.dilate_kernel_size)
            dilated_area = best_mask.sum()
            logger.info(f"Mask已膨胀: kernel={request.dilate_kernel_size}, {original_mask_area} → {dilated_area} pixels (在{resized_w}x{resized_h}尺寸上)")

        # 9. 生成预览图（用缩放后的小图！快！）
        preview_resized = apply_mask_to_image(resized_image, best_mask, color=(0, 255, 0), alpha=0.4)
        preview_b64 = image_to_base64(preview_resized)

        # 10. 缩放mask回原始尺寸
        if (resized_h, resized_w) != (original_h, original_w):
            best_mask = resize_mask_to_original(best_mask, original_size)
            logger.info(f"Mask已缩放回原始尺寸: {resized_w}x{resized_h} → {original_w}x{original_h}")

        # 11. 转换mask为base64
        mask_uint8 = (best_mask * 255).astype(np.uint8)
        pil_mask = Image.fromarray(mask_uint8, mode='L')
        buffer = io.BytesIO()
        pil_mask.save(buffer, format='PNG')
        mask_base64 = base64.b64encode(buffer.getvalue()).decode()

        processing_time = time.time() - start_time
        logger.info(f"✓ 预测完成，最佳score: {best_score:.4f}，耗时: {processing_time:.2f}秒")

        response = {
            "success": True,
            "mask": mask_base64,  # 仅base64字符串，无前缀
            "score": best_score,
            "preview": preview_b64,
            "processing_time": round(processing_time, 3),
            "image_shape": [original_h, original_w],  # 原始尺寸
            "num_points": len(request.points)
        }

        # 如果需要返回logits用于细化
        if request.return_logits:
            response["logits"] = numpy_to_base64(logits[0])

        return response

    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"预测失败: {e}")
        import traceback
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=f"预测失败: {str(e)}")

@app.post("/api/set_image")
async def set_image(request: SetImageRequest):
    """
    设置图像（预先计算embedding，加速后续预测）

    输入：
    - image: 图像（base64）

    输出：
    - success: 是否成功
    - image_shape: 图像尺寸
    """
    if not model_manager.is_ready():
        raise HTTPException(status_code=503, detail="模型未加载，请稍后重试")

    try:
        image = base64_to_image(request.image)
        if image is None:
            raise HTTPException(status_code=400, detail="无效的图像数据")

        model_manager.set_image(image)

        return {
            "success": True,
            "image_shape": image.shape[:2],
            "message": "图像已设置，embedding已计算"
        }

    except Exception as e:
        logger.error(f"设置图像失败: {e}")
        raise HTTPException(status_code=500, detail=f"设置图像失败: {str(e)}")

@app.post("/api/merge_masks")
async def merge_masks(request: MergeMasksRequest):
    """
    合并多个mask（用于画笔修复等场景）

    输入：
    - mask1: 第一个mask（base64）
    - mask2: 第二个mask（base64）
    - operation: 操作类型（union/intersect/diff）

    输出：
    - merged_mask: 合并后的mask（base64）
    """
    try:
        mask1 = base64_to_mask(request.mask1)
        mask2 = base64_to_mask(request.mask2)

        if mask1 is None or mask2 is None:
            raise HTTPException(status_code=400, detail="无效的mask数据")

        # 执行操作
        if request.operation == "union":
            merged = np.logical_or(mask1, mask2)
        elif request.operation == "intersect":
            merged = np.logical_and(mask1, mask2)
        elif request.operation == "diff":
            merged = np.logical_and(mask1, np.logical_not(mask2))
        else:
            raise HTTPException(status_code=400, detail="不支持的操作类型")

        merged_b64 = mask_to_base64(merged)

        return {
            "success": True,
            "merged_mask": merged_b64,
            "operation": request.operation
        }

    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"合并mask失败: {e}")
        raise HTTPException(status_code=500, detail=f"合并mask失败: {str(e)}")

@app.on_event("startup")
async def startup_event():
    """启动时加载模型"""
    logger.info("=" * 60)
    logger.info("API Server 启动中...")
    logger.info("=" * 60)

    # 模型路径配置
    checkpoint_path = "./sam2/checkpoints/sam2.1_hiera_tiny.pt"
    model_cfg = "configs/sam2.1/sam2.1_hiera_t.yaml"

    # 检查文件是否存在
    if not os.path.exists(checkpoint_path):
        logger.warning(f"checkpoint文件不存在: {checkpoint_path}")
        logger.warning("服务将以降级模式运行，请先下载模型文件")
        return

    # 加载模型
    success = model_manager.load_model(checkpoint_path, model_cfg)

    if not success:
        logger.error("=" * 60)
        logger.error("⚠️  模型加载失败，服务将以降级模式运行")
        logger.error("=" * 60)
    else:
        logger.info("=" * 60)
        logger.info("✓  主体分割 API Server 启动成功！")
        logger.info("=" * 60)

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8002, log_level="info")