引入3d模型

2026-04-20 21:52:35 +08:00
parent ee5cec6de4
commit d6e1797f08
23 changed files with 1890 additions and 288 deletions
--- a/art-agent/backend/app/services/video_frame_extractor.py
+++ b/art-agent/backend/app/services/video_frame_extractor.py
@@ -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,
+    )