hair/scripts/batch_grow_v2.py

"""批量调用接口12（/api/v1/hairline/grow_v2）生成对比素材。

20 张女生照片 × 5 种发际线发型 × {高清, 非高清} = 200 张输出。
并发 4，失败的跳过并记录原因。结果图落盘到 static/report_hairline_v2/img/，
元数据落盘 static/report_hairline_v2/results.json，供生成报告用。

用法： python scripts/batch_grow_v2.py
"""
import base64
import json
import os
import time
from concurrent.futures import ThreadPoolExecutor, as_completed

import httpx

API = "http://127.0.0.1:8187/api/v1/hairline/grow_v2"
TOKEN = "dev-shared-secret-2026"
CONCURRENCY = 2

INPUT_DIR = "/home/xsl/hair/image/test"
OUT_DIR = "/home/xsl/hair/static/report_hairline_v2"
IMG_DIR = os.path.join(OUT_DIR, "img")
ORIG_DIR = os.path.join(OUT_DIR, "orig")

# 5 种发际线发型（= change_hair hair_id）
HAIRSTYLES = [
    ("chang_zhixian", "直线"),
    ("chang_tuoyuan", "椭圆"),
    ("chang_bolang", "波浪"),
    ("chang_xinxing", "心形"),
    ("chang_huaban", "花瓣"),
]
HR_OPTIONS = [(True, "hr"), (False, "nohr")]


def list_inputs():
    files = sorted(f for f in os.listdir(INPUT_DIR) if f.lower().endswith((".jpg", ".png")))
    return files


def one_call(stem, face_file, hair_id, hair_cn, is_hr, hr_tag):
    """调用一次接口，落盘结果图。返回结果 dict。"""
    src = os.path.join(INPUT_DIR, face_file)
    out_name = f"{stem}__{hair_id}__{hr_tag}.jpg"
    out_path = os.path.join(IMG_DIR, out_name)
    # 断点续跑：已存在的图直接跳过，不重复调用
    if os.path.exists(out_path) and os.path.getsize(out_path) > 1024:
        return {
            "stem": stem, "face_file": face_file, "hair_id": hair_id, "hair_cn": hair_cn,
            "is_hr": is_hr, "hr_tag": hr_tag, "ok": True,
            "out": f"img/{out_name}", "size": None,
            "ms": 0, "error": None, "skipped": True,
        }
    t0 = time.time()
    try:
        with open(src, "rb") as fh:
            files = {"image_file": (face_file, fh.read(), "image/jpeg")}
        data = {"hairline_id": hair_id, "is_hr": str(is_hr).lower()}
        with httpx.Client(timeout=180.0) as c:
            resp = c.post(API, headers={"X-Internal-Token": TOKEN}, files=files, data=data)
        j = resp.json()
        if j.get("code") != 0 or not j.get("data"):
            raise RuntimeError(f"code={j.get('code')} msg={j.get('message')}")
        b64 = j["data"]["final_base64"].split(",", 1)[1]
        raw = base64.b64decode(b64)
        with open(out_path, "wb") as fh:
            fh.write(raw)
        return {
            "stem": stem, "face_file": face_file, "hair_id": hair_id, "hair_cn": hair_cn,
            "is_hr": is_hr, "hr_tag": hr_tag, "ok": True,
            "out": f"img/{out_name}", "size": j["data"].get("image_size"),
            "ms": int((time.time() - t0) * 1000), "error": None,
        }
    except Exception as ex:  # noqa: BLE001
        return {
            "stem": stem, "face_file": face_file, "hair_id": hair_id, "hair_cn": hair_cn,
            "is_hr": is_hr, "hr_tag": hr_tag, "ok": False,
            "out": None, "size": None, "ms": int((time.time() - t0) * 1000),
            "error": str(ex)[:200],
        }


def main():
    os.makedirs(IMG_DIR, exist_ok=True)
    os.makedirs(ORIG_DIR, exist_ok=True)
    faces = list_inputs()
    print(f"输入 {len(faces)} 张脸 × {len(HAIRSTYLES)} 发型 × {len(HR_OPTIONS)} = "
          f"{len(faces)*len(HAIRSTYLES)*len(HR_OPTIONS)} 次调用，并发 {CONCURRENCY}")

    # 1. 先把原图拷一份到 orig/（报告要用）
    import shutil
    for f in faces:
        stem = os.path.splitext(f)[0]
        dst = os.path.join(ORIG_DIR, f"{stem}.jpg")
        if not os.path.exists(dst):
            shutil.copy2(os.path.join(INPUT_DIR, f), dst)

    # 2. 构造全部任务
    tasks = []
    for f in faces:
        stem = os.path.splitext(f)[0]
        for hair_id, hair_cn in HAIRSTYLES:
            for is_hr, hr_tag in HR_OPTIONS:
                tasks.append((stem, f, hair_id, hair_cn, is_hr, hr_tag))

    results = []
    done = 0
    total = len(tasks)
    t_start = time.time()
    with ThreadPoolExecutor(max_workers=CONCURRENCY) as ex:
        futs = {ex.submit(one_call, *t): t for t in tasks}
        for fut in as_completed(futs):
            r = fut.result()
            results.append(r)
            done += 1
            status = "OK " if r["ok"] else "FAIL"
            if r.get("skipped"):
                print(f"[{done}/{total}] SKIP {r['stem']} {r['hair_cn']} {r['hr_tag']}")
            elif r["ok"]:
                print(f"[{done}/{total}] {status} {r['stem']} {r['hair_cn']} {r['hr_tag']} "
                      f"({r['ms']}ms)", flush=True)
            else:
                print(f"[{done}/{total}] {status} {r['stem']} {r['hair_cn']} {r['hr_tag']} "
                      f"-> {r['error']}")

    elapsed = time.time() - t_start
    ok = sum(1 for r in results if r["ok"])
    fail = len(results) - ok
    # 按稳定顺序排序，报告好看
    order = {s: i for i, s in enumerate(HAIRSTYLES)}
    hr_order = {True: 0, False: 1}
    face_order = {os.path.splitext(f)[0]: i for i, f in enumerate(faces)}
    results.sort(key=lambda r: (face_order.get(r["stem"], 0),
                                order.get((r["hair_id"], r["hair_cn"]), 0),
                                hr_order.get(r["is_hr"], 0)))

    meta = {
        "total": total, "ok": ok, "fail": fail,
        "elapsed_sec": round(elapsed, 1), "concurrency": CONCURRENCY,
        "hairstyles": [{"id": h, "cn": c} for h, c in HAIRSTYLES],
        "hr_options": [{"is_hr": True, "tag": "hr"}, {"is_hr": False, "tag": "nohr"}],
        "faces": [os.path.splitext(f)[0] for f in faces],
        "generated_at": time.strftime("%Y-%m-%d %H:%M:%S"),
    }
    out = {"meta": meta, "results": results}
    with open(os.path.join(OUT_DIR, "results.json"), "w", encoding="utf-8") as fh:
        json.dump(out, fh, ensure_ascii=False, indent=2)

    print(f"\n完成：{ok}/{total} 成功，{fail} 失败，耗时 {elapsed:.1f}s")
    print(f"结果图 -> {IMG_DIR}")
    print(f"元数据 -> {OUT_DIR}/results.json")


if __name__ == "__main__":
    main()