Spaces:
Running on Zero
Running on Zero
File size: 1,356 Bytes
7826761 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 | from __future__ import annotations
from pathlib import Path
import cv2
import numpy as np
import torch
from camera import Camera
from tracker.inference.preprocess import enlarge_bbox
from .visualizer import Visualizer
def render_reconstruction_panels(
visualizer: Visualizer,
camera: Camera,
image_rgb: np.ndarray,
inference,
vertices: torch.Tensor,
alpha: float = 0.5,
) -> np.ndarray:
alignment = visualizer.vis_align(image_rgb, inference.vertex_coord)
alignment = visualizer.vis_contour(alignment, inference.vertex_coord)
bbox = np.round(enlarge_bbox(inference.face_bbox, 1.3)).astype(np.int32)
cv2.rectangle(alignment, tuple(bbox[:2]), tuple(bbox[2:]), (255, 0, 0), 2)
overlay, shading = visualizer.vis_flame_shading_overlay(
camera, image_rgb, vertices, alpha=alpha, return_shading=True
)
return np.ascontiguousarray(
np.concatenate(
[
np.concatenate([image_rgb, alignment], axis=1),
np.concatenate([overlay, shading], axis=1),
],
axis=0,
)
)
def save_reconstruction_image(path: str | Path, image_rgb: np.ndarray) -> None:
path = Path(path)
if not cv2.imwrite(str(path), cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR)):
raise OSError(f"Failed to write visualization image: {path}")
|