引入3d模型
This commit is contained in:
167
art-agent/backend/scripts/test_style_restore.py
Normal file
167
art-agent/backend/scripts/test_style_restore.py
Normal file
@@ -0,0 +1,167 @@
|
||||
"""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()))
|
||||
Reference in New Issue
Block a user