hair/scripts/compare_hairline_height.py

#!/usr/bin/env python3
"""保存接口52旧/新高度规则对比，供 static/test_interface52_compare.html 展示。

默认验证用户清单全部30张，男女各7型，共420组；只跑定位，不调用ComfyUI。
用法：./venv/bin/python scripts/compare_hairline_height.py
可用 --cases 1,2,3,4,5 只验证指定案例。
"""
from __future__ import annotations

import argparse
import base64
from datetime import datetime
import json
import os
from pathlib import Path
import time
from zoneinfo import ZoneInfo

import cv2
import httpx
import numpy as np

ROOT = Path(__file__).resolve().parent.parent


def run(args):
    manifest = json.loads((ROOT / "tests/fixtures/hairline_height_cases.json").read_text())
    available = {c["image_number"]: c for c in manifest["cases"]}
    requested = list(available) if args.cases == "all" else [int(n) for n in args.cases.split(",")]
    cases = [available[n] for n in requested]
    inputs = ROOT / "tests/output/hairline_height_compare/inputs"
    inputs.mkdir(parents=True, exist_ok=True)
    output = ROOT / "static/annotations/hairline_height_compare"
    output.mkdir(parents=True, exist_ok=True)
    report = {
        "created_at": datetime.now(ZoneInfo("Asia/Shanghai")).isoformat(),
        "case_numbers": requested, "groups": [], "errors": [],
    }
    started = time.monotonic()
    with httpx.Client(timeout=180, follow_redirects=True) as client:
        for case in cases:
            number = case["image_number"]
            source = inputs / f"case_{number}.png"
            if not source.exists():
                response = client.get(case["image_url"])
                response.raise_for_status()
                source.write_bytes(response.content)
            photo = cv2.imdecode(np.frombuffer(source.read_bytes(), np.uint8), cv2.IMREAD_COLOR)
            if photo is None:
                raise ValueError(f"案例 {number} 无法解码")
            # Keep the source photo available even when positioning fails.
            original_path = output / f"case_{number}_original.jpg"
            cv2.imwrite(str(original_path), photo, [cv2.IMWRITE_JPEG_QUALITY, 95])
            original_url = "/static/annotations/hairline_height_compare/" + original_path.name
            # Interface 52 frames native-resolution input before any size changes.
            content = cv2.imencode(".png", photo)[1].tobytes()
            for gender in args.genders.split(","):
                step = time.monotonic()
                try:
                    response = client.post(args.base_url.rstrip("/") + "/api/v1/debug/hairline-height-compare",
                        files={"image_file": (f"case_{number}.png", content, "image/png")},
                        data={"gender": gender, "hair_style": "1,2,3,4,5,6,7"})
                    response.raise_for_status()
                    body = response.json()
                    if body["code"] != 0:
                        raise ValueError(body["message"])
                    group = body["data"]
                    group.update(case_number=number, source_url=case["image_url"],
                                 elapsed_seconds=round(time.monotonic() - step, 2))
                    def save_image(obj, field, suffix):
                        value = obj.pop(field + "_base64", None)
                        if not value:
                            return
                        # Both gender sets use the same source and processed
                        # photos; share their URLs and browser image cache.
                        prefix = f"case_{number}" if obj is group else f"case_{number}_{gender}"
                        path = output / f"{prefix}_{suffix}"
                        path.write_bytes(base64.b64decode(value.split(",", 1)[-1]))
                        obj[field + "_url"] = "/static/annotations/hairline_height_compare/" + path.name
                    save_image(group, "input", "input.jpg")
                    save_image(group, "original", "original.jpg")
                    for result in group["results"]:
                        save_image(result, "old_overlay", f"{result['order']}_old.png")
                        save_image(result, "new_overlay", f"{result['order']}_new.png")
                        save_image(result, "grown_image", f"{result['order']}_grown.jpg")
                        if result["height_placement"]["target_error_px"] > 1:
                            raise ValueError(f"{result['hairline_type']} 中心误差超过1像素")
                    if group["framing"]["axis_x_error_px"] > 1:
                        raise ValueError("扶正后的9/152横向误差超过1像素")
                    report["groups"].append(group)
                    print(f"case {number} / {gender}: {len(group['results'])} shapes OK ({group['elapsed_seconds']}s)", flush=True)
                except Exception as ex:
                    report["errors"].append({"case_number": number, "gender": gender, "message": str(ex),
                                             "source_url": case["image_url"], "original_url": original_url,
                                             "width": photo.shape[1], "height": photo.shape[0]})
                    print(f"case {number} / {gender}: FAILED {ex}", flush=True)
                # Atomically expose partial results to the page while running.
                temporary = output / "latest.json.tmp"
                temporary.write_text(json.dumps(report, ensure_ascii=False, indent=2))
                os.replace(temporary, output / "latest.json")
    results = [r for g in report["groups"] for r in g["results"]]
    summary = {
        "cases": len(cases), "successful_groups": len(report["groups"]),
        "comparisons": len(results), "failures": len(report["errors"]),
        "max_center_error_px": max((r["height_placement"]["target_error_px"] for r in results), default=None),
        "elapsed_seconds": round(time.monotonic() - started, 2),
    }
    report["summary"] = summary
    temporary = output / "latest.json.tmp"
    temporary.write_text(json.dumps(report, ensure_ascii=False, indent=2))
    os.replace(temporary, output / "latest.json")
    print(json.dumps(summary, ensure_ascii=False), flush=True)
    if report["errors"]:
        raise SystemExit(1)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--cases", default="all", help="all 或逗号分隔的图片编号")
    parser.add_argument("--genders", default="female,male")
    parser.add_argument("--base-url", default="http://127.0.0.1:8187")
    run(parser.parse_args())