Spaces:
Running on Zero
Running on Zero
| 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}") | |