168 lines
5.9 KiB
Python
168 lines
5.9 KiB
Python
"""Phase 3 / Checkpoint 3 验收脚本：独立验证 Stage 3 风格还原链路。
|
||
|
||
复用 Phase 2 已抽出的 normal 帧，无需再跑 Trellis（省 $0.041）。
|
||
默认会挑最新一组 view_*_el0_normal.png 作为结构参考，并让用户指定原图。
|
||
|
||
用法（在 backend 目录激活 venv 后）：
|
||
# 1. 全 6 帧批量重绘（~6 × $0.07 ≈ $0.42）
|
||
python scripts/test_style_restore.py <原图路径>
|
||
|
||
# 2. 只跑一帧做快速冒烟（便宜 ~$0.07）
|
||
python scripts/test_style_restore.py <原图路径> --single
|
||
|
||
# 3. 手动指定 normal 帧前缀（不指定时用最新一组）
|
||
python scripts/test_style_restore.py <原图路径> --prefix view_xxxxxxxx
|
||
|
||
成功标准：
|
||
- 退出码 0
|
||
- 打印的每一帧都返回 success=True 且 url 以 /generated/restyled_ 开头
|
||
- 打开 generated/restyled_* 图片，目视检查：风格与原图一致、视角符合 az 标注
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import asyncio
|
||
import json
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
_BACKEND_ROOT = Path(__file__).resolve().parent.parent
|
||
if str(_BACKEND_ROOT) not in sys.path:
|
||
sys.path.insert(0, str(_BACKEND_ROOT))
|
||
|
||
from dotenv import load_dotenv # noqa: E402
|
||
|
||
load_dotenv(_BACKEND_ROOT / ".env", override=True)
|
||
|
||
from app.services import style_restorer # noqa: E402
|
||
from app.services.image_gen import GENERATED_DIR # noqa: E402
|
||
|
||
|
||
def _pick_latest_normal_group() -> tuple[str, list[dict]]:
|
||
"""在 generated/ 下挑一组最新的 view_<prefix>_az*_el0_normal.png，按 az 排好序。
|
||
|
||
返回 (prefix, frames)，frames = [{az, el, normal_png, color_png}, ...]
|
||
"""
|
||
all_normals = sorted(
|
||
GENERATED_DIR.glob("view_*_el0_normal.png"),
|
||
key=lambda p: p.stat().st_mtime,
|
||
reverse=True,
|
||
)
|
||
if not all_normals:
|
||
raise FileNotFoundError(
|
||
"generated/ 下找不到 view_*_el0_normal.png；请先跑 scripts/test_frame_extract.py"
|
||
)
|
||
# 最新文件的 prefix：view_<8hex>_az...
|
||
latest_name = all_normals[0].name
|
||
# 形如 view_63731b77_az0_el0_normal.png → prefix = view_63731b77
|
||
parts = latest_name.split("_")
|
||
prefix = "_".join(parts[:2])
|
||
return prefix, _load_group_by_prefix(prefix)
|
||
|
||
|
||
def _load_group_by_prefix(prefix: str) -> list[dict]:
|
||
normals = sorted(GENERATED_DIR.glob(f"{prefix}_az*_el0_normal.png"))
|
||
if not normals:
|
||
raise FileNotFoundError(f"找不到前缀为 {prefix} 的 normal 帧")
|
||
frames = []
|
||
for n in normals:
|
||
# 从文件名解析 az：view_63731b77_az120_el0_normal.png → az=120
|
||
tokens = n.stem.split("_")
|
||
az = 0
|
||
el = 0
|
||
for t in tokens:
|
||
if t.startswith("az"):
|
||
try:
|
||
az = int(t[2:])
|
||
except ValueError:
|
||
pass
|
||
elif t.startswith("el"):
|
||
try:
|
||
el = int(t[2:])
|
||
except ValueError:
|
||
pass
|
||
color = n.with_name(n.name.replace("_normal.png", "_color.png"))
|
||
frames.append({
|
||
"az": az,
|
||
"el": el,
|
||
"normal_png": f"/generated/{n.name}",
|
||
"color_png": f"/generated/{color.name}" if color.exists() else None,
|
||
"frame_idx": -1,
|
||
})
|
||
frames.sort(key=lambda f: f["az"])
|
||
return frames
|
||
|
||
|
||
async def main() -> int:
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument("original", help="用户原图路径（本地相对/绝对路径均可）")
|
||
parser.add_argument("--prefix", default=None, help="normal 帧的文件名前缀（如 view_63731b77）")
|
||
parser.add_argument("--single", action="store_true", help="只跑第一帧做快速冒烟")
|
||
parser.add_argument("--scale", type=float, default=None, help="覆盖 IP-Adapter 权重（0-1）")
|
||
parser.add_argument("--cn-scale", type=float, default=None, help="覆盖 ControlNet 权重（0-1）")
|
||
args = parser.parse_args()
|
||
|
||
original = args.original
|
||
original_path = Path(original)
|
||
if not original_path.is_absolute():
|
||
original_path = (_BACKEND_ROOT / original).resolve()
|
||
if not original_path.exists():
|
||
print(f"[test_style_restore] 原图不存在：{original_path}")
|
||
return 2
|
||
|
||
if args.prefix:
|
||
frames = _load_group_by_prefix(args.prefix)
|
||
prefix = args.prefix
|
||
else:
|
||
prefix, frames = _pick_latest_normal_group()
|
||
|
||
if args.single:
|
||
frames = frames[:1]
|
||
|
||
override_params = {}
|
||
if args.scale is not None:
|
||
override_params["scale"] = args.scale
|
||
if args.cn_scale is not None:
|
||
override_params["controlnet_conditioning_scale"] = args.cn_scale
|
||
|
||
print(f"[test_style_restore] 原图：{original_path}")
|
||
print(f"[test_style_restore] normal 帧组 prefix={prefix}，共 {len(frames)} 帧")
|
||
print(f"[test_style_restore] 角度列表：{[f['az'] for f in frames]}")
|
||
if override_params:
|
||
print(f"[test_style_restore] 覆盖参数：{override_params}")
|
||
print(
|
||
f"[test_style_restore] 预计耗时 ~{72}s × {len(frames)}，"
|
||
f"成本 ~${0.07 * len(frames):.2f}（并发 3 实际会更快）"
|
||
)
|
||
|
||
results = await style_restorer.restore_batch(
|
||
original_image=str(original_path),
|
||
frames=frames,
|
||
params=override_params or None,
|
||
concurrency=3,
|
||
)
|
||
|
||
print("[test_style_restore] 重绘结果：")
|
||
print(json.dumps(results, ensure_ascii=False, indent=2))
|
||
|
||
ok = sum(1 for r in results if r.get("success"))
|
||
print(f"[test_style_restore] 成功 {ok}/{len(results)} 帧")
|
||
|
||
if ok == 0:
|
||
print("[test_style_restore] 全军覆没，判定失败")
|
||
return 1
|
||
if ok < len(results):
|
||
print("[test_style_restore] 部分帧失败，视为部分通过；请检查上方错误信息")
|
||
return 0
|
||
|
||
print(
|
||
"[test_style_restore] 全部成功；请目视检查 generated/restyled_* 图片："
|
||
"风格应贴近原图、视角按 az 顺序环绕。"
|
||
)
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(asyncio.run(main()))
|