Spaces:
Running on Zero
Running on Zero
| from __future__ import annotations | |
| from pathlib import Path | |
| import cv2 | |
| import numpy as np | |
| from camera import IntrinsicsCamera | |
| from .visualizer import Visualizer | |
| def save_nersemble_tracking_visualizations( | |
| tracking_results_path: Path, | |
| video_dir: Path, | |
| output_dir: Path, | |
| intrinsics: np.ndarray, | |
| world_to_camera: dict[str, np.ndarray], | |
| camera_serials: list[str], | |
| visualizer: Visualizer, | |
| downsample: int, | |
| num_frames: int = 8, | |
| ) -> None: | |
| with np.load(tracking_results_path, allow_pickle=False) as flame_data: | |
| frame_ids = flame_data["frame_ids"].astype(np.int32, copy=False) | |
| posed_vertices = flame_data["posed_vertices"].astype(np.float32, copy=False) | |
| width = int(flame_data["image_width"]) | |
| height = int(flame_data["image_height"]) | |
| if len(frame_ids) == 0: | |
| return | |
| selected = np.linspace(0, len(frame_ids) - 1, min(num_frames, len(frame_ids)), dtype=np.int32) | |
| output_dir.mkdir(parents=True, exist_ok=True) | |
| for stale_image in output_dir.glob("frame_*.jpg"): | |
| stale_image.unlink() | |
| captures: dict[str, cv2.VideoCapture] = {} | |
| cameras: dict[str, IntrinsicsCamera] = {} | |
| scaled_intrinsics = np.asarray(intrinsics, dtype=np.float32).copy() | |
| scaled_intrinsics[:2, :] /= float(downsample) | |
| for serial in camera_serials: | |
| video_path = video_dir / f"cam_{serial}.mp4" | |
| capture = cv2.VideoCapture(str(video_path)) | |
| if not capture.isOpened(): | |
| for opened_capture in captures.values(): | |
| opened_capture.release() | |
| raise RuntimeError(f"Failed to open visualization video: {video_path}") | |
| captures[serial] = capture | |
| transform = world_to_camera[serial] | |
| cameras[serial] = IntrinsicsCamera( | |
| K=scaled_intrinsics, | |
| R=transform[:3, :3], | |
| T=transform[:3, 3], | |
| width=width, | |
| height=height, | |
| ) | |
| try: | |
| for sample_index in selected.tolist(): | |
| frame_id = int(frame_ids[sample_index]) | |
| panels = [] | |
| for serial in camera_serials: | |
| capture = captures[serial] | |
| capture.set(cv2.CAP_PROP_POS_FRAMES, frame_id) | |
| valid, frame_bgr = capture.read() | |
| if not valid: | |
| continue | |
| frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB) | |
| panels.append( | |
| visualizer.vis_flame_shading_overlay( | |
| cameras[serial], | |
| frame_rgb, | |
| posed_vertices[sample_index], | |
| alpha=0.5, | |
| ) | |
| ) | |
| if panels: | |
| composite = np.concatenate(panels, axis=1) | |
| output_path = output_dir / f"frame_{frame_id:06d}.jpg" | |
| cv2.imwrite( | |
| str(output_path), | |
| cv2.cvtColor(composite, cv2.COLOR_RGB2BGR), | |
| [cv2.IMWRITE_JPEG_QUALITY, 90], | |
| ) | |
| finally: | |
| for capture in captures.values(): | |
| capture.release() | |