Files
hair/gateway/stub_worker.py
T
xsl 5b2b73f5d5 feat: 新增52号发际线接口并限制GPU worker来源IP
保留5号历史三档实现,新增中下庭高度定位、扶正裁切、矩形叠加与30案例对比页。

worker仅允许本机和配置白名单访问;网关不增加IP限制。空图片URL字段改写限定为52号,保持旧接口兼容。同步接口、裁切与部署文档。

验证:113项测试通过;1项历史数值基线失败,在未修改的a3fddbb版本复现。浏览器验证5/52及30案例对比页通过。
2026-11-04 10:20:29 +08:00

276 lines
11 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""Stub Worker — 模拟高性能后端,用于网关开发与测试。
每个 stub 监听一个端口,提供:
- GET /health → 200(可用 ?fail=1 模拟故障)
- POST /api/v1/face/measure → mock 响应(含 annotated_image_base64)
- POST /api/v1/hair/grow → mock 响应
- POST /api/v1/hair/grow-b → mock 响应
- POST /api/v1/face/features → mock 响应
- POST /api/v1/hairline/generate → mock 响应
- POST /api/v1/hairline/generate-52 → 新规则 mock 响应
鉴权:校验 X-Internal-Token 头,不匹配返回 401。
并发模拟:可通过 ?delay=N 让接口 sleep N 秒(默认 1)。
故障模拟:GET /health?fail=1 返回 503。
用法:
python gateway/stub_worker.py --port 28187
python gateway/stub_worker.py --port 28188
"""
import argparse
import base64
import json
import os
import sys
import time
import uuid
from typing import Optional
# 确保项目根在 sys.path
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import uvicorn
from fastapi import FastAPI, File, Form, Header, Request, UploadFile
from fastapi.responses import JSONResponse
# ---------------------------------------------------------------------------
# 一小张 PNG 用于 mock base64 返回
# 1x1 白色像素 PNG(最小合法 PNG)
# ---------------------------------------------------------------------------
TINY_PNG_BASE64 = (
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg=="
)
TINY_PNG_BYTES = base64.b64decode(TINY_PNG_BASE64)
# ---------------------------------------------------------------------------
# 应用工厂
# ---------------------------------------------------------------------------
def create_app(port: int, password: str) -> FastAPI:
app = FastAPI(
title=f"Stub Worker :{port}",
version="0.1.0",
docs_url=None,
redoc_url=None,
)
# --- 鉴权依赖 ---
def check_token(x_internal_token: Optional[str] = Header(default=None)):
if x_internal_token != password:
return False
return True
# --- 通用 mock data 生成 ---
def mock_ok(data: dict, delay: float = 1.0):
"""模拟处理延迟后返回标准响应。"""
if delay > 0:
time.sleep(delay)
return {
"code": 0,
"message": "success",
"request_id": f"stub-{port}-{uuid.uuid4().hex[:8]}",
"data": data,
}
def mock_err(code: int, message: str):
return {
"code": code,
"message": message,
"request_id": f"stub-{port}-{uuid.uuid4().hex[:8]}",
"data": None,
}
# --- /health ---
@app.get("/health", include_in_schema=False)
async def health(request: Request, x_internal_token: Optional[str] = Header(default=None)):
# /health 也校验 token(生产环境 worker 行为)
if x_internal_token != password:
return JSONResponse(status_code=401, content={"detail": "unauthorized"})
fail = request.query_params.get("fail")
if fail == "1":
return JSONResponse(status_code=503, content={"status": "error"})
return {"status": "ok", "worker": f"stub-{port}"}
# --- 接口1 ---
@app.post("/api/v1/face/measure")
async def face_measure(
request: Request,
x_internal_token: Optional[str] = Header(default=None),
image_file: Optional[UploadFile] = File(default=None),
image_url: Optional[str] = Form(default=None),
image_base64: Optional[str] = Form(default=None),
):
if x_internal_token != password:
return JSONResponse(status_code=401, content={"detail": "unauthorized"})
delay = float(request.query_params.get("delay", "1"))
data = {
"annotated_image_base64": f"data:image/png;base64,{TINY_PNG_BASE64}",
"face_total_height_cm": 13.76,
"four_courts": {
"top_court_cm": 3.44,
"upper_court_cm": 3.44,
"middle_court_cm": 3.44,
"lower_court_cm": 3.44,
"ratios": {
"top_court": 0.25,
"upper_court": 0.25,
"middle_court": 0.25,
"lower_court": 0.25,
},
},
"seven_eyes": {
"eye_width_cm": 3.44,
"face_width_cm": 24.08,
"inter_eye_distance_cm": 3.44,
"ratios": {"eye_width": 0.143, "inter_eye_distance": 0.143},
},
"landmarks": {
"hair_top": {"x": 540, "y": 120},
"hairline": {"x": 540, "y": 430},
"brow_center": {"x": 540, "y": 740},
"nose_bottom": {"x": 540, "y": 1050},
"chin_tip": {"x": 540, "y": 1360},
},
}
return mock_ok(data, delay=delay)
# --- 接口2 ---
@app.post("/api/v1/hair/grow")
async def hair_grow(
request: Request,
x_internal_token: Optional[str] = Header(default=None),
image_file: Optional[UploadFile] = File(default=None),
image_url: Optional[str] = Form(default=None),
image_base64: Optional[str] = Form(default=None),
beauty_enabled: bool = Form(default=False),
):
if x_internal_token != password:
return JSONResponse(status_code=401, content={"detail": "unauthorized"})
delay = float(request.query_params.get("delay", "1"))
data = {
"results": [
{
"image_base64": f"data:image/png;base64,{TINY_PNG_BASE64}",
"hairline_type": "花瓣形",
"order": 1,
},
{
"image_base64": f"data:image/png;base64,{TINY_PNG_BASE64}",
"hairline_type": "波浪形",
"order": 2,
},
]
}
return mock_ok(data, delay=delay)
# --- 接口3 ---
@app.post("/api/v1/hair/grow-b")
async def hair_grow_b(
request: Request,
x_internal_token: Optional[str] = Header(default=None),
marked_image_file: Optional[UploadFile] = File(default=None),
marked_image_url: Optional[str] = Form(default=None),
marked_image_base64: Optional[str] = Form(default=None),
original_image_file: Optional[UploadFile] = File(default=None),
original_image_url: Optional[str] = Form(default=None),
original_image_base64: Optional[str] = Form(default=None),
):
if x_internal_token != password:
return JSONResponse(status_code=401, content={"detail": "unauthorized"})
delay = float(request.query_params.get("delay", "1"))
data = {
"best_hairline_image_base64": f"data:image/png;base64,{TINY_PNG_BASE64}",
"hair_growth_image_base64": f"data:image/png;base64,{TINY_PNG_BASE64}",
"hairline_type": "花瓣形",
}
return mock_ok(data, delay=delay)
# --- 接口4 ---
@app.post("/api/v1/face/features")
async def face_features(
request: Request,
x_internal_token: Optional[str] = Header(default=None),
image_file: Optional[UploadFile] = File(default=None),
image_url: Optional[str] = Form(default=None),
image_base64: Optional[str] = Form(default=None),
):
if x_internal_token != password:
return JSONResponse(status_code=401, content={"detail": "unauthorized"})
delay = float(request.query_params.get("delay", "1"))
features = json.dumps(
{
"face_shape": "鹅蛋脸",
"eyebrow_shape": "柳叶眉",
"facial_age": 26,
"dynamic_static_type": "静态",
"gender": "女",
"gene_style": {"label": "面部特征标签", "style": "基因风格示例"},
},
ensure_ascii=False,
)
return mock_ok({"features": features}, delay=delay)
# --- 接口5 ---
@app.post("/api/v1/hairline/generate")
@app.post("/api/v1/hairline/generate-52")
async def hairline_generate(
request: Request,
x_internal_token: Optional[str] = Header(default=None),
image_file: Optional[UploadFile] = File(default=None),
image_url: Optional[str] = Form(default=None),
image_base64: Optional[str] = Form(default=None),
):
if x_internal_token != password:
return JSONResponse(status_code=401, content={"detail": "unauthorized"})
delay = float(request.query_params.get("delay", "1"))
_png = f"data:image/png;base64,{TINY_PNG_BASE64}"
data = {
"hairline_images": [
{
"hairline_type": "ellipse",
"image_middle_base64": _png,
"image_high_base64": _png,
"image_low_base64": _png,
"grown_image_base64": _png,
"order": 1,
},
{
"hairline_type": "heart",
"image_middle_base64": _png,
"image_high_base64": _png,
"image_low_base64": _png,
"grown_image_base64": None,
"order": 3,
},
],
"best_hairline_center_point": {"x": 540, "y": 430},
}
if request.url.path.endswith("generate-52"):
for item in data["hairline_images"]:
item["image_high_base64"] = None
item["image_low_base64"] = None
data.update(high_hairline_center_point=None, low_hairline_center_point=None,
face_measure=None, framing=None, height_rule=None,
guide_landmarks={}, input_image_base64=_png, original_image_base64=_png,
coordinate_space="original_image", width=1, height=1)
return mock_ok(data, delay=delay)
return app
# ---------------------------------------------------------------------------
# 入口
# ---------------------------------------------------------------------------
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Stub Worker for gateway testing")
parser.add_argument("--port", type=int, default=28187, help="监听端口(默认 28187)")
parser.add_argument("--password", type=str, default="dev-shared-secret-2026", help="共享密码")
args = parser.parse_args()
app = create_app(port=args.port, password=args.password)
print(f"[stub_worker] 启动在 :{args.port},密码={args.password[:4]}...")
uvicorn.run(app, host="127.0.0.1", port=args.port, log_level="warning")