Files
hair/tests/test_api.py
T
xslandClaude Opus 4.8 94ad95850e feat(接口2): 新增生发后图片(ComfyUI/Flux inpaint)
在发际线预览基础上,每种发际线再出一张「植发3个月」生发图:

- hairline/mask.py: headmark 5步法遮罩(额头上部区域∩SegFormer头部=ROI,
  取发际线曲线以上闭合区域);用 hairline_texture_black 渲染黑线替代手绘检测;
  compose_comfy_rgba 合成 RGBA(alpha=255-mask, 透明=重绘区, 对齐 ComfyUI mask=1-alpha)
- hairline/comfyui.py: ComfyUI 客户端(默认8182),/upload/image+/prompt(改节点26+随机seed)
  +轮询/history+/view 取回生发图
- hairline/render.py: 抽出 build_overlay_layer 供遮罩取曲线像素
- hairline/service.py: extract_context 一次出 landmarks/parse_map/502点;
  generate_grow_results 每种=预览+生发图(同步串行N张,单张ComfyUI失败则grown置空不拖垮整请求)
- app.py: /hair/grow 返回 results[].grown_image_base64;重活放线程池避免卡事件循环
- add_hair.json 工作流 + hairline_texture_black/ 黑贴图入库
- 测试: test_mask.py(遮罩几何) + test_api mock ComfyUI 验 grown 字段,35 全绿

实测(5090): female 5张生发图同步约18s;预览/生发图人物五官服饰背景保持、黑线已清除。

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-06-14 23:42:30 +08:00

124 lines
4.2 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.
"""接口集成测试(FastAPI TestClient):错误码 + 鉴权 + 正常用例结构。"""
import base64
import pytest
from fastapi.testclient import TestClient
from conftest import fixture
import app as app_module
# 测试自带固定密码,避免依赖 worker_config.json 的实际值(中间件运行期读模块全局)
app_module.ACCEPT_PASSWORDS = ["testpass"]
URL = "/api/v1/face/measure"
H = {"X-Internal-Token": "testpass"}
@pytest.fixture(scope="module")
def client():
# with 触发 lifespan:加载模型单例(detector + 尽力加载 segmenter
with TestClient(app_module.app) as c:
yield c
def _post(client, fixture_name=None, headers=H, data=None, extra_files=None):
files = {}
if fixture_name:
files["image_file"] = (fixture_name, open(fixture(fixture_name), "rb"), "application/octet-stream")
if extra_files:
files.update(extra_files)
return client.post(URL, headers=headers, files=files or None, data=data)
def test_auth_missing_token_401(client):
r = _post(client, "frontal.jpg", headers={})
assert r.status_code == 401
def test_health_no_token_200(client):
r = client.get("/health")
assert r.status_code == 200
assert r.json()["status"] == "ok"
def test_param_none_provided_1007(client):
r = client.post(URL, headers=H)
assert r.json()["code"] == 1007
def test_param_multiple_provided_1007(client):
r = _post(client, "frontal.jpg", data={"image_url": "http://example.com/x.jpg"})
assert r.json()["code"] == 1007
def test_oversize_1006(client, oversize_file):
files = {"image_file": ("oversize.bin", open(oversize_file, "rb"), "application/octet-stream")}
r = client.post(URL, headers=H, files=files)
assert r.json()["code"] == 1006
def test_lowres_1002(client):
r = _post(client, "lowres.png")
assert r.json()["code"] == 1002
def test_no_face_1001(client):
r = _post(client, "landscape.jpg")
assert r.json()["code"] == 1001
def test_corrupt_1008(client):
r = _post(client, "corrupt.bin")
assert r.json()["code"] == 1008
GROW = "/api/v1/hair/grow"
def test_grow_missing_gender_1004(client):
files = {"image_file": ("frontal.jpg", open(fixture("frontal.jpg"), "rb"), "application/octet-stream")}
r = client.post(GROW, headers=H, files=files)
assert r.json()["code"] == 1004
_PNG_1x1 = base64.b64decode(
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
)
def test_grow_female_returns_5(client, monkeypatch):
# mock ComfyUI:不依赖 8182、不跑 Flux,只验证管线接线 + grown 字段
import hairline.comfyui as comfy
monkeypatch.setattr(comfy, "run", lambda *a, **k: _PNG_1x1)
files = {"image_file": ("frontal.jpg", open(fixture("frontal.jpg"), "rb"), "application/octet-stream")}
r = client.post(GROW, headers=H, files=files, data={"gender": "female"})
body = r.json()
assert body["code"] == 0, body
results = body["data"]["results"]
assert [x["hairline_type"] for x in results] == ["ellipse", "flower", "heart", "straight", "wave"]
assert [x["order"] for x in results] == [1, 2, 3, 4, 5]
assert base64.b64decode(results[0]["image_base64"])[:8] == b"\x89PNG\r\n\x1a\n"
assert base64.b64decode(results[0]["grown_image_base64"])[:8] == b"\x89PNG\r\n\x1a\n"
assert "image_url" not in results[0]
def test_success_structure(client):
r = _post(client, "frontal.jpg")
body = r.json()
assert body["code"] == 0, body
data = body["data"]
# 业务字段对齐文档
assert set(["face_total_height_cm", "four_courts", "seven_eyes",
"landmarks", "hairline_source", "head_pose",
"annotated_image_base64"]).issubset(data.keys())
assert set(["top_court_cm", "upper_court_cm", "middle_court_cm",
"lower_court_cm", "ratios"]).issubset(data["four_courts"].keys())
assert set(["eye_width_cm", "face_width_cm", "inter_eye_distance_cm",
"ratios"]).issubset(data["seven_eyes"].keys())
assert data["hairline_source"] in ("segmentation", "estimated")
# base64 解码为合法 PNG(非 URL
assert "annotated_image_url" not in data
png = base64.b64decode(data["annotated_image_base64"])
assert png[:8] == b"\x89PNG\r\n\x1a\n"