""" SAM3 Segmentation API Service (SAM3 图像分割 API 服务) -------------------------------------------------- 本服务提供基于 SAM3 模型的图像分割能力。 包含通用分割、塔罗牌分析和人脸分析等专用接口。 This service provides image segmentation capabilities using the SAM3 model. It includes specialized endpoints for Tarot card analysis and Face analysis. """ # ========================================== # 1. Imports (导入) # ========================================== # Standard Library Imports (标准库) import os import uuid import time import json import traceback import re import asyncio import shutil import subprocess from datetime import datetime from typing import Optional, List, Dict, Any from contextlib import asynccontextmanager # Third-Party Imports (第三方库) import cv2 import torch import numpy as np import requests import matplotlib matplotlib.use('Agg') # 使用非交互式后端，防止在服务器上报错 import matplotlib.pyplot as plt from PIL import Image # FastAPI Imports from fastapi import FastAPI, File, UploadFile, Form, HTTPException, Request, Depends, status, APIRouter, Cookie from fastapi.security import APIKeyHeader from fastapi.staticfiles import StaticFiles from fastapi.responses import JSONResponse, HTMLResponse, Response # Dashscope (Aliyun Qwen) Imports import dashscope from dashscope import MultiModalConversation # Local Imports (本地模块) from sam3.model_builder import build_sam3_image_model from sam3.model.sam3_image_processor import Sam3Processor from sam3.visualization_utils import plot_results import human_analysis_service # 引入新服务: 人脸分析 # ========================================== # 2. Configuration & Constants (配置与常量) # ========================================== # 路径配置 STATIC_DIR = "static" RESULT_IMAGE_DIR = os.path.join(STATIC_DIR, "results") os.makedirs(RESULT_IMAGE_DIR, exist_ok=True) # API Key 配置 VALID_API_KEY = "123quant-speed" API_KEY_HEADER_NAME = "X-API-Key" # Admin Config ADMIN_PASSWORD = "admin_secure_password" # 可以根据需求修改 HISTORY_FILE = "history.json" # Dashscope (Qwen-VL) 配置 dashscope.api_key = 'sk-ce2404f55f744a1987d5ece61c6bac58' QWEN_MODEL = 'qwen-vl-max' # Default model AVAILABLE_QWEN_MODELS = ["qwen-vl-max", "qwen-vl-plus"] # 清理配置 (Cleanup Config) CLEANUP_CONFIG = { "enabled": os.getenv("AUTO_CLEANUP_ENABLED", "True").lower() == "true", "lifetime": int(os.getenv("FILE_LIFETIME_SECONDS", "3600")), "interval": int(os.getenv("CLEANUP_INTERVAL_SECONDS", "600")) } # 提示词配置 (Prompt Config) PROMPTS = { "translate": "请将以下描述翻译成简洁、精准的英文，用于图像分割模型(SAM)的提示词。直接返回英文，不要包含任何解释或其他文字。\n\n输入: {text}", "tarot_card_dual": """这是一张塔罗牌的两个方向： 图1：原始方向 图2：旋转180度后的方向 请仔细对比两张图片的牌面内容（文字方向、人物站立方向、图案逻辑）： 1. 识别这张牌的名字（中文）。 2. 判断哪一张图片展示了正确的“正位”（Upright）状态。 - 如果图1是正位，说明原图就是正位。 - 如果图2是正位，说明原图是逆位。 请以JSON格式返回，包含 'name' 和 'position' 两个字段。 例如：{'name': '愚者', 'position': '正位'} 或 {'name': '倒吊人', 'position': '逆位'}。 不要包含Markdown代码块标记。""", "tarot_card_single": "这是一张塔罗牌。请识别它的名字（中文），并判断它是正位还是逆位。请以JSON格式返回，包含 'name' 和 'position' 两个字段。例如：{'name': '愚者', 'position': '正位'}。不要包含Markdown代码块标记。", "tarot_spread": "这是一张包含多张塔罗牌的图片。请根据牌的排列方式识别这是什么牌阵（例如：圣三角、凯尔特十字、三张牌等）。如果看不出明显的正规牌阵，请返回“不是正规牌阵”。请以JSON格式返回，包含 'spread_name' 和 'description' 两个字段。例如：{'spread_name': '圣三角', 'description': '常见的时间流占卜法'}。不要包含Markdown代码块标记。", "face_analysis": """请仔细观察这张图片中的人物头部/面部特写： 1. 识别性别 (Gender)：男性/女性 2. 预估年龄 (Age)：请给出一个合理的年龄范围，例如 "25-30岁" 3. 简要描述：发型、发色、是否有眼镜等显著特征。 请以 JSON 格式返回，包含 'gender', 'age', 'description' 字段。 不要包含 Markdown 标记。""" } # API Tags (用于文档分类) TAG_GENERAL = "General Segmentation (通用分割)" TAG_TAROT = "Tarot Analysis (塔罗牌分析)" TAG_FACE = "Face Analysis (人脸分析)" # ========================================== # 3. Security & Middleware (安全与中间件) # ========================================== # 定义 Header 认证 Scheme api_key_header = APIKeyHeader(name=API_KEY_HEADER_NAME, auto_error=False) async def verify_api_key(api_key: Optional[str] = Depends(api_key_header)): """ 强制验证 API Key (Enforce API Key Verification) 1. 检查 Header 中是否存在 API Key 2. 验证 API Key 是否匹配 """ if not api_key: raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Missing API Key. Please provide it in the header." ) if api_key != VALID_API_KEY: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, detail="Invalid API Key." ) return True # ========================================== # 4. Lifespan Management (生命周期管理) # ========================================== async def cleanup_old_files(directory: str): """ 后台任务：定期清理过期的图片文件 - 动态读取 CLEANUP_CONFIG 配置 """ print(f"🧹 自动清理任务已启动 | 目录: {directory}") while True: try: # 动态读取配置 interval = CLEANUP_CONFIG["interval"] lifetime = CLEANUP_CONFIG["lifetime"] enabled = CLEANUP_CONFIG["enabled"] await asyncio.sleep(interval) if not enabled: continue current_time = time.time() count = 0 # 遍历所有文件（包括子目录） for root, dirs, files in os.walk(directory): for file in files: file_path = os.path.join(root, file) # 获取文件修改时间 try: file_mtime = os.path.getmtime(file_path) if current_time - file_mtime > lifetime: os.remove(file_path) count += 1 except OSError: pass # 文件可能已被删除 # 尝试清理空目录 (可选，仅清理二级目录) for root, dirs, files in os.walk(directory, topdown=False): for dir in dirs: dir_path = os.path.join(root, dir) try: if not os.listdir(dir_path): # 如果目录为空 os.rmdir(dir_path) except OSError: pass if count > 0: print(f"🧹 已清理 {count} 个过期文件 (Lifetime: {lifetime}s)") except asyncio.CancelledError: print("🛑 清理任务已停止") break except Exception as e: print(f"⚠️ 清理任务出错: {e}") await asyncio.sleep(60) # 出错后等待一分钟再试 @asynccontextmanager async def lifespan(app: FastAPI): """ FastAPI 生命周期管理器 - 启动时: 加载模型到 GPU/CPU - 关闭时: 清理资源 """ print("="*40) print("✅ API Key 保护已激活") print(f"✅ 有效 Key: {VALID_API_KEY}") print("="*40) print("正在加载 SAM3 模型到 GPU...") device = torch.device("cuda" if torch.cuda.is_available() else "cpu") if not torch.cuda.is_available(): print("警告: 未检测到 GPU，将使用 CPU，速度会较慢。") # 加载模型 model = build_sam3_image_model() model = model.to(device) model.eval() # 初始化处理器 processor = Sam3Processor(model) # 存储到 app.state 以便全局访问 app.state.model = model app.state.processor = processor app.state.device = device print(f"模型加载完成，设备: {device}") # --- 启动后台清理任务 --- cleanup_task_handle = None try: cleanup_task_handle = asyncio.create_task( cleanup_old_files(RESULT_IMAGE_DIR) ) except Exception as e: print(f"启动清理任务失败: {e}") # ----------------------- yield print("正在清理资源...") # --- 停止后台任务 --- if cleanup_task_handle: cleanup_task_handle.cancel() try: await cleanup_task_handle except asyncio.CancelledError: pass # ------------------ # 这里可以添加释放显存的逻辑，如果需要 # ========================================== # 5. Helper Functions (辅助函数) # ========================================== def is_english(text: str) -> bool: """ 判断文本是否为纯英文 - 如果包含中文字符范围 (\u4e00-\u9fff)，则返回 False """ for char in text: if '\u4e00' <= char <= '\u9fff': return False return True def append_to_history(req_type: str, prompt: str, status: str, result_path: str = None, details: str = "", final_prompt: str = None, duration: float = 0.0): """ 记录请求历史到 history.json """ record = { "timestamp": time.time(), "type": req_type, "prompt": prompt, "final_prompt": final_prompt, "status": status, "result_path": result_path, "details": details, "duration": duration } try: with open(HISTORY_FILE, "a", encoding="utf-8") as f: f.write(json.dumps(record, ensure_ascii=False) + "\n") except Exception as e: print(f"Failed to write history: {e}") def translate_to_sam3_prompt(text: str) -> str: """ 使用 Qwen 模型将中文提示词翻译为英文 - SAM3 模型对英文 Prompt 支持更好 """ print(f"正在翻译提示词: {text}") try: prompt_template = PROMPTS["translate"] messages = [ { "role": "user", "content": [ {"text": prompt_template.format(text=text)} ] } ] response = MultiModalConversation.call(model=QWEN_MODEL, messages=messages) if response.status_code == 200: translated_text = response.output.choices[0].message.content[0]['text'].strip() # 去除可能的 markdown 标记或引号 translated_text = translated_text.replace('"', '').replace("'", "").strip() print(f"翻译结果: {translated_text}") return translated_text else: print(f"翻译失败: {response.code} - {response.message}") return text # 失败则回退到原始文本 except Exception as e: print(f"翻译异常: {e}") return text def order_points(pts): """ 对四个坐标点进行排序：左上，右上，右下，左下 用于透视变换前的点位整理 """ rect = np.zeros((4, 2), dtype="float32") s = pts.sum(axis=1) rect[0] = pts[np.argmin(s)] rect[2] = pts[np.argmax(s)] diff = np.diff(pts, axis=1) rect[1] = pts[np.argmin(diff)] rect[3] = pts[np.argmax(diff)] return rect def four_point_transform(image, pts): """ 根据四个点进行透视变换 (Perspective Transform) 用于将倾斜的卡片矫正为矩形 """ rect = order_points(pts) (tl, tr, br, bl) = rect # 计算新图像的宽度 widthA = np.sqrt(((br[0] - bl[0]) ** 2) + ((br[1] - bl[1]) ** 2)) widthB = np.sqrt(((tr[0] - tl[0]) ** 2) + ((tr[1] - tl[1]) ** 2)) maxWidth = max(int(widthA), int(widthB)) # 计算新图像的高度 heightA = np.sqrt(((tr[0] - br[0]) ** 2) + ((tr[1] - br[1]) ** 2)) heightB = np.sqrt(((tl[0] - bl[0]) ** 2) + ((tl[1] - bl[1]) ** 2)) maxHeight = max(int(heightA), int(heightB)) dst = np.array([ [0, 0], [maxWidth - 1, 0], [maxWidth - 1, maxHeight - 1], [0, maxHeight - 1]], dtype="float32") M = cv2.getPerspectiveTransform(rect, dst) warped = cv2.warpPerspective(image, M, (maxWidth, maxHeight)) return warped def load_image_from_url(url: str) -> Image.Image: """ 从 URL 下载图片并转换为 PIL Image 对象 """ try: headers = {'User-Agent': 'Mozilla/5.0'} response = requests.get(url, headers=headers, stream=True, timeout=10) response.raise_for_status() image = Image.open(response.raw).convert("RGB") return image except Exception as e: raise HTTPException(status_code=400, detail=f"无法下载图片: {str(e)}") def crop_and_save_objects(image: Image.Image, masks, boxes, output_dir: str = RESULT_IMAGE_DIR, is_tarot: bool = True, cutout: bool = False, perspective_correction: bool = False) -> list[dict]: """ 根据 mask 和 box 进行处理并保存独立的对象图片 参数: - image: 原始图片 - masks: 分割掩码列表 - boxes: 边界框列表 - output_dir: 输出目录 - is_tarot: 是否为塔罗牌模式 (会影响文件名前缀和旋转逻辑) - cutout: 如果为 True，则进行轮廓抠图（透明背景）；否则进行透视矫正（主要用于卡片） - perspective_correction: 是否进行梯度透视矫正 返回: - 保存的对象信息列表 """ saved_objects = [] # Convert image to numpy array (RGB) img_arr = np.array(image) for i, (mask, box) in enumerate(zip(masks, boxes)): # Handle tensor/numpy conversions if isinstance(mask, torch.Tensor): mask_np = mask.cpu().numpy().squeeze() else: mask_np = mask.squeeze() # Handle box conversion if isinstance(box, torch.Tensor): box_np = box.cpu().numpy() else: box_np = np.array(box) # Ensure mask is uint8 binary for OpenCV/Pillow if mask_np.dtype == bool: mask_uint8 = (mask_np * 255).astype(np.uint8) else: mask_uint8 = (mask_np > 0.5).astype(np.uint8) * 255 # --- 准备基础图像 --- if cutout: # 1. 准备 RGBA 原图 if image.mode != "RGBA": img_rgba = image.convert("RGBA") else: img_rgba = image.copy() # 2. 准备 Alpha Mask mask_img = Image.fromarray(mask_uint8, mode='L') # 3. 将 Mask 应用到 Alpha 通道 base_img_pil = Image.new("RGBA", img_rgba.size, (0, 0, 0, 0)) base_img_pil.paste(image.convert("RGB"), (0, 0), mask=mask_img) # Convert to numpy for potential warping base_img_arr = np.array(base_img_pil) else: base_img_pil = image.convert("RGB") base_img_arr = img_arr # RGB numpy array # --- 透视矫正 vs 简单裁剪 --- final_img_pil = None is_rotated = False note = "" if perspective_correction: # --- 透视矫正模式 (矩形矫正) --- contours, _ = cv2.findContours(mask_uint8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) pts = None if contours: c = max(contours, key=cv2.contourArea) peri = cv2.arcLength(c, True) approx = cv2.approxPolyDP(c, 0.04 * peri, True) if len(approx) == 4: pts = approx.reshape(4, 2) else: rect = cv2.minAreaRect(c) pts = cv2.boxPoints(rect) if pts is not None: warped = four_point_transform(base_img_arr, pts) note = "Geometric correction applied." else: # Fallback to simple crop if no contours found x1, y1, x2, y2 = map(int, box_np) # Ensure bounds h, w = base_img_arr.shape[:2] x1 = max(0, x1); y1 = max(0, y1) x2 = min(w, x2); y2 = min(h, y2) warped = base_img_arr[y1:y2, x1:x2] note = "Correction failed, fallback to crop." # Check orientation (Portrait vs Landscape) - Only for Tarot usually h, w = warped.shape[:2] # 强制竖屏逻辑 (塔罗牌通常是竖屏) if is_tarot and w > h: warped = cv2.rotate(warped, cv2.ROTATE_90_CLOCKWISE) is_rotated = True final_img_pil = Image.fromarray(warped) else: # --- 简单裁剪模式 (Simple Crop) --- x1, y1, x2, y2 = map(int, box_np) w, h = base_img_pil.size x1 = max(0, x1); y1 = max(0, y1) x2 = min(w, x2); y2 = min(h, y2) if x2 > x1 and y2 > y1: final_img_pil = base_img_pil.crop((x1, y1, x2, y2)) else: final_img_pil = base_img_pil # Fallback note = "Simple crop applied." # --- 保存图片 --- prefix = "cutout" if cutout else ("tarot" if is_tarot else "segment") filename = f"{prefix}_{uuid.uuid4().hex}_{i}.png" save_path = os.path.join(output_dir, filename) final_img_pil.save(save_path) saved_objects.append({ "filename": filename, "is_rotated_by_algorithm": is_rotated, "note": note }) return saved_objects def generate_and_save_result(image: Image.Image, inference_state, output_dir: str = RESULT_IMAGE_DIR) -> str: """ 生成并保存包含 Mask 叠加效果的完整结果图 """ filename = f"seg_{uuid.uuid4().hex}.jpg" save_path = os.path.join(output_dir, filename) plot_results(image, inference_state) plt.savefig(save_path, dpi=150, bbox_inches='tight') plt.close() return filename def recognize_card_with_qwen(image_path: str) -> dict: """ 调用 Qwen-VL 识别塔罗牌 (采用正逆位对比策略) """ try: abs_path = os.path.abspath(image_path) file_url = f"file://{abs_path}" try: # 1. 打开原图 img = Image.open(abs_path) # 2. 生成旋转图 (180度) rotated_img = img.rotate(180) # 3. 保存旋转图 dir_name = os.path.dirname(abs_path) file_name = os.path.basename(abs_path) rotated_name = f"rotated_{file_name}" rotated_path = os.path.join(dir_name, rotated_name) rotated_img.save(rotated_path) rotated_file_url = f"file://{rotated_path}" # 4. 构建对比 Prompt messages = [ { "role": "user", "content": [ {"image": file_url}, # 图1 (原图) {"image": rotated_file_url}, # 图2 (旋转180度) {"text": PROMPTS["tarot_card_dual"]} ] } ] except Exception as e: print(f"对比图生成失败，回退到单图模式: {e}") messages = [ { "role": "user", "content": [ {"image": file_url}, {"text": PROMPTS["tarot_card_single"]} ] } ] response = MultiModalConversation.call(model=QWEN_MODEL, messages=messages) if response.status_code == 200: content = response.output.choices[0].message.content[0]['text'] import json try: clean_content = content.replace("```json", "").replace("```", "").strip() result = json.loads(clean_content) return result except: return {"raw_response": content} else: return {"error": f"API Error: {response.code} - {response.message}"} except Exception as e: return {"error": f"识别失败: {str(e)}"} def recognize_spread_with_qwen(image_path: str) -> dict: """ 调用 Qwen-VL 识别塔罗牌牌阵 """ try: abs_path = os.path.abspath(image_path) file_url = f"file://{abs_path}" messages = [ { "role": "user", "content": [ {"image": file_url}, {"text": PROMPTS["tarot_spread"]} ] } ] response = MultiModalConversation.call(model=QWEN_MODEL, messages=messages) if response.status_code == 200: content = response.output.choices[0].message.content[0]['text'] import json try: clean_content = content.replace("```json", "").replace("```", "").strip() result = json.loads(clean_content) return result except: return {"raw_response": content, "spread_name": "Unknown"} else: return {"error": f"API Error: {response.code} - {response.message}"} except Exception as e: return {"error": f"牌阵识别失败: {str(e)}"} # ========================================== # 6. FastAPI App Setup (应用初始化) # ========================================== app = FastAPI( lifespan=lifespan, title="SAM3 Segmentation API", description="## 🔒 受 API Key 保护\n请点击右上角 **Authorize** 并输入: `123quant-*****`", openapi_tags=[ {"name": TAG_GENERAL, "description": "General purpose image segmentation"}, {"name": TAG_TAROT, "description": "Specialized endpoints for Tarot card recognition"}, {"name": TAG_FACE, "description": "Face detection and analysis"}, ] ) # 手动添加 OpenAPI 安全配置 app.openapi_schema = None def custom_openapi(): if app.openapi_schema: return app.openapi_schema from fastapi.openapi.utils import get_openapi openapi_schema = get_openapi( title=app.title, version=app.version, description=app.description, routes=app.routes, ) openapi_schema["components"]["securitySchemes"] = { "APIKeyHeader": { "type": "apiKey", "in": "header", "name": API_KEY_HEADER_NAME, } } for path in openapi_schema["paths"]: for method in openapi_schema["paths"][path]: openapi_schema["paths"][path][method]["security"] = [{"APIKeyHeader": []}] app.openapi_schema = openapi_schema return app.openapi_schema app.openapi = custom_openapi app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static") # ========================================== # 7. API Endpoints (接口定义) # ========================================== # ------------------------------------------ # Group 1: General Segmentation (通用分割) # ------------------------------------------ @app.post("/segment", tags=[TAG_GENERAL], dependencies=[Depends(verify_api_key)]) async def segment( request: Request, prompt: str = Form(..., description="Text prompt for segmentation (e.g., 'cat', 'person')"), file: Optional[UploadFile] = File(None, description="Image file to upload"), image_url: Optional[str] = Form(None, description="URL of the image"), save_segment_images: bool = Form(False, description="Whether to save and return individual segmented objects"), cutout: bool = Form(False, description="If True, returns transparent background PNGs; otherwise returns original crops"), perspective_correction: bool = Form(False, description="If True, applies perspective correction (warping) to the segmented object."), highlight: bool = Form(False, description="If True, darkens the background to highlight the subject (周边变黑放大)."), confidence: float = Form(0.7, description="Confidence threshold (0.0-1.0). Default is 0.7.") ): """ **通用图像分割接口** - 支持上传图片或提供图片 URL - 支持自动将中文 Prompt 翻译为英文 - 支持周边变黑放大效果 (Highlight Mode) - 支持手动设置置信度 (Confidence Threshold) - 支持透视矫正 (Perspective Correction) """ if not file and not image_url: raise HTTPException(status_code=400, detail="必须提供 file (图片文件) 或 image_url (图片链接)") start_time = time.time() # 1. Prompt 处理 final_prompt = prompt if not is_english(prompt): try: translated = translate_to_sam3_prompt(prompt) if translated: final_prompt = translated except Exception as e: print(f"Prompt翻译失败，使用原始Prompt: {e}") print(f"最终使用的 Prompt: {final_prompt}") # 2. 图片加载 try: if file: image = Image.open(file.file).convert("RGB") elif image_url: image = load_image_from_url(image_url) except Exception as e: duration = time.time() - start_time append_to_history("general", prompt, "failed", details=f"Image Load Error: {str(e)}", final_prompt=final_prompt, duration=duration) raise HTTPException(status_code=400, detail=f"图片解析失败: {str(e)}") processor = request.app.state.processor # 3. 模型推理 try: # 设置图片 inference_state = processor.set_image(image) # 处理 Confidence Threshold # 如果用户提供了 confidence，临时修改 processor 的阈值 original_confidence = processor.confidence_threshold if confidence is not None: # 简单校验 if not (0.0 <= confidence <= 1.0): raise HTTPException(status_code=400, detail="Confidence must be between 0.0 and 1.0") processor.confidence_threshold = confidence print(f"Using manual confidence threshold: {confidence}") try: # 执行推理 output = processor.set_text_prompt(state=inference_state, prompt=final_prompt) masks, boxes, scores = output["masks"], output["boxes"], output["scores"] finally: # 恢复默认阈值，防止影响其他请求 (虽然在单进程模型下可能不重要，但保持状态一致性是个好习惯) if confidence is not None: processor.confidence_threshold = original_confidence except Exception as e: duration = time.time() - start_time append_to_history("general", prompt, "failed", details=f"Inference Error: {str(e)}", final_prompt=final_prompt, duration=duration) raise HTTPException(status_code=500, detail=f"模型推理错误: {str(e)}") # 4. 结果可视化与保存 try: if highlight: filename = f"seg_highlight_{uuid.uuid4().hex}.jpg" save_path = os.path.join(RESULT_IMAGE_DIR, filename) # 使用 human_analysis_service 中的可视化函数 (周边变黑) human_analysis_service.create_highlighted_visualization(image, masks, save_path) else: filename = generate_and_save_result(image, inference_state) except Exception as e: duration = time.time() - start_time append_to_history("general", prompt, "failed", details=f"Save Error: {str(e)}", final_prompt=final_prompt, duration=duration) raise HTTPException(status_code=500, detail=f"绘图保存错误: {str(e)}") file_url = request.url_for("static", path=f"results/{filename}") # 5. 保存分割子图 (Optional) saved_segments_info = [] if save_segment_images: try: request_id = f"{int(time.time())}_{uuid.uuid4().hex[:8]}" output_dir = os.path.join(RESULT_IMAGE_DIR, request_id) os.makedirs(output_dir, exist_ok=True) saved_objects = crop_and_save_objects( image, masks, boxes, output_dir=output_dir, is_tarot=False, cutout=cutout, perspective_correction=perspective_correction ) for obj in saved_objects: fname = obj["filename"] seg_url = str(request.url_for("static", path=f"results/{request_id}/{fname}")) saved_segments_info.append({ "url": seg_url, "filename": fname, "note": obj.get("note", "") }) except Exception as e: print(f"Error saving segments: {e}") # Don't fail the whole request just for this part, but log it? Or fail? Usually fail. duration = time.time() - start_time append_to_history("general", prompt, "partial_success", result_path=f"results/{filename}", details="Segments save failed", final_prompt=final_prompt, duration=duration) raise HTTPException(status_code=500, detail=f"保存分割图片失败: {str(e)}") response_content = { "status": "success", "result_image_url": str(file_url), "detected_count": len(masks), "scores": scores.tolist() if torch.is_tensor(scores) else scores } if save_segment_images: response_content["segmented_images"] = saved_segments_info duration = time.time() - start_time append_to_history("general", prompt, "success", result_path=f"results/{filename}", details=f"Detected: {len(masks)}", final_prompt=final_prompt, duration=duration) return JSONResponse(content=response_content) # ------------------------------------------ # Group 2: Tarot Analysis (塔罗牌分析) # ------------------------------------------ @app.post("/segment_tarot", tags=[TAG_TAROT], dependencies=[Depends(verify_api_key)]) async def segment_tarot( request: Request, file: Optional[UploadFile] = File(None), image_url: Optional[str] = Form(None), expected_count: int = Form(3) ): """ **塔罗牌分割接口** 1. 检测是否包含指定数量的塔罗牌 (默认为 3) 2. 对检测到的卡片进行透视矫正和裁剪 3. 返回矫正后的图片 URL """ if not file and not image_url: raise HTTPException(status_code=400, detail="必须提供 file (图片文件) 或 image_url (图片链接)") start_time = time.time() try: if file: image = Image.open(file.file).convert("RGB") elif image_url: image = load_image_from_url(image_url) except Exception as e: duration = time.time() - start_time append_to_history("tarot", f"expected: {expected_count}", "failed", details=f"Image Load Error: {str(e)}", duration=duration) raise HTTPException(status_code=400, detail=f"图片解析失败: {str(e)}") processor = request.app.state.processor try: inference_state = processor.set_image(image) # 固定 Prompt 检测塔罗牌 output = processor.set_text_prompt(state=inference_state, prompt="tarot card") masks, boxes, scores = output["masks"], output["boxes"], output["scores"] except Exception as e: duration = time.time() - start_time append_to_history("tarot", f"expected: {expected_count}", "failed", details=f"Inference Error: {str(e)}", duration=duration) raise HTTPException(status_code=500, detail=f"模型推理错误: {str(e)}") # 核心逻辑：判断数量 detected_count = len(masks) # 创建本次请求的独立文件夹 request_id = f"{int(time.time())}_{uuid.uuid4().hex[:8]}" output_dir = os.path.join(RESULT_IMAGE_DIR, request_id) os.makedirs(output_dir, exist_ok=True) if detected_count != expected_count: # 保存一张图用于调试/反馈 try: filename = generate_and_save_result(image, inference_state, output_dir=output_dir) file_url = request.url_for("static", path=f"results/{request_id}/{filename}") except: file_url = None duration = time.time() - start_time append_to_history("tarot", f"expected: {expected_count}", "failed", result_path=f"results/{request_id}/{filename}" if file_url else None, details=f"Detected {detected_count} cards, expected {expected_count}", duration=duration) return JSONResponse( status_code=400, content={ "status": "failed", "message": f"检测到 {detected_count} 个目标，需要严格的 {expected_count} 张塔罗牌。请调整拍摄角度或背景。", "detected_count": detected_count, "debug_image_url": str(file_url) if file_url else None } ) # 数量正确，执行抠图 try: saved_objects = crop_and_save_objects(image, masks, boxes, output_dir=output_dir, is_tarot=True, perspective_correction=True) except Exception as e: duration = time.time() - start_time append_to_history("tarot", f"expected: {expected_count}", "failed", details=f"Crop Error: {str(e)}", duration=duration) raise HTTPException(status_code=500, detail=f"抠图处理错误: {str(e)}") # 生成 URL 列表和元数据 tarot_cards = [] for obj in saved_objects: fname = obj["filename"] file_url = str(request.url_for("static", path=f"results/{request_id}/{fname}")) tarot_cards.append({ "url": file_url, "is_rotated": obj["is_rotated_by_algorithm"], "orientation_status": "corrected_to_portrait" if obj["is_rotated_by_algorithm"] else "original_portrait", "note": obj["note"] }) # 生成整体效果图 try: main_filename = generate_and_save_result(image, inference_state, output_dir=output_dir) main_file_url = str(request.url_for("static", path=f"results/{request_id}/{main_filename}")) except: main_file_url = None duration = time.time() - start_time append_to_history("tarot", f"expected: {expected_count}", "success", result_path=f"results/{request_id}/{main_filename}" if main_file_url else None, details=f"Successfully segmented {expected_count} cards", duration=duration) return JSONResponse(content={ "status": "success", "message": f"成功识别并分割 {expected_count} 张塔罗牌 (已执行透视矫正)", "tarot_cards": tarot_cards, "full_visualization": main_file_url, "scores": scores.tolist() if torch.is_tensor(scores) else scores }) @app.post("/recognize_tarot", tags=[TAG_TAROT], dependencies=[Depends(verify_api_key)]) async def recognize_tarot( request: Request, file: Optional[UploadFile] = File(None), image_url: Optional[str] = Form(None), expected_count: int = Form(3) ): """ **塔罗牌全流程接口: 分割 + 矫正 + 识别** 1. 检测是否包含指定数量的塔罗牌 (SAM3) 2. 分割并透视矫正 3. 调用 Qwen-VL 识别每张牌的名称和正逆位 """ if not file and not image_url: raise HTTPException(status_code=400, detail="必须提供 file (图片文件) 或 image_url (图片链接)") start_time = time.time() try: if file: image = Image.open(file.file).convert("RGB") elif image_url: image = load_image_from_url(image_url) except Exception as e: duration = time.time() - start_time append_to_history("tarot-recognize", f"expected: {expected_count}", "failed", details=f"Image Load Error: {str(e)}", duration=duration) raise HTTPException(status_code=400, detail=f"图片解析失败: {str(e)}") processor = request.app.state.processor try: inference_state = processor.set_image(image) output = processor.set_text_prompt(state=inference_state, prompt="tarot card") masks, boxes, scores = output["masks"], output["boxes"], output["scores"] except Exception as e: duration = time.time() - start_time append_to_history("tarot-recognize", f"expected: {expected_count}", "failed", details=f"Inference Error: {str(e)}", duration=duration) raise HTTPException(status_code=500, detail=f"模型推理错误: {str(e)}") detected_count = len(masks) request_id = f"{int(time.time())}_{uuid.uuid4().hex[:8]}" output_dir = os.path.join(RESULT_IMAGE_DIR, request_id) os.makedirs(output_dir, exist_ok=True) # 保存整体效果图 try: main_filename = generate_and_save_result(image, inference_state, output_dir=output_dir) main_file_path = os.path.join(output_dir, main_filename) main_file_url = str(request.url_for("static", path=f"results/{request_id}/{main_filename}")) except: main_filename = None main_file_path = None main_file_url = None # Step 0: 牌阵识别 spread_info = {"spread_name": "Unknown"} if main_file_path: # 使用原始图的一份拷贝给 Qwen 识别牌阵 temp_raw_path = os.path.join(output_dir, "raw_for_spread.jpg") image.save(temp_raw_path) spread_info = recognize_spread_with_qwen(temp_raw_path) if detected_count != expected_count: duration = time.time() - start_time append_to_history("tarot-recognize", f"expected: {expected_count}", "failed", result_path=f"results/{request_id}/{main_filename}" if main_file_url else None, details=f"Detected {detected_count}, expected {expected_count}", duration=duration) return JSONResponse( status_code=400, content={ "status": "failed", "message": f"检测到 {detected_count} 个目标，需要严格的 {expected_count} 张塔罗牌。请调整拍摄角度或背景。", "detected_count": detected_count, "spread_info": spread_info, "debug_image_url": str(main_file_url) if main_file_url else None } ) # 数量正确，执行抠图 + 矫正 try: saved_objects = crop_and_save_objects(image, masks, boxes, output_dir=output_dir, is_tarot=True, perspective_correction=True) except Exception as e: duration = time.time() - start_time append_to_history("tarot-recognize", f"expected: {expected_count}", "failed", details=f"Crop Error: {str(e)}", duration=duration) raise HTTPException(status_code=500, detail=f"抠图处理错误: {str(e)}") # 遍历每张卡片进行识别 tarot_cards = [] for obj in saved_objects: fname = obj["filename"] file_path = os.path.join(output_dir, fname) # 调用 Qwen-VL 识别 (串行) recognition_res = recognize_card_with_qwen(file_path) file_url = str(request.url_for("static", path=f"results/{request_id}/{fname}")) tarot_cards.append({ "url": file_url, "is_rotated": obj["is_rotated_by_algorithm"], "orientation_status": "corrected_to_portrait" if obj["is_rotated_by_algorithm"] else "original_portrait", "recognition": recognition_res, "note": obj["note"] }) duration = time.time() - start_time append_to_history("tarot-recognize", f"expected: {expected_count}", "success", result_path=f"results/{request_id}/{main_filename}" if main_file_url else None, details=f"Spread: {spread_info.get('spread_name', 'Unknown')}", duration=duration) return JSONResponse(content={ "status": "success", "message": f"成功识别并分割 {expected_count} 张塔罗牌 (含Qwen识别结果)", "spread_info": spread_info, "tarot_cards": tarot_cards, "full_visualization": main_file_url, "scores": scores.tolist() if torch.is_tensor(scores) else scores }) # ------------------------------------------ # Group 3: Face Analysis (人脸分析) # ------------------------------------------ @app.post("/segment_face", tags=[TAG_FACE], dependencies=[Depends(verify_api_key)]) async def segment_face( request: Request, file: Optional[UploadFile] = File(None), image_url: Optional[str] = Form(None), prompt: str = Form("face and hair", description="Prompt for face detection") ): """ **人脸/头部检测与属性分析接口** 1. 调用 SAM3 分割出头部区域 (含头发) 2. 裁剪并保存 3. 调用 Qwen-VL 识别性别和年龄 """ if not file and not image_url: raise HTTPException(status_code=400, detail="必须提供 file (图片文件) 或 image_url (图片链接)") start_time = time.time() # Prompt 翻译/优化 final_prompt = prompt if not is_english(prompt): try: translated = translate_to_sam3_prompt(prompt) if translated: final_prompt = translated except Exception as e: print(f"Prompt翻译失败，使用原始Prompt: {e}") print(f"Face Segment 最终使用的 Prompt: {final_prompt}") # 加载图片 try: if file: image = Image.open(file.file).convert("RGB") elif image_url: image = load_image_from_url(image_url) except Exception as e: duration = time.time() - start_time append_to_history("face", prompt, "failed", details=f"Image Load Error: {str(e)}", final_prompt=final_prompt, duration=duration) raise HTTPException(status_code=400, detail=f"图片解析失败: {str(e)}") processor = request.app.state.processor # 调用独立服务进行处理 try: result = human_analysis_service.process_face_segmentation_and_analysis( processor=processor, image=image, prompt=final_prompt, output_base_dir=RESULT_IMAGE_DIR, qwen_model=QWEN_MODEL, analysis_prompt=PROMPTS["face_analysis"] ) except Exception as e: import traceback traceback.print_exc() duration = time.time() - start_time append_to_history("face", prompt, "failed", details=f"Process Error: {str(e)}", final_prompt=final_prompt, duration=duration) raise HTTPException(status_code=500, detail=f"处理失败: {str(e)}") # 补全 URL if result["status"] == "success": if result.get("full_visualization"): full_vis_rel_path = result["full_visualization"] result["full_visualization"] = str(request.url_for("static", path=full_vis_rel_path)) for item in result["results"]: relative_path = item.pop("relative_path") item["url"] = str(request.url_for("static", path=relative_path)) duration = time.time() - start_time append_to_history("face", prompt, result["status"], details=f"Results: {len(result.get('results', []))}", final_prompt=final_prompt, duration=duration) return JSONResponse(content=result) # ========================================== # 9. Admin Management APIs (管理后台接口) # ========================================== @app.get("/admin", include_in_schema=False) async def admin_page(): """ Serve the admin HTML page """ # 检查 static/admin.html 是否存在，不存在则返回错误或简易页面 admin_html_path = os.path.join(STATIC_DIR, "admin.html") if os.path.exists(admin_html_path): with open(admin_html_path, "r", encoding="utf-8") as f: content = f.read() return HTMLResponse(content=content) else: return HTMLResponse(content="

