Update main.py

main.py

@@ -1,99 +1,247 @@
-import gc
-import os
-from typing import List
-
-import cv2
-import numpy as np
-from captcha_recognizer.slider import Slider
-from fastapi import FastAPI, File, UploadFile
-from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel
+"""
+滑块验证码识别 API 服务
+========================
+基于 FastAPI + ONNX 模型，提供 HTTP 接口识别滑块验证码的滑块和缺口位置。
+
+接口说明：
+  GET  /                  - 健康检查
+  GET  /health            - 健康检查
+  POST /captcha           - 旧版：上传图片文件，只返回缺口（向后兼容）
+  POST /captcha/base64    - 旧版：上传 base64 图片，只返回缺口
+  POST /captcha/v2        - 新版：上传图片文件，返回滑块+缺口+偏移量
+  POST /captcha/v2/base64 - 新版：上传 base64 图片，返回滑块+缺口+偏移量
+
+旧版返回格式（/captcha、/captcha/base64）：
+  {
+    "box": [x1, y1, x2, y2],  // 缺口坐标
+    "confidence": 0.95,        // 置信度
+    "message": null            // 错误信息
+  }
+
+新版返回格式（/captcha/v2、/captcha/v2/base64）：
+  {
+    "slider": [x1, y1, x2, y2],       // 滑块坐标（空列表=未检测到）
+    "slider_confidence": 0.92,         // 滑块置信度
+    "gap": [x1, y1, x2, y2],          // 缺口坐标（空列表=未检测到）
+    "gap_confidence": 0.88,            // 缺口置信度
+    "offset": 162,                     // 滑动距离 = 缺口x1 - 滑块x1
+    "message": null                    // 错误信息
+  }
+"""
+
+import gc       # 垃圾回收，手动释放内存（HF免费层512MB）
+import os       # 读取环境变量（端口号）
+from typing import List, Optional
+
+import cv2      # OpenCV，图片解码
+import numpy as np  # 数组操作
+
+from captcha_recognizer.slider import Slider  # 核心：滑块识别模型
+from fastapi import FastAPI, File, UploadFile  # Web 框架
+from fastapi.middleware.cors import CORSMiddleware  # 跨域支持
+from pydantic import BaseModel  # 数据模型定义
+
+
+# ============================================================
+# FastAPI 应用初始化
+# ============================================================
 
 app = FastAPI(
-    docs_url=None,
-    redoc_url=None,
+    docs_url=None,   # 关闭 /docs（Swagger UI），避免暴露接口文档
+    redoc_url=None,  # 关闭 /redoc，同上
 )
 
-# 允许跨域
+# 允许跨域请求（任何域名都能调用这个 API）
+# 如果你只想允许特定域名调用，把 "*" 改成具体域名列表
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],
-    allow_methods=["*"],
-    allow_headers=["*"],
+    allow_origins=["*"],    # 允许的来源域名，* 表示全部
+    allow_methods=["*"],    # 允许的 HTTP 方法
+    allow_headers=["*"],    # 允许的请求头
 )
 
-# 启动时只加载一次模型
+
+# ============================================================
+# 模型加载（服务启动时只加载一次）
+# ============================================================
+# Slider() 会加载 slider.onnx 模型文件到内存
+# 放在全局意味着整个服务生命周期只加载一次，不用每次请求都读模型
 slider = Slider()
 
 
+# ============================================================
+# 健康检查接口
+# ============================================================
+
 @app.get("/")
 def hello_captcha():
+    """根路径，简单返回服务状态，可用于判断服务是否存活"""
     return {"Hello": "Captcha", "status": "running"}
 
 
 @app.get("/health")
 def health():
+    """健康检查接口，部署平台（如 HF Spaces）会定期调用"""
     return {"status": "ok"}
 
 
-class DetectionResult(BaseModel):
-    box: List[int]  # [x1, y1, x2, y2]
-    confidence: float
-    message: str = None
-
-
-@app.post("/captcha", response_model=DetectionResult)
-async def captcha(file: UploadFile = File(...)):
-    contents = await file.read()
+# ============================================================
+# 返回数据模型
+# ============================================================
 
-    nparr = np.frombuffer(contents, np.uint8)
+class DetectionResult(BaseModel):
+    """旧版接口返回格式（只返回缺口，保持向后兼容）"""
+    box: List[int]          # 缺口坐标 [x1, y1, x2, y2]
+    confidence: float       # 置信度 0~1
+    message: str = None     # 错误信息，正常为 null
+
+
+class SliderGapResult(BaseModel):
+    """新版接口返回格式（同时返回滑块和缺口）"""
+    slider: List[int] = []           # 滑块坐标 [x1, y1, x2, y2]，空列表=未检测到
+    slider_confidence: float = 0.0   # 滑块置信度
+    gap: List[int] = []              # 缺口坐标 [x1, y1, x2, y2]，空列表=未检测到
+    gap_confidence: float = 0.0      # 缺口置信度
+    offset: int = 0                  # 滑动距离 = 缺口x1 - 滑块x1，0=无法计算
+    message: Optional[str] = None    # 错误信息
+
+
+# ============================================================
+# 内部公共方法：图片解码 + 识别
+# ============================================================
+
+def _decode_and_recognize(image_bytes: bytes) -> SliderGapResult:
+    """
+    图片字节 → 识别结果（滑块+缺口）
+    
+    流程：字节 → OpenCV图片 → 模型推理 → 返回结构化结果
+    所有接口共用这个方法，避免重复代码
+    """
+    nparr = np.frombuffer(image_bytes, np.uint8)
     image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
 
     if image is None:
