358 lines
13 KiB
Python
358 lines
13 KiB
Python
"""视角变换服务 — 调度单图→多视角生成，支持两种管道：
|
||
|
||
- grid 管道：Zero123++ 直接输出 6 宫格拼接图，切分后得到固定 6 视角。
|
||
- mesh 管道：Trellis/Hunyuan3D → .glb + 预渲染视频 → 抽帧 → 可选风格还原。
|
||
"""
|
||
|
||
import base64
|
||
import logging
|
||
import mimetypes
|
||
import os
|
||
import uuid
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
import httpx
|
||
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__)
|
||
|
||
_replicate_client = _make_replicate_client()
|
||
|
||
|
||
def _split_grid_image(
|
||
image_path: Path,
|
||
cols: int,
|
||
rows: int,
|
||
view_labels: list[dict[str, int]],
|
||
) -> list[dict[str, Any]]:
|
||
"""将 Zero123++ 输出的 grid 拼接图切分为独立的视角图片。
|
||
|
||
返回格式：[{"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
|
||
|
||
results: list[dict[str, Any]] = []
|
||
for idx, view in enumerate(view_labels):
|
||
row_idx = idx // cols
|
||
col_idx = idx % cols
|
||
left = col_idx * cell_w
|
||
upper = row_idx * cell_h
|
||
right = left + cell_w
|
||
lower = upper + cell_h
|
||
|
||
cropped = img.crop((left, upper, right, lower))
|
||
filename = f"view_{uuid.uuid4().hex[:8]}_az{view['azimuth']}_el{view['elevation']}.png"
|
||
filepath = GENERATED_DIR / filename
|
||
cropped.save(filepath, "PNG")
|
||
|
||
results.append({
|
||
"url": f"/generated/{filename}",
|
||
"azimuth": view["azimuth"],
|
||
"elevation": view["elevation"],
|
||
})
|
||
|
||
return results
|
||
|
||
|
||
async def _download_to_local(url: str) -> Path:
|
||
"""下载远程图片到本地 generated/ 目录，返回本地路径。"""
|
||
filename = f"grid_{uuid.uuid4().hex}.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(120.0)) as client:
|
||
resp = await client.get(url, follow_redirects=True)
|
||
resp.raise_for_status()
|
||
filepath.write_bytes(resp.content)
|
||
return filepath
|
||
|
||
|
||
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 / trellis / hunyuan3d）
|
||
azimuths: mesh 管道专用，方位角数组
|
||
elevations: mesh 管道专用，仰角数组（当前一期仅 0 有效）
|
||
preserve_style: mesh 管道专用，是否走阶段 3 风格还原
|
||
|
||
Returns:
|
||
统一字段（不同管道可能缺省）：
|
||
{
|
||
"success": bool,
|
||
"images": [{"url": str, "azimuth": int, "elevation": int}, ...],
|
||
"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"]
|
||
|
||
try:
|
||
image_uri = to_data_uri(image_path)
|
||
input_params = {"image": image_uri}
|
||
|
||
logger.info("调用 %s 进行视角变换...", model_name)
|
||
output = await _replicate_client.async_run(
|
||
replicate_model_id, input=input_params, wait=False
|
||
)
|
||
|
||
# Zero123++ 返回单张 grid 拼接图
|
||
items = output if isinstance(output, list) else [output]
|
||
if not items:
|
||
return {
|
||
"success": False,
|
||
"images": [],
|
||
"grid_image": None,
|
||
"model_name": model_name,
|
||
"error": "模型未返回任何输出",
|
||
}
|
||
|
||
grid_url = str(items[0])
|
||
if not (grid_url.startswith("http://") or grid_url.startswith("https://")):
|
||
return {
|
||
"success": False,
|
||
"images": [],
|
||
"grid_image": None,
|
||
"model_name": model_name,
|
||
"error": f"模型返回非图片内容: {grid_url[:200]}",
|
||
}
|
||
|
||
grid_local_path = await _download_to_local(grid_url)
|
||
grid_local_url = f"/generated/{grid_local_path.name}"
|
||
|
||
view_images = _split_grid_image(
|
||
grid_local_path,
|
||
cols=grid_layout["cols"],
|
||
rows=grid_layout["rows"],
|
||
view_labels=output_views,
|
||
)
|
||
|
||
return {
|
||
"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,
|
||
}
|
||
|
||
except Exception as e:
|
||
detail = str(e) or f"{type(e).__name__}: {repr(e)}"
|
||
logger.error("视角变换失败: %s", detail, exc_info=True)
|
||
return {
|
||
"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),
|
||
}
|