引入3d模型
This commit is contained in:
@@ -39,17 +39,34 @@ TOOL_DEFINITIONS = [
|
||||
"function": {
|
||||
"name": "transform_view",
|
||||
"description": (
|
||||
"将一张图片转换为多个不同视角。"
|
||||
"使用 Zero123++ 模型从单张图片生成 6 个固定视角的图片。"
|
||||
"适用于建筑、物体等需要从不同角度查看的场景。"
|
||||
"输入图片必须是正方形（或会被自动裁切为正方形），建议分辨率 >= 320x320。"
|
||||
"输出 6 个视角：方位角 30°/90°/150°/210°/270°/330°，"
|
||||
"仰角交替为 30°/-20°（正俯视/微仰视）。"
|
||||
"此工具需要参考图作为输入——必须先有用户上传的图片才能使用。"
|
||||
"将一张图片转换为多个不同视角。需要用户先上传参考图。\n"
|
||||
"可选模型：\n"
|
||||
"- zero123plus（默认）：~15s/次，Zero123++ 直接输出 6 个固定视角"
|
||||
"（方位角 30/90/150/210/270/330°，仰角交替 30/-20°），适合快速预览，"
|
||||
"对建筑大角度可能有错位。\n"
|
||||
"- trellis：~105s/次，先 3D 重建再重绘，"
|
||||
"支持任意方位角 + 原画风保留，适合卡通描边风格建筑。"
|
||||
),
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {},
|
||||
"properties": {
|
||||
"model_id": {
|
||||
"type": "string",
|
||||
"enum": ["zero123plus", "trellis"],
|
||||
"default": "zero123plus",
|
||||
"description": "视角变换模型；默认 zero123plus（快），trellis 画风还原更强但慢",
|
||||
},
|
||||
"azimuths": {
|
||||
"type": "array",
|
||||
"items": {"type": "integer"},
|
||||
"description": "mesh 管道专用：自定义方位角数组（0-359），默认 [0,60,120,180,240,300]",
|
||||
},
|
||||
"preserve_style": {
|
||||
"type": "boolean",
|
||||
"default": True,
|
||||
"description": "mesh 管道专用：是否对抽帧结果走 IP-Adapter+ControlNet 重绘还原原画风",
|
||||
},
|
||||
},
|
||||
"required": [],
|
||||
},
|
||||
},
|
||||
@@ -92,7 +109,13 @@ async def execute_tool(
|
||||
"success": False,
|
||||
"error": "视角变换需要一张输入图片，请先上传参考图",
|
||||
}
|
||||
result = await transform_view(ref_image_urls[0])
|
||||
result = await transform_view(
|
||||
ref_image_urls[0],
|
||||
model_id=arguments.get("model_id"),
|
||||
azimuths=arguments.get("azimuths"),
|
||||
elevations=arguments.get("elevations"),
|
||||
preserve_style=arguments.get("preserve_style", True),
|
||||
)
|
||||
return result
|
||||
|
||||
return {"success": False, "error": f"未知工具: {tool_name}"}
|
||||
|
||||
@@ -225,9 +225,35 @@ IMAGE_MODELS: dict[str, dict[str, Any]] = {
|
||||
"output_format": "png",
|
||||
},
|
||||
},
|
||||
# Mesh 管道专用的内部模型（internal=True 不展示到前端下拉）：
|
||||
# 输入两张图：image=原图做 IP-Adapter 风格参考；controlnet_input=normal/depth 结构参考。
|
||||
# 用途：在 Stage 3 对 Trellis 抽出来的 normal 帧做二次重绘，还原原画风。
|
||||
"ip-adapter-controlnet-depth": {
|
||||
"id": "ip-adapter-controlnet-depth",
|
||||
"name": "IP-Adapter + ControlNet Depth (内部)",
|
||||
"provider": "replicate",
|
||||
"model_id": "chigozienri/ip_adapter-sdxl-controlnet-depth:0436c8702ef52616be5c30948551b3af6a86c821cca9b01f11ac297624fff14c",
|
||||
"description": "Mesh 管道 Stage 3 专用：IP-Adapter 保风格 + ControlNet 保几何",
|
||||
"internal": True,
|
||||
"supports_ref_image": True,
|
||||
"ref_image_param": "image", # IP-Adapter 参考图字段
|
||||
"default_params": {
|
||||
"scale": 0.75, # IP-Adapter 权重（建议 0.7-0.8）
|
||||
"controlnet_conditioning_scale": 0.8, # 结构约束（建议 0.7-0.9）
|
||||
"prompt": "same building, same art style, consistent with reference",
|
||||
"negative_prompt": "blurry, distorted, different style, realistic photo, photograph",
|
||||
"num_outputs": 1,
|
||||
"num_inference_steps": 30,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def get_style_restore_model_id() -> str:
|
||||
"""Mesh 管道 Stage 3 使用的内部模型短 ID。"""
|
||||
return "ip-adapter-controlnet-depth"
|
||||
|
||||
|
||||
# ─── 视角变换模型注册表 ─────────────────────────────────
|
||||
#
|
||||
# 视角变换模型不用于生图，而是将已有图片转换为不同视角。
|
||||
@@ -240,6 +266,7 @@ VIEW_TRANSFORM_MODELS: dict[str, dict[str, Any]] = {
|
||||
"provider": "replicate",
|
||||
"model_id": "jd7h/zero123plusplus:c69c6559a29011b576f1ff0371b3bc1add2856480c60520c7e9ce0b40a6e9052",
|
||||
"description": "单图生成 6 个固定视角，适合建筑/物体的多角度预览",
|
||||
"pipeline": "grid",
|
||||
"output_views": [
|
||||
{"azimuth": 30, "elevation": 30},
|
||||
{"azimuth": 90, "elevation": -20},
|
||||
@@ -248,7 +275,46 @@ VIEW_TRANSFORM_MODELS: dict[str, dict[str, Any]] = {
|
||||
{"azimuth": 270, "elevation": 30},
|
||||
{"azimuth": 330, "elevation": -20},
|
||||
],
|
||||
"grid_layout": {"cols": 3, "rows": 2},
|
||||
"grid_layout": {"cols": 2, "rows": 3},
|
||||
"enabled": True,
|
||||
},
|
||||
"trellis": {
|
||||
"id": "trellis",
|
||||
"name": "Trellis (Mesh Pipeline)",
|
||||
"provider": "replicate",
|
||||
"model_id": "firtoz/trellis:e8f6c45206993f297372f5436b90350817bd9b4a0d52d2a76df50c1c8afa2b3c",
|
||||
"description": "3D 重建 + 任意视角 + 风格还原，建筑友好（~30s 重建 + 72s/张重绘）",
|
||||
"pipeline": "mesh",
|
||||
"default_params": {
|
||||
"texture_size": 1024,
|
||||
"mesh_simplify": 0.95,
|
||||
"generate_color": True,
|
||||
"generate_normal": True,
|
||||
"generate_model": True,
|
||||
"save_gaussian_ply": False,
|
||||
"ss_sampling_steps": 12,
|
||||
"slat_sampling_steps": 12,
|
||||
"ss_guidance_strength": 7.5,
|
||||
"slat_guidance_strength": 3.0,
|
||||
},
|
||||
"default_azimuths": [0, 60, 120, 180, 240, 300],
|
||||
"default_elevations": [0, 0, 0, 0, 0, 0],
|
||||
"enabled": True,
|
||||
},
|
||||
"hunyuan3d": {
|
||||
"id": "hunyuan3d",
|
||||
"name": "Hunyuan3D-2 (高质量)",
|
||||
"provider": "replicate",
|
||||
"model_id": "tencent/hunyuan3d-2:b1b9449a1277e10402781c5d41eb30c0a0683504fb23fab591ca9dfc2aabe1cb",
|
||||
"description": "几何质量最佳（~127s），需自定义渲染器（二期启用）",
|
||||
"pipeline": "mesh",
|
||||
"default_params": {
|
||||
"steps": 50,
|
||||
"guidance_scale": 5.5,
|
||||
"octree_resolution": 256,
|
||||
"remove_background": True,
|
||||
},
|
||||
"enabled": False,
|
||||
},
|
||||
}
|
||||
|
||||
@@ -261,14 +327,16 @@ def get_view_transform_model_config(model_id: str | None = None) -> dict[str, An
|
||||
|
||||
|
||||
def get_view_transform_models_list() -> list[dict]:
|
||||
"""返回前端下拉列表所需的视角变换模型摘要信息。"""
|
||||
"""返回前端下拉列表所需的视角变换模型摘要信息（只包含 enabled=True 的）。"""
|
||||
return [
|
||||
{
|
||||
"id": cfg["id"],
|
||||
"name": cfg["name"],
|
||||
"description": cfg["description"],
|
||||
"pipeline": cfg.get("pipeline", "grid"),
|
||||
}
|
||||
for cfg in VIEW_TRANSFORM_MODELS.values()
|
||||
if cfg.get("enabled", True)
|
||||
]
|
||||
|
||||
|
||||
@@ -297,7 +365,7 @@ def get_ref_image_model_id() -> str | None:
|
||||
|
||||
|
||||
def get_image_models_list() -> list[dict]:
|
||||
"""返回前端下拉列表所需的模型摘要信息。"""
|
||||
"""返回前端下拉列表所需的模型摘要信息（过滤掉 internal=True 的内部模型）。"""
|
||||
return [
|
||||
{
|
||||
"id": cfg["id"],
|
||||
@@ -306,6 +374,7 @@ def get_image_models_list() -> list[dict]:
|
||||
"supports_ref_image": cfg.get("supports_ref_image", False),
|
||||
}
|
||||
for cfg in IMAGE_MODELS.values()
|
||||
if not cfg.get("internal", False)
|
||||
]
|
||||
|
||||
|
||||
|
||||
189
art-agent/backend/app/services/mesh_generator.py
Normal file
189
art-agent/backend/app/services/mesh_generator.py
Normal file
@@ -0,0 +1,189 @@
|
||||
"""3D Mesh 生成服务 — 阶段 1：单图 → textured .glb + 360° 预渲染视频。
|
||||
|
||||
当前支持：
|
||||
- Trellis (firtoz/trellis)：~30s，$0.041/次，自带 color_video + normal_video + .glb
|
||||
- Hunyuan3D-2 (tencent/hunyuan3d-2)：~127s，$0.12/次，仅 .glb 无预渲染视频（二期启用）
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from app.config import get_view_transform_model_config
|
||||
from app.services.image_gen import _make_replicate_client, GENERATED_DIR, to_data_uri
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_replicate_client = _make_replicate_client()
|
||||
|
||||
SUPPORTED_EXTS = {
|
||||
"model_file": ".glb",
|
||||
"color_video": ".mp4",
|
||||
"normal_video": ".mp4",
|
||||
"combined_video": ".mp4",
|
||||
"gaussian_ply": ".ply",
|
||||
"mesh": ".glb",
|
||||
}
|
||||
|
||||
|
||||
async def _download_asset(url: str, ext: str, prefix: str) -> Path:
|
||||
"""下载远程资源到 generated/ 目录，返回本地路径。"""
|
||||
filename = f"{prefix}_{uuid.uuid4().hex[:12]}{ext}"
|
||||
filepath = GENERATED_DIR / filename
|
||||
proxy = os.environ.get("HTTPS_PROXY") or os.environ.get("HTTP_PROXY")
|
||||
async with httpx.AsyncClient(proxy=proxy, timeout=httpx.Timeout(300.0)) as client:
|
||||
resp = await client.get(url, follow_redirects=True)
|
||||
resp.raise_for_status()
|
||||
filepath.write_bytes(resp.content)
|
||||
return filepath
|
||||
|
||||
|
||||
def _to_url(value: Any) -> str | None:
|
||||
"""把 Replicate FileOutput / str / None 统一转为 URL 字符串。"""
|
||||
if value is None:
|
||||
return None
|
||||
s = str(value)
|
||||
if s.startswith("http://") or s.startswith("https://"):
|
||||
return s
|
||||
return None
|
||||
|
||||
|
||||
async def generate_with_trellis(
|
||||
image_path: str,
|
||||
params: dict[str, Any] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""调用 Trellis 做单图 3D 重建，下载所有输出资产到本地。
|
||||
|
||||
Args:
|
||||
image_path: 输入图（本地路径 / data URI / URL）
|
||||
params: 覆盖 default_params 的可选参数字典
|
||||
|
||||
Returns:
|
||||
{
|
||||
"success": bool,
|
||||
"glb_path": "/generated/trellis_xxx.glb" | None,
|
||||
"color_video_path": "/generated/trellis_xxx_color.mp4" | None,
|
||||
"normal_video_path": "/generated/trellis_xxx_normal.mp4" | None,
|
||||
"combined_video_path": ...,
|
||||
"raw_output_keys": [...], # 调试用
|
||||
"error": str | None,
|
||||
}
|
||||
"""
|
||||
config = get_view_transform_model_config("trellis")
|
||||
replicate_model_id = config["model_id"]
|
||||
default_params = dict(config.get("default_params", {}))
|
||||
if params:
|
||||
default_params.update(params)
|
||||
|
||||
GENERATED_DIR.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
try:
|
||||
image_uri = to_data_uri(image_path)
|
||||
# firtoz/trellis:e8f6c452... 这个部署版的 schema：
|
||||
# - 输入必须用 `images` 数组（传 `image` 单数会 422）
|
||||
# - 输出字段 `color_video` 和 `normal_video` 实测恒为 null
|
||||
# - 实际可用的视频只有 `combined_video`（color+normal 左右并排）
|
||||
# 因此 Phase 2 抽帧时从 combined_video 切左半得 color、切右半得 normal。
|
||||
input_params: dict[str, Any] = {
|
||||
"images": [image_uri],
|
||||
"randomize_seed": False,
|
||||
"seed": 0,
|
||||
**default_params,
|
||||
}
|
||||
|
||||
logger.info(
|
||||
"调用 Trellis 进行 3D 重建：model=%s, params=%s",
|
||||
replicate_model_id,
|
||||
{k: v for k, v in input_params.items() if k != "images"},
|
||||
)
|
||||
output = await _replicate_client.async_run(
|
||||
replicate_model_id, input=input_params, wait=False
|
||||
)
|
||||
|
||||
if not isinstance(output, dict):
|
||||
return {
|
||||
"success": False,
|
||||
"error": f"Trellis 返回非字典输出：{type(output).__name__} / {str(output)[:200]}",
|
||||
"raw_output_keys": [],
|
||||
}
|
||||
|
||||
raw_keys = list(output.keys())
|
||||
logger.info("Trellis 返回字段：%s", raw_keys)
|
||||
|
||||
# 按已知字段挨个下载
|
||||
result: dict[str, Any] = {
|
||||
"success": True,
|
||||
"raw_output_keys": raw_keys,
|
||||
"error": None,
|
||||
}
|
||||
download_map = [
|
||||
("model_file", "glb_path", ".glb", "trellis"),
|
||||
("color_video", "color_video_path", ".mp4", "trellis_color"),
|
||||
("normal_video", "normal_video_path", ".mp4", "trellis_normal"),
|
||||
("combined_video", "combined_video_path", ".mp4", "trellis_combined"),
|
||||
]
|
||||
|
||||
for src_key, dst_key, ext, prefix in download_map:
|
||||
if src_key not in output:
|
||||
result[dst_key] = None
|
||||
continue
|
||||
url = _to_url(output[src_key])
|
||||
if not url:
|
||||
logger.warning("Trellis 字段 %s 不是有效 URL：%s", src_key, output[src_key])
|
||||
result[dst_key] = None
|
||||
continue
|
||||
try:
|
||||
local_path = await _download_asset(url, ext, prefix)
|
||||
result[dst_key] = f"/generated/{local_path.name}"
|
||||
logger.info("已下载 %s → %s", src_key, local_path.name)
|
||||
except Exception as e:
|
||||
logger.error("下载 %s 失败：%s", src_key, e)
|
||||
result[dst_key] = None
|
||||
|
||||
# 阶段 2 抽帧的优先级：color_video > combined_video（切左半）
|
||||
# 只要其中一个可用即算成功
|
||||
has_color_source = bool(result.get("color_video_path")) or bool(
|
||||
result.get("combined_video_path")
|
||||
)
|
||||
if not has_color_source:
|
||||
result["success"] = False
|
||||
result["error"] = (
|
||||
"Trellis 未返回 color_video 或 combined_video，无法进入阶段 2 抽帧；"
|
||||
f"实际返回字段：{raw_keys}"
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
detail = str(e) or f"{type(e).__name__}: {repr(e)}"
|
||||
logger.error("Trellis 调用失败：%s", detail, exc_info=True)
|
||||
return {
|
||||
"success": False,
|
||||
"error": detail,
|
||||
"raw_output_keys": [],
|
||||
}
|
||||
|
||||
|
||||
async def generate_mesh(
|
||||
image_path: str,
|
||||
model_id: str = "trellis",
|
||||
params: dict[str, Any] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""统一入口：按 model_id 分发到具体实现。"""
|
||||
if model_id == "trellis":
|
||||
return await generate_with_trellis(image_path, params)
|
||||
if model_id == "hunyuan3d":
|
||||
return {
|
||||
"success": False,
|
||||
"error": "Hunyuan3D-2 (一期未启用，待二期接入自定义渲染)",
|
||||
"raw_output_keys": [],
|
||||
}
|
||||
return {
|
||||
"success": False,
|
||||
"error": f"未知 mesh 模型：{model_id}",
|
||||
"raw_output_keys": [],
|
||||
}
|
||||
180
art-agent/backend/app/services/style_restorer.py
Normal file
180
art-agent/backend/app/services/style_restorer.py
Normal file
@@ -0,0 +1,180 @@
|
||||
"""阶段 3：对 Trellis 抽帧结果做 IP-Adapter + ControlNet 重绘，还原原画风。
|
||||
|
||||
输入：
|
||||
- original_image：用户上传的原图（提供风格 / IP-Adapter 参考）
|
||||
- structure_image：Stage 2 抽出的 normal 帧（提供几何 / ControlNet 结构约束）
|
||||
|
||||
输出：
|
||||
- 单张保留原画风且视角正确的 PNG，落盘到 GENERATED_DIR/ 下
|
||||
|
||||
模型：chigozienri/ip_adapter-sdxl-controlnet-depth（配置见 config.IMAGE_MODELS）
|
||||
- $0.07/张，~72s，硬件 L40S
|
||||
- `image` 字段 = IP-Adapter 参考图（风格源）
|
||||
- `controlnet_input` 字段 = 结构参考图（当前一期传 normal 帧作为 depth 近似）
|
||||
|
||||
调参建议（可在 IMAGE_MODELS 默认值基础上覆盖）：
|
||||
- scale (IP-Adapter 权重): 0.7-0.8，越高越贴近原图风格
|
||||
- controlnet_conditioning_scale (结构权重): 0.7-0.9，越高越保几何
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from app.config import get_image_model_config, get_style_restore_model_id
|
||||
from app.services.image_gen import GENERATED_DIR, _make_replicate_client, to_data_uri
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_replicate_client = _make_replicate_client()
|
||||
|
||||
|
||||
async def _download_image(url: str, prefix: str) -> Path:
|
||||
"""把 Replicate 返回的图片下载到 generated/。"""
|
||||
filename = f"{prefix}_{uuid.uuid4().hex[:12]}.png"
|
||||
filepath = GENERATED_DIR / filename
|
||||
proxy = os.environ.get("HTTPS_PROXY") or os.environ.get("HTTP_PROXY")
|
||||
async with httpx.AsyncClient(proxy=proxy, timeout=httpx.Timeout(180.0)) as client:
|
||||
resp = await client.get(url, follow_redirects=True)
|
||||
resp.raise_for_status()
|
||||
filepath.write_bytes(resp.content)
|
||||
return filepath
|
||||
|
||||
|
||||
def _first_url(output: Any) -> str | None:
|
||||
"""Replicate 输出有时是 list[FileOutput]、有时是单个；统一取第一张。"""
|
||||
if output is None:
|
||||
return None
|
||||
if isinstance(output, list):
|
||||
if not output:
|
||||
return None
|
||||
return str(output[0])
|
||||
return str(output)
|
||||
|
||||
|
||||
async def restore(
|
||||
original_image: str,
|
||||
structure_image: str,
|
||||
azimuth: int = 0,
|
||||
elevation: int = 0,
|
||||
params: dict[str, Any] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""对单张 normal 帧做 IP-Adapter + ControlNet 重绘。
|
||||
|
||||
Args:
|
||||
original_image: 用户原图，本地路径 / data URI / URL 均可
|
||||
structure_image: Stage 2 抽出的 normal 帧，同上
|
||||
azimuth / elevation: 仅用于日志与返回元数据，不参与模型调用
|
||||
params: 覆盖默认参数的字典（scale / controlnet_conditioning_scale / prompt 等）
|
||||
|
||||
Returns:
|
||||
{
|
||||
"success": bool,
|
||||
"url": "/generated/restyled_xxx.png" | None,
|
||||
"azimuth": int,
|
||||
"elevation": int,
|
||||
"error": str | None,
|
||||
}
|
||||
"""
|
||||
config = get_image_model_config(get_style_restore_model_id())
|
||||
replicate_model_id = config["model_id"]
|
||||
merged_params = dict(config.get("default_params", {}))
|
||||
if params:
|
||||
merged_params.update(params)
|
||||
|
||||
GENERATED_DIR.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
try:
|
||||
image_uri = to_data_uri(original_image)
|
||||
structure_uri = to_data_uri(structure_image)
|
||||
|
||||
input_params: dict[str, Any] = {
|
||||
"image": image_uri, # IP-Adapter 参考图 = 用户原图
|
||||
"controlnet_input": structure_uri, # 结构参考 = normal 帧
|
||||
**merged_params,
|
||||
}
|
||||
|
||||
logger.info(
|
||||
"Stage 3 重绘 az=%d el=%d: scale=%s, cn_scale=%s",
|
||||
azimuth, elevation,
|
||||
merged_params.get("scale"),
|
||||
merged_params.get("controlnet_conditioning_scale"),
|
||||
)
|
||||
|
||||
output = await _replicate_client.async_run(
|
||||
replicate_model_id, input=input_params, wait=False
|
||||
)
|
||||
|
||||
url = _first_url(output)
|
||||
if not url or not (url.startswith("http://") or url.startswith("https://")):
|
||||
return {
|
||||
"success": False,
|
||||
"url": None,
|
||||
"azimuth": azimuth,
|
||||
"elevation": elevation,
|
||||
"error": f"模型未返回有效图片 URL: {str(output)[:200]}",
|
||||
}
|
||||
|
||||
local = await _download_image(url, prefix=f"restyled_az{azimuth}_el{elevation}")
|
||||
return {
|
||||
"success": True,
|
||||
"url": f"/generated/{local.name}",
|
||||
"azimuth": azimuth,
|
||||
"elevation": elevation,
|
||||
"error": None,
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
detail = str(e) or f"{type(e).__name__}: {repr(e)}"
|
||||
logger.error(
|
||||
"Stage 3 重绘失败 az=%d el=%d: %s", azimuth, elevation, detail, exc_info=True
|
||||
)
|
||||
return {
|
||||
"success": False,
|
||||
"url": None,
|
||||
"azimuth": azimuth,
|
||||
"elevation": elevation,
|
||||
"error": detail,
|
||||
}
|
||||
|
||||
|
||||
async def restore_batch(
|
||||
original_image: str,
|
||||
frames: list[dict[str, Any]],
|
||||
params: dict[str, Any] | None = None,
|
||||
concurrency: int = 3,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""对一批 frames 做 Stage 3 重绘，带并发控制。
|
||||
|
||||
Args:
|
||||
original_image: 用户原图（所有帧共用同一张风格参考）
|
||||
frames: [{"az", "el", "normal_png", ...}, ...]，normal_png 作为结构参考
|
||||
params: 透传给 restore() 的覆盖参数
|
||||
concurrency: 最大并发数；Replicate 并发打太猛容易撞 rate limit，3 比较稳
|
||||
|
||||
Returns:
|
||||
列表，顺序与 frames 对齐，元素与 restore() 的返回一致。
|
||||
"""
|
||||
if not frames:
|
||||
return []
|
||||
|
||||
sem = asyncio.Semaphore(max(1, concurrency))
|
||||
|
||||
async def _one(f: dict[str, Any]) -> dict[str, Any]:
|
||||
async with sem:
|
||||
return await restore(
|
||||
original_image=original_image,
|
||||
structure_image=f["normal_png"],
|
||||
azimuth=int(f.get("az", 0)),
|
||||
elevation=int(f.get("el", 0)),
|
||||
params=params,
|
||||
)
|
||||
|
||||
return await asyncio.gather(*[_one(f) for f in frames])
|
||||
191
art-agent/backend/app/services/video_frame_extractor.py
Normal file
191
art-agent/backend/app/services/video_frame_extractor.py
Normal file
@@ -0,0 +1,191 @@
|
||||
"""阶段 2：从 Trellis combined_video 抽帧 + 切左右半得到 color / normal 帧。
|
||||
|
||||
Trellis (firtoz/trellis) 部署版 `e8f6c452...` 实测只返回 `combined_video`
|
||||
（color 和 normal 左右并排 mp4，约 5 秒 360° 水平环绕），不再单独返回
|
||||
`color_video` / `normal_video`（详见 PF-20260420-1600）。
|
||||
|
||||
本模块职责：
|
||||
1. 对给定 combined_video 做一次全帧抽取（imageio-ffmpeg 自带二进制，无需系统 ffmpeg）
|
||||
2. 按 azimuth / total_rotation_deg 映射到帧索引并挑出对应帧
|
||||
3. 用 PIL 切左半 → color png、右半 → normal png，落盘到 GENERATED_DIR
|
||||
4. 返回 [{az, el, frame_idx, color_png, normal_png}, ...]
|
||||
|
||||
一期仅支持 elevation=0（Trellis 的 combined_video 就是平视一圈，本身没有仰俯）。
|
||||
需要自由 elevation 时走 Phase 5 的 mesh_renderer 路径（pyrender / Blender headless）。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import subprocess
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import imageio_ffmpeg # 提供静态 ffmpeg 二进制，避免依赖系统 ffmpeg
|
||||
from PIL import Image
|
||||
|
||||
from app.services.image_gen import GENERATED_DIR
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
FFMPEG_EXE = imageio_ffmpeg.get_ffmpeg_exe()
|
||||
|
||||
|
||||
def _resolve_local_path(path_or_url: str) -> Path:
|
||||
"""把 mesh_generator 返回的 '/generated/xxx.mp4' 解析为真实磁盘路径。
|
||||
|
||||
兼容三种输入：
|
||||
- '/generated/xxx.mp4' → GENERATED_DIR / xxx.mp4
|
||||
- 'D:/abs/path/xxx.mp4' → Path 原样
|
||||
- 'xxx.mp4' → Path 原样（调用方自行保证）
|
||||
"""
|
||||
if not path_or_url:
|
||||
raise ValueError("combined_video_path 为空")
|
||||
if path_or_url.startswith("/generated/"):
|
||||
return GENERATED_DIR / Path(path_or_url).name
|
||||
return Path(path_or_url)
|
||||
|
||||
|
||||
def _extract_all_frames(video_path: Path, temp_dir: Path) -> list[Path]:
|
||||
"""把 video 的所有帧抽到 temp_dir，返回按文件名排序的帧列表。
|
||||
|
||||
对 Trellis ~5s 的 combined_video（~100-150 帧）这种量级是 OK 的，
|
||||
比起"用 -ss 逐帧抽取"更省心、角度映射也更稳。
|
||||
"""
|
||||
temp_dir.mkdir(parents=True, exist_ok=True)
|
||||
out_pattern = str(temp_dir / "f%05d.png")
|
||||
cmd = [
|
||||
FFMPEG_EXE,
|
||||
"-y",
|
||||
"-hide_banner",
|
||||
"-loglevel", "error",
|
||||
"-i", str(video_path),
|
||||
"-vsync", "0",
|
||||
out_pattern,
|
||||
]
|
||||
logger.info("ffmpeg extract-all: %s", " ".join(cmd))
|
||||
proc = subprocess.run(cmd, capture_output=True, text=True)
|
||||
if proc.returncode != 0:
|
||||
raise RuntimeError(
|
||||
f"ffmpeg 抽帧失败 (code {proc.returncode}): stderr={proc.stderr[-500:]}"
|
||||
)
|
||||
frames = sorted(temp_dir.glob("f*.png"))
|
||||
if not frames:
|
||||
raise RuntimeError(
|
||||
f"ffmpeg 执行成功但没产出帧：{video_path}；stderr={proc.stderr[-500:]}"
|
||||
)
|
||||
return frames
|
||||
|
||||
|
||||
def _split_combined_frame(
|
||||
frame_path: Path,
|
||||
out_prefix: str,
|
||||
az: int,
|
||||
el: int,
|
||||
) -> tuple[Path, Path]:
|
||||
"""把 combined 帧按中线切成 color（左）+ normal（右）两张图。"""
|
||||
img = Image.open(frame_path)
|
||||
w, h = img.size
|
||||
mid = w // 2
|
||||
color = img.crop((0, 0, mid, h))
|
||||
normal = img.crop((mid, 0, w, h))
|
||||
|
||||
color_name = f"view_{out_prefix}_az{az}_el{el}_color.png"
|
||||
normal_name = f"view_{out_prefix}_az{az}_el{el}_normal.png"
|
||||
color_path = GENERATED_DIR / color_name
|
||||
normal_path = GENERATED_DIR / normal_name
|
||||
color.save(color_path, "PNG")
|
||||
normal.save(normal_path, "PNG")
|
||||
return color_path, normal_path
|
||||
|
||||
|
||||
def _cleanup_dir(dir_path: Path) -> None:
|
||||
"""静默删除临时目录及其所有文件，失败不抛。"""
|
||||
if not dir_path.exists():
|
||||
return
|
||||
for f in dir_path.glob("*"):
|
||||
try:
|
||||
f.unlink()
|
||||
except OSError:
|
||||
pass
|
||||
try:
|
||||
dir_path.rmdir()
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
def _do_extract(
|
||||
video: Path,
|
||||
azimuths: list[int],
|
||||
elevations: list[int],
|
||||
total_rotation_deg: int,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""同步版抽帧逻辑，供 async 版本用 run_in_executor 包起来跑。"""
|
||||
batch_id = uuid.uuid4().hex[:8]
|
||||
temp_dir = GENERATED_DIR / f"_frames_{batch_id}"
|
||||
|
||||
try:
|
||||
frames = _extract_all_frames(video, temp_dir)
|
||||
total = len(frames)
|
||||
logger.info("combined_video %s 抽出 %d 帧，总时长~%.2fs (按 24fps 估)",
|
||||
video.name, total, total / 24.0)
|
||||
|
||||
results: list[dict[str, Any]] = []
|
||||
for az, el in zip(azimuths, elevations):
|
||||
# az 归一化到 [0, total_rotation_deg)，再映射到 [0, total) 整数索引
|
||||
az_mod = az % total_rotation_deg if total_rotation_deg > 0 else 0
|
||||
frame_idx = int(round(az_mod / total_rotation_deg * total)) % total if total > 0 else 0
|
||||
src_frame = frames[frame_idx]
|
||||
color_path, normal_path = _split_combined_frame(src_frame, batch_id, az, el)
|
||||
results.append({
|
||||
"az": az,
|
||||
"el": el,
|
||||
"frame_idx": frame_idx,
|
||||
"color_png": f"/generated/{color_path.name}",
|
||||
"normal_png": f"/generated/{normal_path.name}",
|
||||
})
|
||||
return results
|
||||
finally:
|
||||
_cleanup_dir(temp_dir)
|
||||
|
||||
|
||||
async def extract_by_azimuth(
|
||||
combined_video_path: str,
|
||||
azimuths: list[int],
|
||||
elevations: list[int] | None = None,
|
||||
total_rotation_deg: int = 360,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""按 azimuth 列表从 combined_video 抽帧 + 切左右半。
|
||||
|
||||
Args:
|
||||
combined_video_path: '/generated/trellis_combined_xxx.mp4' 或本地磁盘路径
|
||||
azimuths: 方位角数组，单位度（0-360，允许大于 360 会自动取模）
|
||||
elevations: 仰角数组；长度须等于 azimuths；一期所有值被忽略，保留字段以便后续扩展
|
||||
total_rotation_deg: combined_video 覆盖的水平旋转范围；Trellis 默认 360
|
||||
|
||||
Returns:
|
||||
[{"az": 60, "el": 0, "frame_idx": 20,
|
||||
"color_png": "/generated/view_xxx_az60_el0_color.png",
|
||||
"normal_png": "/generated/view_xxx_az60_el0_normal.png"}, ...]
|
||||
"""
|
||||
video = _resolve_local_path(combined_video_path)
|
||||
if not video.exists():
|
||||
raise FileNotFoundError(f"combined_video 文件不存在：{video}")
|
||||
if not azimuths:
|
||||
return []
|
||||
if elevations is None:
|
||||
elevations = [0] * len(azimuths)
|
||||
if len(elevations) != len(azimuths):
|
||||
raise ValueError(
|
||||
f"azimuths / elevations 长度不一致：{len(azimuths)} vs {len(elevations)}"
|
||||
)
|
||||
|
||||
# ffmpeg + PIL 是 CPU/IO 阻塞的，丢到线程池里跑避免堵死事件循环
|
||||
loop = asyncio.get_running_loop()
|
||||
return await loop.run_in_executor(
|
||||
None,
|
||||
_do_extract,
|
||||
video, azimuths, elevations, total_rotation_deg,
|
||||
)
|
||||
@@ -1,4 +1,8 @@
|
||||
"""视角变换服务 — 调用 Zero123++ 等模型，将单张图片转为多视角图片。"""
|
||||
"""视角变换服务 — 调度单图→多视角生成，支持两种管道：
|
||||
|
||||
- grid 管道：Zero123++ 直接输出 6 宫格拼接图，切分后得到固定 6 视角。
|
||||
- mesh 管道：Trellis/Hunyuan3D → .glb + 预渲染视频 → 抽帧 → 可选风格还原。
|
||||
"""
|
||||
|
||||
import base64
|
||||
import logging
|
||||
@@ -13,6 +17,9 @@ from PIL import Image
|
||||
|
||||
from app.config import get_view_transform_model_config
|
||||
from app.services.image_gen import _make_replicate_client, GENERATED_DIR, to_data_uri
|
||||
from app.services import mesh_generator, style_restorer, video_frame_extractor
|
||||
|
||||
DEFAULT_MESH_AZIMUTHS = [0, 60, 120, 180, 240, 300]
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -29,9 +36,30 @@ def _split_grid_image(
|
||||
|
||||
返回格式：[{"url": "/generated/xxx.png", "azimuth": 30, "elevation": 30}, ...]
|
||||
Grid 布局从左到右、从上到下依次对应 view_labels 中的视角。
|
||||
|
||||
由于不同 Replicate 部署版本输出的拼接图布局可能不同（有的是 3列×2行 的
|
||||
宽图，有的是 2列×3行 的高图），这里会根据实际图像长宽比自动纠正 cols/rows——
|
||||
假定每个单视角为正方形，按 aspect ratio 推断真实布局，防止切图错位。
|
||||
"""
|
||||
img = Image.open(image_path)
|
||||
w, h = img.size
|
||||
|
||||
expected_cells = cols * rows
|
||||
if expected_cells > 0 and len(view_labels) == expected_cells:
|
||||
# 假定每个视角为正方形：真实 aspect ratio = cols / rows
|
||||
configured_ratio = cols / rows
|
||||
actual_ratio = w / h if h else configured_ratio
|
||||
# 如果实际 ratio 与配置差异较大（超过 20%），判定为布局反了
|
||||
if actual_ratio > 0 and (
|
||||
max(configured_ratio, actual_ratio) / min(configured_ratio, actual_ratio) > 1.2
|
||||
):
|
||||
logger.warning(
|
||||
"Grid 布局与实际图像不匹配：配置 cols=%d rows=%d（ratio=%.2f），"
|
||||
"实际图像 %dx%d（ratio=%.2f），自动交换 cols/rows。",
|
||||
cols, rows, configured_ratio, w, h, actual_ratio,
|
||||
)
|
||||
cols, rows = rows, cols
|
||||
|
||||
cell_w = w // cols
|
||||
cell_h = h // rows
|
||||
|
||||
@@ -73,30 +101,60 @@ async def _download_to_local(url: str) -> Path:
|
||||
async def transform_view(
|
||||
image_path: str,
|
||||
model_id: str | None = None,
|
||||
azimuths: list[int] | None = None,
|
||||
elevations: list[int] | None = None,
|
||||
preserve_style: bool = True,
|
||||
) -> dict[str, Any]:
|
||||
"""调用视角变换模型，返回多视角图片列表。
|
||||
|
||||
Args:
|
||||
image_path: 输入图片路径（本地路径、data URI 或 URL）
|
||||
model_id: 视角变换模型 ID，默认 zero123plus
|
||||
model_id: 视角变换模型 ID（zero123plus / trellis / hunyuan3d）
|
||||
azimuths: mesh 管道专用，方位角数组
|
||||
elevations: mesh 管道专用，仰角数组（当前一期仅 0 有效）
|
||||
preserve_style: mesh 管道专用，是否走阶段 3 风格还原
|
||||
|
||||
Returns:
|
||||
统一字段（不同管道可能缺省）：
|
||||
{
|
||||
"success": bool,
|
||||
"images": [{"url": str, "azimuth": int, "elevation": int}, ...],
|
||||
"grid_image": str, # 原始拼接图的本地 URL
|
||||
"grid_image": str | None, # grid 管道：原始拼接图 URL
|
||||
"mesh_url": str | None, # mesh 管道：.glb URL
|
||||
"color_video": str | None, # mesh 管道：360° color video URL
|
||||
"normal_video": str | None, # mesh 管道：法线 video URL
|
||||
"pipeline": "grid" | "mesh",
|
||||
"model_name": str,
|
||||
"error": str | None,
|
||||
}
|
||||
"""
|
||||
config = get_view_transform_model_config(model_id)
|
||||
pipeline = config.get("pipeline", "grid")
|
||||
|
||||
GENERATED_DIR.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
if pipeline == "mesh":
|
||||
return await _transform_view_mesh(
|
||||
image_path=image_path,
|
||||
config=config,
|
||||
azimuths=azimuths,
|
||||
elevations=elevations,
|
||||
preserve_style=preserve_style,
|
||||
)
|
||||
|
||||
return await _transform_view_grid(image_path=image_path, config=config)
|
||||
|
||||
|
||||
async def _transform_view_grid(
|
||||
image_path: str,
|
||||
config: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
"""原 Zero123++ 逻辑：单次调用 → grid 图 → 切 6 张。"""
|
||||
model_name = config["name"]
|
||||
replicate_model_id = config["model_id"]
|
||||
grid_layout = config["grid_layout"]
|
||||
output_views = config["output_views"]
|
||||
|
||||
GENERATED_DIR.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
try:
|
||||
image_uri = to_data_uri(image_path)
|
||||
input_params = {"image": image_uri}
|
||||
@@ -141,6 +199,10 @@ async def transform_view(
|
||||
"success": True,
|
||||
"images": view_images,
|
||||
"grid_image": grid_local_url,
|
||||
"mesh_url": None,
|
||||
"color_video": None,
|
||||
"normal_video": None,
|
||||
"pipeline": "grid",
|
||||
"model_name": model_name,
|
||||
"error": None,
|
||||
}
|
||||
@@ -152,6 +214,144 @@ async def transform_view(
|
||||
"success": False,
|
||||
"images": [],
|
||||
"grid_image": None,
|
||||
"mesh_url": None,
|
||||
"color_video": None,
|
||||
"normal_video": None,
|
||||
"pipeline": "grid",
|
||||
"model_name": model_name,
|
||||
"error": detail,
|
||||
}
|
||||
|
||||
|
||||
async def _transform_view_mesh(
|
||||
image_path: str,
|
||||
config: dict[str, Any],
|
||||
azimuths: list[int] | None,
|
||||
elevations: list[int] | None,
|
||||
preserve_style: bool,
|
||||
) -> dict[str, Any]:
|
||||
"""Mesh 管道调度器：阶段 1 生成 mesh + video；阶段 2/3 后续 Phase 接入。
|
||||
|
||||
Phase 1 验收状态：只跑阶段 1，返回 .glb / color_video / normal_video 本地 URL，
|
||||
images 列表为空。Phase 2 完成后会填充 images（raw 抽帧结果），
|
||||
Phase 3 完成后 images 会是经 IP-Adapter+ControlNet 重绘的最终图。
|
||||
"""
|
||||
model_name = config["name"]
|
||||
model_id = config["id"]
|
||||
|
||||
# === Stage 1：生成 mesh + 预渲染视频 ===
|
||||
stage1 = await mesh_generator.generate_mesh(image_path, model_id=model_id)
|
||||
if not stage1.get("success"):
|
||||
return {
|
||||
"success": False,
|
||||
"images": [],
|
||||
"grid_image": None,
|
||||
"mesh_url": None,
|
||||
"color_video": None,
|
||||
"normal_video": None,
|
||||
"combined_video": None,
|
||||
"pipeline": "mesh",
|
||||
"model_name": model_name,
|
||||
"error": stage1.get("error") or "阶段 1 失败",
|
||||
}
|
||||
|
||||
# === Stage 2：抽帧 + 切左右半（color + normal）===
|
||||
# azimuths / elevations 缺省时用默认 6 视角；长度对齐
|
||||
if not azimuths:
|
||||
azimuths = list(DEFAULT_MESH_AZIMUTHS)
|
||||
if elevations is None or not elevations:
|
||||
elevations = [0] * len(azimuths)
|
||||
elif len(elevations) < len(azimuths):
|
||||
elevations = list(elevations) + [0] * (len(azimuths) - len(elevations))
|
||||
elif len(elevations) > len(azimuths):
|
||||
elevations = list(elevations[: len(azimuths)])
|
||||
|
||||
# 优先 color_video（当前 Trellis 部署版恒为 None），缺则 fallback 到 combined_video
|
||||
combined_video_path = stage1.get("combined_video_path")
|
||||
frames: list[dict[str, Any]] = []
|
||||
stage2_error: str | None = None
|
||||
if combined_video_path:
|
||||
try:
|
||||
frames = await video_frame_extractor.extract_by_azimuth(
|
||||
combined_video_path=combined_video_path,
|
||||
azimuths=azimuths,
|
||||
elevations=elevations,
|
||||
total_rotation_deg=360,
|
||||
)
|
||||
except Exception as e:
|
||||
stage2_error = f"Stage 2 抽帧失败：{e}"
|
||||
logger.error(stage2_error, exc_info=True)
|
||||
else:
|
||||
stage2_error = "Trellis 未返回 combined_video，无法执行 Stage 2 抽帧"
|
||||
logger.error(stage2_error)
|
||||
|
||||
# === Stage 3：风格还原（可选）===
|
||||
# preserve_style=True 时对每一帧的 normal_png 做 IP-Adapter+ControlNet 重绘，
|
||||
# 用户原图做风格参考；失败的单帧会回退为 Stage 2 的 color 帧，不阻塞其它帧。
|
||||
stage3_error: str | None = None
|
||||
stage3_results: list[dict[str, Any]] = []
|
||||
did_stage3 = bool(preserve_style) and bool(frames)
|
||||
if did_stage3:
|
||||
try:
|
||||
stage3_results = await style_restorer.restore_batch(
|
||||
original_image=image_path,
|
||||
frames=frames,
|
||||
concurrency=3,
|
||||
)
|
||||
except Exception as e:
|
||||
stage3_error = f"Stage 3 风格还原失败：{e}"
|
||||
logger.error(stage3_error, exc_info=True)
|
||||
|
||||
# 组装最终 images：
|
||||
# - Stage 3 成功的帧 → 用重绘后的 URL
|
||||
# - Stage 3 失败 / 未跑 → 回退到 Stage 2 的 color 帧（保证前端至少能看到点东西）
|
||||
images: list[dict[str, Any]] = []
|
||||
for i, f in enumerate(frames):
|
||||
restyled = stage3_results[i] if i < len(stage3_results) else None
|
||||
if restyled and restyled.get("success") and restyled.get("url"):
|
||||
final_url = restyled["url"]
|
||||
frame_error = None
|
||||
else:
|
||||
final_url = f["color_png"]
|
||||
frame_error = (restyled or {}).get("error") if did_stage3 else None
|
||||
images.append({
|
||||
"url": final_url,
|
||||
"azimuth": f["az"],
|
||||
"elevation": f["el"],
|
||||
"normal_url": f["normal_png"],
|
||||
"color_url": f["color_png"],
|
||||
"frame_idx": f["frame_idx"],
|
||||
"restyled": bool(restyled and restyled.get("success")),
|
||||
**({"restyle_error": frame_error} if frame_error else {}),
|
||||
})
|
||||
|
||||
# Stage 2 失败不把整体置为失败——Stage 1 的 mesh / video 产物仍有交付价值
|
||||
# Stage 3 部分失败也不置为整体失败（每帧都有 Stage 2 回退）
|
||||
success = bool(stage1.get("success")) and (stage2_error is None)
|
||||
|
||||
# 合并错误信息（Stage 2 和 Stage 3 错误都要暴露给前端）
|
||||
error_parts = [e for e in (stage2_error, stage3_error) if e]
|
||||
error_out = "; ".join(error_parts) if error_parts else None
|
||||
|
||||
if did_stage3:
|
||||
restyled_ok = sum(1 for img in images if img.get("restyled"))
|
||||
stage_label = f"stage3 ({restyled_ok}/{len(images)} restyled)"
|
||||
elif frames:
|
||||
stage_label = "stage2"
|
||||
else:
|
||||
stage_label = "stage1"
|
||||
|
||||
return {
|
||||
"success": success,
|
||||
"images": images,
|
||||
"grid_image": None,
|
||||
"mesh_url": stage1.get("glb_path"),
|
||||
"color_video": stage1.get("color_video_path"),
|
||||
"normal_video": stage1.get("normal_video_path"),
|
||||
"combined_video": combined_video_path,
|
||||
"pipeline": "mesh",
|
||||
"model_name": model_name,
|
||||
"error": error_out,
|
||||
"_stage": stage_label,
|
||||
"_preserve_style_requested": bool(preserve_style),
|
||||
}
|
||||
|
||||
Binary file not shown.
@@ -1,5 +1,6 @@
|
||||
fastapi==0.135.3
|
||||
uvicorn==0.44.0
|
||||
watchfiles==0.24.0
|
||||
openai==2.32.0
|
||||
replicate==1.0.7
|
||||
sse-starlette==3.3.4
|
||||
@@ -7,6 +8,7 @@ python-multipart==0.0.26
|
||||
httpx==0.28.1
|
||||
python-dotenv==1.2.2
|
||||
Pillow==12.2.0
|
||||
imageio-ffmpeg==0.5.1
|
||||
mem0ai==1.0.11
|
||||
ollama==0.6.1
|
||||
sqlmodel==0.0.38
|
||||
|
||||
107
art-agent/backend/scripts/test_frame_extract.py
Normal file
107
art-agent/backend/scripts/test_frame_extract.py
Normal file
@@ -0,0 +1,107 @@
|
||||
"""Phase 2 验收脚本：独立验证 combined_video 抽帧 + 左右切分逻辑。
|
||||
|
||||
不需要再跑一次 Trellis（省钱），直接用 generated/ 下已有的
|
||||
`trellis_combined_*.mp4` 做验证。
|
||||
|
||||
用法（在 backend 目录激活 venv 之后）：
|
||||
python scripts/test_frame_extract.py
|
||||
python scripts/test_frame_extract.py generated/trellis_combined_211176d3940b.mp4
|
||||
|
||||
成功标准：
|
||||
- 退出码 0
|
||||
- 打印的 color_png / normal_png 均为 /generated/view_xxx.png
|
||||
- 到 backend/generated/ 能看到 6 对 color + normal 文件（12 张）且尺寸接近正方形
|
||||
- 目视检查：每组 color 图按 0/60/120/180/240/300° 依次环绕建筑一圈
|
||||
"""
|
||||
|
||||
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 video_frame_extractor # noqa: E402
|
||||
from app.services.image_gen import GENERATED_DIR # noqa: E402
|
||||
|
||||
|
||||
def _pick_default_video() -> str:
|
||||
candidates = sorted(
|
||||
GENERATED_DIR.glob("trellis_combined_*.mp4"),
|
||||
key=lambda p: p.stat().st_mtime,
|
||||
reverse=True,
|
||||
)
|
||||
if not candidates:
|
||||
raise FileNotFoundError(
|
||||
"generated/ 下找不到 trellis_combined_*.mp4；请先跑一次 Phase 1 "
|
||||
"（scripts/test_trellis.py）或显式传入视频路径"
|
||||
)
|
||||
return f"/generated/{candidates[0].name}"
|
||||
|
||||
|
||||
async def main() -> int:
|
||||
video_arg = sys.argv[1] if len(sys.argv) > 1 else _pick_default_video()
|
||||
# 允许传相对路径或纯文件名
|
||||
if not video_arg.startswith("/generated/"):
|
||||
if video_arg.startswith("generated/"):
|
||||
video_arg = "/" + video_arg.replace("\\", "/")
|
||||
elif not video_arg.startswith("/") and not Path(video_arg).is_absolute():
|
||||
# 如果只给了个文件名，假定在 generated/ 下
|
||||
if (GENERATED_DIR / video_arg).exists():
|
||||
video_arg = f"/generated/{video_arg}"
|
||||
|
||||
print(f"[test_frame_extract] 输入视频：{video_arg}")
|
||||
print("[test_frame_extract] 抽取 6 视角：az=[0,60,120,180,240,300]，el 全 0")
|
||||
|
||||
azimuths = [0, 60, 120, 180, 240, 300]
|
||||
elevations = [0] * len(azimuths)
|
||||
|
||||
frames = await video_frame_extractor.extract_by_azimuth(
|
||||
combined_video_path=video_arg,
|
||||
azimuths=azimuths,
|
||||
elevations=elevations,
|
||||
total_rotation_deg=360,
|
||||
)
|
||||
|
||||
print("[test_frame_extract] 抽帧结果：")
|
||||
print(json.dumps(frames, ensure_ascii=False, indent=2))
|
||||
|
||||
# 简单完整性校验
|
||||
all_exist = True
|
||||
for f in frames:
|
||||
for key in ("color_png", "normal_png"):
|
||||
rel = f.get(key)
|
||||
if not rel:
|
||||
print(f"[test_frame_extract] 警告：帧 az={f['az']} 缺少 {key}")
|
||||
all_exist = False
|
||||
continue
|
||||
local = GENERATED_DIR / Path(rel).name
|
||||
size = local.stat().st_size if local.exists() else -1
|
||||
status = "OK" if size > 0 else "MISSING"
|
||||
print(f" [{status}] az={f['az']:3d} {key}: {rel} ({size} bytes)")
|
||||
if size <= 0:
|
||||
all_exist = False
|
||||
|
||||
if not all_exist:
|
||||
print("[test_frame_extract] 部分文件缺失，判定失败")
|
||||
return 1
|
||||
|
||||
print(
|
||||
f"[test_frame_extract] 成功：共 {len(frames)} 个视角，"
|
||||
f"color+normal 合计 {len(frames) * 2} 张已落盘到 {GENERATED_DIR}"
|
||||
)
|
||||
print(
|
||||
"[test_frame_extract] 请目视检查：打开 generated/ 下任意 view_*_color.png，"
|
||||
"确认 az=0 正面、az=180 背面、其余按顺序环绕；normal 帧应为蓝紫调法线图。"
|
||||
)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(asyncio.run(main()))
|
||||
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()))
|
||||
72
art-agent/backend/scripts/test_trellis.py
Normal file
72
art-agent/backend/scripts/test_trellis.py
Normal file
@@ -0,0 +1,72 @@
|
||||
"""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()))
|
||||
Reference in New Issue
Block a user