-        return DetectionResult(box=[], confidence=0, message="不支持的图片")
+        return SliderGapResult(message="不支持的图片")
 
-    box, confidence = slider.identify(source=image)
-    box = [int(x) for x in box]
+    # 调用模型，同时获取滑块和缺口
+    result = slider.identify_both(source=image)
 
-    del image, nparr, contents
+    # 手动释放内存（HF 免费层内存有限，防止 OOM）
+    del image, nparr
     gc.collect()
 
-    return DetectionResult(box=box, confidence=confidence)
+    return SliderGapResult(
+        slider=result['slider'],
+        slider_confidence=result['slider_confidence'],
+        gap=result['gap'],
+        gap_confidence=result['gap_confidence'],
+        offset=result['offset'],
+    )
+
+
+# ============================================================
+# 旧版接口（保持向后兼容，只返回缺口）
+# ============================================================
+
+@app.post("/captcha", response_model=DetectionResult)
+async def captcha(file: UploadFile = File(...)):
+    """
+    旧版：上传图片文件识别，只返回缺口
+    用法：POST /captcha，body 里传图片文件（form-data）
+    """
+    contents = await file.read()
+    result = _decode_and_recognize(contents)
+    if result.message:
+        return DetectionResult(box=[], confidence=0, message=result.message)
+    return DetectionResult(box=result.gap, confidence=result.gap_confidence)
 
 
 @app.post("/captcha/base64", response_model=DetectionResult)
 async def captcha_base64(data: dict):
-    """支持base64图片上传，用于懒人精灵调用"""
+    """
+    旧版：上传 base64 图片识别，只返回缺口
+    用法：POST /captcha/base64，JSON body: {"image": "base64字符串"}
+    """
     import base64
 
     image_b64 = data.get("image", "")
     if not image_b64:
         return DetectionResult(box=[], confidence=0, message="缺少image参数")
 
-    # 去掉base64头（如果有）
+    # 去掉 base64 头（如果有），如 "data:image/png;base64,"
     if "," in image_b64:
         image_b64 = image_b64.split(",")[1]
 
     try:
         image_bytes = base64.b64decode(image_b64)
-        nparr = np.frombuffer(image_bytes, np.uint8)
-        image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
     except Exception as e:
         return DetectionResult(box=[], confidence=0, message=f"base64解码失败: {str(e)}")
 
-    if image is None:
-        return DetectionResult(box=[], confidence=0, message="不支持的图片")
+    result = _decode_and_recognize(image_bytes)
+    if result.message:
+        return DetectionResult(box=[], confidence=0, message=result.message)
+    return DetectionResult(box=result.gap, confidence=result.gap_confidence)
+
+
+# ============================================================
+# 新版接口（返回滑块 + 缺口 + 偏移量）
+# ============================================================
+
+@app.post("/captcha/v2", response_model=SliderGapResult)
+async def captcha_v2(file: UploadFile = File(...)):
+    """
+    新版：上传图片文件识别，返回滑块+缺口+偏移量
+    用法：POST /captcha/v2，body 里传图片文件（form-data）
+    
+    返回示例：
+    {
+      "slider": [23, 78, 62, 117],
+      "slider_confidence": 0.92,
+      "gap": [185, 76, 224, 115],
+      "gap_confidence": 0.88,
+      "offset": 162,
+      "message": null
+    }
+    """
+    contents = await file.read()
+    return _decode_and_recognize(contents)
+
+
+@app.post("/captcha/v2/base64", response_model=SliderGapResult)
+async def captcha_v2_base64(data: dict):
+    """
+    新版：上传 base64 图片识别，返回滑块+缺口+偏移量
+    
+    请求格式：{"image": "base64字符串"}
+    适合懒人精灵等脚本通过 HTTP POST 调用
+    
+    base64 字符串支持两种格式：
+    - 带 data:image 前缀的：data:image/png;base64,iVBORw0KGgo...
+    - 纯 base64：iVBORw0KGgo...
+    """
+    import base64
+
+    image_b64 = data.get("image", "")
+    if not image_b64:
+        return SliderGapResult(message="缺少image参数")
 
-    box, confidence = slider.identify(source=image)
-    box = [int(x) for x in box]
+    if "," in image_b64:
+        image_b64 = image_b64.split(",")[1]
 
-    del image, nparr, image_bytes
-    gc.collect()
+    try:
+        image_bytes = base64.b64decode(image_b64)
+    except Exception as e:
+        return SliderGapResult(message=f"base64解码失败: {str(e)}")
 
-    return DetectionResult(box=box, confidence=confidence)
+    return _decode_and_recognize(image_bytes)
 
 
+# ============================================================
+# 启动入口
+# ============================================================
 if __name__ == "__main__":
     import uvicorn
+    # HF Spaces 要求从环境变量读端口号，默认 7860
+    # 本地开发时直接 python main.py 即可启动
     port = int(os.environ.get("PORT", 7860))
     uvicorn.run(app, host="0.0.0.0", port=port)