Admin page not found

", status_code=404) @app.post("/admin/login", include_in_schema=False) async def admin_login(password: str = Form(...)): """ Simple Admin Login """ if password == ADMIN_PASSWORD: content = {"status": "success", "message": "Logged in"} response = JSONResponse(content=content) # Set a simple cookie response.set_cookie(key="admin_token", value="logged_in", httponly=True) return response else: return JSONResponse(status_code=401, content={"status": "error", "message": "Invalid password"}) async def verify_admin(request: Request): # Check cookie or header token = request.cookies.get("admin_token") # Also allow local dev without strict check if needed, but here we enforce password if token != "logged_in": raise HTTPException(status_code=401, detail="Unauthorized") @app.get("/admin/api/history", dependencies=[Depends(verify_admin)]) async def get_history(): """ Get request history """ if not os.path.exists(HISTORY_FILE): return [] records = [] try: with open(HISTORY_FILE, "r", encoding="utf-8") as f: for line in f: if line.strip(): try: records.append(json.loads(line)) except: pass # Limit to last 100 records return records[-100:] except Exception as e: return {"error": str(e)} @app.get("/admin/api/files", dependencies=[Depends(verify_admin)]) async def list_files(path: str = ""): """ List files in static/results """ # Security check: prevent directory traversal if ".." in path or path.startswith("/"): raise HTTPException(status_code=400, detail="Invalid path") target_dir = os.path.join(RESULT_IMAGE_DIR, path) if not os.path.exists(target_dir): return [] items = [] try: for entry in os.scandir(target_dir): is_dir = entry.is_dir() item = { "name": entry.name, "is_dir": is_dir, "path": os.path.join(path, entry.name), "mtime": entry.stat().st_mtime } if is_dir: try: item["count"] = len(os.listdir(entry.path)) except: item["count"] = 0 else: item["size"] = entry.stat().st_size # Construct URL # Assuming static mount is /static # path is relative to results/ # so url is /static/results/path/name rel_path = os.path.join("results", path, entry.name) # Ensure forward slashes for URL if os.sep != "/": rel_path = rel_path.replace(os.sep, "/") item["url"] = f"/static/{rel_path}" items.append(item) # Sort: Directories first, then by time (newest first) items.sort(key=lambda x: (not x["is_dir"], -x["mtime"])) return items except Exception as e: raise HTTPException(status_code=500, detail=str(e)) @app.delete("/admin/api/files/{file_path:path}", dependencies=[Depends(verify_admin)]) async def delete_file(file_path: str): """ Delete file or directory """ # Security check if ".." in file_path: raise HTTPException(status_code=400, detail="Invalid path") # file_path is relative to static/results/ (or passed as full relative path from API) # The API is called with relative path from current view target_path = os.path.join(RESULT_IMAGE_DIR, file_path) if not os.path.exists(target_path): raise HTTPException(status_code=404, detail="Not found") try: if os.path.isdir(target_path): shutil.rmtree(target_path) else: os.remove(target_path) return {"status": "success", "message": f"Deleted {file_path}"} except Exception as e: raise HTTPException(status_code=500, detail=str(e)) @app.post("/admin/api/cleanup", dependencies=[Depends(verify_admin)]) async def trigger_cleanup(): """ Manually trigger cleanup """ try: # Use global config lifetime = CLEANUP_CONFIG["lifetime"] count = 0 current_time = time.time() for root, dirs, files in os.walk(RESULT_IMAGE_DIR): for file in files: file_path = os.path.join(root, file) try: file_mtime = os.path.getmtime(file_path) if current_time - file_mtime > lifetime: os.remove(file_path) count += 1 except: pass # Cleanup empty dirs for root, dirs, files in os.walk(RESULT_IMAGE_DIR, topdown=False): for dir in dirs: dir_path = os.path.join(root, dir) try: if not os.listdir(dir_path): os.rmdir(dir_path) except: pass return {"status": "success", "message": f"Cleaned {count} files (Lifetime: {lifetime}s)"} except Exception as e: return {"status": "error", "message": str(e)} @app.get("/admin/api/config", dependencies=[Depends(verify_admin)]) async def get_config(request: Request): """ Get system config info """ device_str = "Unknown" gpu_status = {} if hasattr(request.app.state, "device"): device_str = str(request.app.state.device) # 获取 GPU 详细信息 if torch.cuda.is_available(): try: device_id = torch.cuda.current_device() props = torch.cuda.get_device_properties(device_id) total_mem = props.total_memory reserved_mem = torch.cuda.memory_reserved(device_id) allocated_mem = torch.cuda.memory_allocated(device_id) gpu_status = { "available": True, "name": props.name, "total_memory": f"{total_mem / 1024**3:.2f} GB", "reserved_memory": f"{reserved_mem / 1024**3:.2f} GB", "allocated_memory": f"{allocated_mem / 1024**3:.2f} GB", "memory_usage_percent": round((reserved_mem / total_mem) * 100, 1) } except Exception as e: gpu_status = {"available": True, "error": str(e)} else: gpu_status = {"available": False, "reason": "No CUDA device detected"} return { "device": device_str, "gpu_status": gpu_status, "cleanup_config": CLEANUP_CONFIG, "current_qwen_model": QWEN_MODEL, "available_qwen_models": AVAILABLE_QWEN_MODELS } @app.post("/admin/api/config/cleanup", dependencies=[Depends(verify_admin)]) async def update_cleanup_config( enabled: bool = Form(...), lifetime: int = Form(...), interval: int = Form(...) ): """ Update cleanup configuration """ global CLEANUP_CONFIG CLEANUP_CONFIG["enabled"] = enabled CLEANUP_CONFIG["lifetime"] = lifetime CLEANUP_CONFIG["interval"] = interval print(f"Updated Cleanup Config: {CLEANUP_CONFIG}") return {"status": "success", "message": "Cleanup configuration updated", "config": CLEANUP_CONFIG} @app.post("/admin/api/config/model", dependencies=[Depends(verify_admin)]) async def set_model(model: str = Form(...)): """ Set the Qwen model """ global QWEN_MODEL if model not in AVAILABLE_QWEN_MODELS: raise HTTPException(status_code=400, detail="Invalid model") QWEN_MODEL = model return {"status": "success", "message": f"Model switched to {model}", "current_model": QWEN_MODEL} @app.get("/admin/api/prompts", dependencies=[Depends(verify_admin)]) async def get_prompts(): """ Get all prompts """ return PROMPTS @app.post("/admin/api/prompts", dependencies=[Depends(verify_admin)]) async def update_prompts( key: str = Form(...), content: str = Form(...) ): """ Update a specific prompt """ if key not in PROMPTS: raise HTTPException(status_code=400, detail="Invalid prompt key") PROMPTS[key] = content return {"status": "success", "message": f"Prompt '{key}' updated"} # ------------------------------------------ # GPU Status Helper & API # ------------------------------------------ def get_gpu_status_smi(): """ Get detailed GPU status using nvidia-smi Returns: dict with utilization, memory, temp, power, etc. """ cuda_version = "Unknown" try: import torch if torch.version.cuda: cuda_version = torch.version.cuda except: pass try: # Check if nvidia-smi is available # Fields: utilization.gpu, utilization.memory, temperature.gpu, power.draw, power.limit, memory.total, memory.used, memory.free, name, driver_version result = subprocess.run( ['nvidia-smi', '--query-gpu=utilization.gpu,utilization.memory,temperature.gpu,power.draw,power.limit,memory.total,memory.used,memory.free,name,driver_version', '--format=csv,noheader,nounits'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, encoding='utf-8' ) if result.returncode != 0: raise Exception("nvidia-smi failed") # Parse the first line (assuming single GPU for now, or take the first one) line = result.stdout.strip().split('\n')[0] vals = [x.strip() for x in line.split(',')] return { "available": True, "gpu_util": float(vals[0]), # % "mem_util": float(vals[1]), # % (controller utilization) "temperature": float(vals[2]), # C "power_draw": float(vals[3]), # W "power_limit": float(vals[4]), # W "mem_total": float(vals[5]), # MB "mem_used": float(vals[6]), # MB "mem_free": float(vals[7]), # MB "name": vals[8], "driver_version": vals[9], "cuda_version": cuda_version, "source": "nvidia-smi", "timestamp": time.time() } except Exception as e: # Fallback to torch if available if torch.cuda.is_available(): try: device_id = torch.cuda.current_device() props = torch.cuda.get_device_properties(device_id) mem_reserved = torch.cuda.memory_reserved(device_id) / 1024**2 # MB mem_total = props.total_memory / 1024**2 # MB return { "available": True, "gpu_util": 0, # Torch can't get this easily "mem_util": (mem_reserved / mem_total) * 100, "temperature": 0, "power_draw": 0, "power_limit": 0, "mem_total": mem_total, "mem_used": mem_reserved, "mem_free": mem_total - mem_reserved, "name": props.name, "driver_version": "Unknown", "cuda_version": cuda_version, "source": "torch", "timestamp": time.time() } except: pass return {"available": False, "error": str(e)} @app.get("/admin/api/gpu/status", dependencies=[Depends(verify_admin)]) async def get_gpu_status_api(): """ Get real-time GPU status """ return get_gpu_status_smi() # ========================================== # 10. Main Entry Point (启动入口) # ========================================== if __name__ == "__main__": import uvicorn # ========================================== # 自动清理图片配置 (Auto Cleanup Config) # ========================================== # 设置是否开启自动清理 (True/False) os.environ["AUTO_CLEANUP_ENABLED"] = "True" # 设置图片生命周期 (秒)，超过此时间的图片将被删除 # 例如: 3600 = 1小时, 86400 = 1天 os.environ["FILE_LIFETIME_SECONDS"] = "3600" # 设置检查间隔 (秒)，每隔多久检查一次 os.environ["CLEANUP_INTERVAL_SECONDS"] = "600" # ========================================== # 启动服务器 uvicorn.run( "fastAPI_tarot:app", host="127.0.0.1", port=55600, proxy_headers=True, forwarded_allow_ips="*", reload=False )