File size: 3,994 Bytes
d82bbe4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 | from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import Optional, List, Any, Dict
import numpy as np
import base64
import zlib
import os
import sys
# Add project root to path to ensure imports work
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.abspath(os.path.join(current_dir, "../../../"))
if project_root not in sys.path:
sys.path.insert(0, project_root)
from dataflow_agent.toolkits.imtool.sam_tool import run_sam_auto, free_sam_model
try:
import torch
except ImportError:
torch = None
app = FastAPI(title="SAM Model Server")
# Check CUDA device on startup
@app.on_event("startup")
async def startup_event():
print("SAM Server Startup Check:")
print(f"CUDA_VISIBLE_DEVICES: {os.environ.get('CUDA_VISIBLE_DEVICES', 'Not Set')}")
if torch and torch.cuda.is_available():
print(f"Torch CUDA available: {torch.cuda.is_available()}")
print(f"Current Device Count: {torch.cuda.device_count()}")
print(f"Current Device Name: {torch.cuda.get_device_name(0)}")
else:
print("CUDA NOT AVAILABLE")
class SAMRequest(BaseModel):
image_path: str
checkpoint: str = "sam_b.pt"
device: str = "cuda"
class SAMItemResponse(BaseModel):
mask_b64: str
mask_shape: List[int]
bbox: List[float]
score: Optional[float] = None
area: int
class SAMResponse(BaseModel):
items: List[SAMItemResponse]
@app.post("/predict", response_model=SAMResponse)
async def predict(req: SAMRequest):
"""
Run SAM auto segmentation on the given image path.
"""
if not os.path.exists(req.image_path):
raise HTTPException(status_code=404, detail=f"Image path not found: {req.image_path}")
try:
# Force device to cuda:0 if it's generic cuda, to align with CUDA_VISIBLE_DEVICES env
target_device = req.device
if target_device == "cuda":
target_device = "cuda:0"
# 调用本地的 sam_tool 函数
# 注意：这里会利用 sam_tool 内部的 caching 机制
# 如果启动了多个 sam_server 进程，每个进程会维护自己的 cache
items = run_sam_auto(
image_path=req.image_path,
checkpoint=req.checkpoint,
device=target_device
)
# 序列化结果
serialized_items = []
for it in items:
mask = it.get("mask")
if mask is None:
continue
# Convert mask to base64
# mask is numpy array (bool or uint8)
if not isinstance(mask, np.ndarray):
mask = np.array(mask)
# Use bool type for serialization consistency
mask_bool = mask.astype(bool)
mask_bytes = mask_bool.tobytes()
# Compress using zlib to reduce payload size
compressed_bytes = zlib.compress(mask_bytes)
mask_b64 = base64.b64encode(compressed_bytes).decode('utf-8')
serialized_items.append(SAMItemResponse(
mask_b64=mask_b64,
mask_shape=list(mask.shape),
bbox=it.get("bbox", []),
score=it.get("score"),
area=it.get("area", 0)
))
return SAMResponse(items=serialized_items)
except Exception as e:
import traceback
traceback.print_exc()
raise HTTPException(status_code=500, detail=str(e))
finally:
# Aggressive cleanup to prevent OOM
if torch and torch.cuda.is_available():
torch.cuda.empty_cache()
@app.post("/free_model")
async def free_model(checkpoint: str = "sam_b.pt"):
try:
free_sam_model(checkpoint)
return {"status": "ok", "message": f"Model {checkpoint} freed"}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.get("/health")
def health():
return {"status": "ok"}
|