import os import subprocess import traceback import cv2 import io import open_clip from fastapi import FastAPI, UploadFile, File, HTTPException, Form, Request from fastapi.responses import JSONResponse import numpy as np import onnxruntime as ort from PIL import Image from open_clip import image_transform from transformers import AutoTokenizer from open_clip import image_transform from insightface.model_zoo import get_model from insightface.app.common import Face from starlette.middleware.base import BaseHTTPMiddleware app = FastAPI(title="Photo AI Feature Extraction API", version="1.0.0") # 全局模型实例 detector = None recognizer = None clip_vision_session = None clip_text_session = None # CLIP 专属的预处理和分词器 clip_preprocess = None clip_tokenizer = None # 判断模型目录是否存在 # if not os.path.exists('/code/antelopev2'): # # 调用GIT lfs 下载模型到 /code/models # # https://huggingface.co/ImgSpace/iLookModels # # 获取环境变量判断是通过hf下载还是modelsp下载 # try: # repo_url = "https://huggingface.co/ImgSpace/iLookModels" # if os.environ.get("MODEL_SOURCE") == "modelscope": # repo_url = "https://modelscope.cn/ImgSpace/iLookModels" # subprocess.run(["git", "lfs", "install"], cwd="/code") # subprocess.run(["git", "clone", repo_url, "/code/models"]) # subprocess.run(["mv", "/code/models/antelopev2", "/code/"]) # subprocess.run(["mv", "/code/models/clip_model", "/code/"]) # subprocess.run(["rm", "-rf", "/code/models"]) # except Exception as e: # print("❌ 下载模型失败, 错误如下:\n", traceback.format_exc()) API_TOKEN = os.getenv("TOKEN", "my-secret-token") # 不需要验证的路径 EXCLUDED_PATHS = ["/health"] class TokenAuthMiddleware(BaseHTTPMiddleware): """简单的Token验证中间件""" async def dispatch(self, request: Request, call_next): # 跳过不需要验证的路径 if request.url.path in EXCLUDED_PATHS: return await call_next(request) # 从请求头获取 Authorization auth_header = request.headers.get("Authorization") # 兼容 Bearer 前缀 if auth_header and auth_header.startswith("Bearer "): auth_header = auth_header[7:] # 移除 "Bearer " 前缀 # 验证token是否匹配 if auth_header != API_TOKEN: return JSONResponse( status_code=401, content={"detail": "未授权：无效的认证令牌"} ) # token验证通过，继续处理请求 response = await call_next(request) return response # 添加中间件 app.add_middleware(TokenAuthMiddleware) @app.on_event("startup") async def startup_event(): global detector, recognizer, clip_vision_session, clip_text_session, clip_preprocess, clip_tokenizer BASE_DIR = os.path.dirname(os.path.abspath(__file__)) ANTELOPE_DIR = os.path.join(BASE_DIR, 'antelopev2') CLIP_DIR = os.path.join(BASE_DIR, 'clip_model') # 配置 ONNX Runtime CPU 推理线程，防止单个请求占满所有核心 sess_options = ort.SessionOptions() sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL sess_options.intra_op_num_threads = 4 # 根据实际 CPU 核心数调整 print("====== 开始加载离线模型 ======") try: # 1. 加载 InsightFace 人脸检测 det_path = os.path.join(ANTELOPE_DIR, 'scrfd_10g_bnkps.onnx') if not os.path.exists(det_path): raise FileNotFoundError(f"Missing: {det_path}") detector = get_model(det_path, providers=['CPUExecutionProvider']) detector.prepare(ctx_id=0, input_size=(640, 640)) # 2. 加载 InsightFace 人脸识别 rec_path = os.path.join(ANTELOPE_DIR, 'glintr100.onnx') if not os.path.exists(rec_path): raise FileNotFoundError(f"Missing: {rec_path}") recognizer = get_model(rec_path, providers=['CPUExecutionProvider']) recognizer.prepare(ctx_id=0) # 3. 加载 CLIP 视觉与文本 ONNX 引擎 vi_path = os.path.join(CLIP_DIR, 'vision_model.onnx') tx_path = os.path.join(CLIP_DIR, 'text_model.onnx') if not os.path.exists(vi_path) or not os.path.exists(tx_path): raise FileNotFoundError("CLIP ONNX 模型或其 .data 配置文件缺失。") clip_vision_session = ort.InferenceSession( vi_path, sess_options, providers=['CPUExecutionProvider']) clip_text_session = ort.InferenceSession( tx_path, sess_options, providers=['CPUExecutionProvider']) # 4. 初始化 CLIP 预处理流与分词器 (保持与训练时一致) # 这里借用 open_clip 结构获取标准配置，不下载权重 # 4. 初始化预处理与纯离线分词器 # 视觉预处理：224 是 ViT-B 要求的标准输入尺寸，纯本地计算 clip_preprocess = image_transform(224, is_train=False) # 文本分词器：直接指向刚才创建的本地离线文件夹，绝对不会产生网络请求 TOKENIZER_DIR = os.path.join(CLIP_DIR, 'tokenizer') if not os.path.exists(TOKENIZER_DIR): raise FileNotFoundError(f"找不到离线分词器目录: {TOKENIZER_DIR}") clip_tokenizer = AutoTokenizer.from_pretrained( TOKENIZER_DIR, local_files_only=True) print("====== 🎉 所有模型（人脸 + CLIP）加载成功，API 已就绪！ ======") except Exception as e: print("❌ 初始化失败, 错误如下:\n", traceback.format_exc()) @app.get("/health") def health_check(): models_ready = all( [detector, recognizer, clip_vision_session, clip_text_session]) return {"status": "healthy" if models_ready else "unhealthy"} # ========================================== # 接口 1：仅做人脸检测与识别聚类 # ========================================== @app.post("/api/v1/extract/face") async def extract_face_only(file: UploadFile = File(...)): if not all([detector, recognizer]): raise HTTPException( status_code=500, detail="Face models are not ready.") try: contents = await file.read() nparr = np.frombuffer(contents, np.uint8) img_cv = cv2.imdecode(nparr, cv2.IMREAD_COLOR) if img_cv is None: raise HTTPException(status_code=400, detail="Invalid image file.") bboxes, kpss = detector.detect(img_cv, max_num=0, metric='default') face_results = [] if bboxes is not None and len(bboxes) > 0: for i in range(len(bboxes)): bbox = bboxes[i] kp = kpss[i] if kpss is not None else None face = Face(bbox=bbox[:4], kps=kp, det_score=bbox[4]) recognizer.get(img_cv, face) face_results.append({ "bbox": face.bbox.tolist(), "det_score": float(face.det_score), "kps": face.kps.tolist() if face.kps is not None else None, "embedding": face.embedding.tolist() if face.embedding is not None else None }) return { "face_count": len(face_results), "faces": face_results } except Exception as e: raise HTTPException( status_code=500, detail=f"Face extraction error: {str(e)}") # ========================================== # 接口 2：仅做 CLIP 整图语义特征提取 # ========================================== @app.post("/api/v1/extract/clip_image") async def extract_clip_image(file: UploadFile = File(...)): if clip_vision_session is None: raise HTTPException( status_code=500, detail="CLIP Vision model is not ready.") try: contents = await file.read() img_pil = Image.open(io.BytesIO(contents)).convert("RGB") clip_tensor = clip_preprocess(img_pil).unsqueeze(0).numpy() image_embedding = clip_vision_session.run( None, {"image": clip_tensor})[0] # L2 归一化 image_embedding = image_embedding / \ np.linalg.norm(image_embedding, axis=-1, keepdims=True) return { "image_embedding": image_embedding[0].tolist() } except Exception as e: raise HTTPException( status_code=500, detail=f"CLIP Image extraction error: {str(e)}") # ========================================== # 接口 3：文本搜图特征提取 (保持不变) # ========================================== @app.post("/api/v1/extract/text") async def extract_text_features(text: str = Form(...)): if clip_text_session is None or clip_tokenizer is None: raise HTTPException( status_code=500, detail="Text model is not initialized.") try: if not text.strip(): raise HTTPException( status_code=400, detail="Text query cannot be empty.") # 使用 transformers 的本地分词器，强制对齐 77 长度，并输出 numpy 格式 tokenized = clip_tokenizer( [text], padding='max_length', truncation=True, max_length=77, return_tensors='np' ) text_tokens = tokenized['input_ids'] # 文本特征推理 text_embedding = clip_text_session.run(None, {"text": text_tokens})[0] # L2 归一化 text_embedding = text_embedding / \ np.linalg.norm(text_embedding, axis=-1, keepdims=True) return { "text": text, "text_embedding": text_embedding[0].tolist() } except Exception as e: raise HTTPException( status_code=500, detail=f"Text extraction error: {str(e)}")