192 lines
6.7 KiB
Python
192 lines
6.7 KiB
Python
"""阶段 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,
|
||
)
|