Files
EPEEAIKit/art-agent/backend/scripts/test_trellis.py
2026-04-20 21:52:35 +08:00

73 lines
2.5 KiB
Python
Raw Permalink 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.
"""Phase 1 验收脚本:直接调用 Trellis验证 .glb + color_video + normal_video 能下到本地。
用法(在 backend 目录激活 venv 之后):
python scripts/test_trellis.py uploads/0e6fd7e47a524bb9ad18de0b72495e0e.png
不带参数时会尝试使用 uploads/ 下最新的 png 作为测试图。
成功标准:
- 退出码 0
- 终端打印 glb_path / color_video_path / normal_video_path 均为 /generated/ 开头
- 实际到 backend/generated/ 目录能看到 3 个文件且 size > 0
"""
import asyncio
import json
import sys
from pathlib import Path
# 允许直接 `python scripts/test_trellis.py` 跑
_BACKEND_ROOT = Path(__file__).resolve().parent.parent
if str(_BACKEND_ROOT) not in sys.path:
sys.path.insert(0, str(_BACKEND_ROOT))
# .env 加载(复用项目 main.py 的习惯)
from dotenv import load_dotenv # noqa: E402
load_dotenv(_BACKEND_ROOT / ".env", override=True)
from app.services import mesh_generator # noqa: E402
def _pick_default_image() -> str:
uploads = _BACKEND_ROOT / "uploads"
pngs = sorted(uploads.glob("*.png"), key=lambda p: p.stat().st_mtime, reverse=True)
if not pngs:
raise FileNotFoundError(f"uploads/ 下没有 png 文件,请显式传入图片路径")
return str(pngs[0].relative_to(_BACKEND_ROOT)).replace("\\", "/")
async def main() -> int:
image_arg = sys.argv[1] if len(sys.argv) > 1 else _pick_default_image()
# 规范化成 / 开头的 "本地相对 URL" 以便 to_data_uri 识别
if not image_arg.startswith("/") and not image_arg.startswith("data:") and not image_arg.startswith("http"):
image_arg = "/" + image_arg.lstrip("./").replace("\\", "/")
print(f"[test_trellis] 输入图:{image_arg}")
print("[test_trellis] 开始调用 Trellis预计 ~30s 起)...")
result = await mesh_generator.generate_with_trellis(image_arg)
print("[test_trellis] Trellis 返回:")
print(json.dumps(result, ensure_ascii=False, indent=2))
if not result.get("success"):
print("[test_trellis] 失败:", result.get("error"))
return 1
generated_dir = _BACKEND_ROOT / "generated"
for key in ("glb_path", "color_video_path", "normal_video_path"):
rel = result.get(key)
if not rel:
print(f"[test_trellis] 警告:缺少 {key}")
continue
local = generated_dir / Path(rel).name
size = local.stat().st_size if local.exists() else -1
print(f"[test_trellis] {key}: {rel} (本地大小: {size} bytes)")
return 0
if __name__ == "__main__":
sys.exit(asyncio.run(main()))