Spaces:
Running on Zero
Running on Zero
| from __future__ import annotations | |
| from pathlib import Path | |
| import cv2 | |
| import numpy as np | |
| import torch | |
| from tqdm import tqdm | |
| from camera import Camera | |
| from tracker.inference.preprocess import enlarge_bbox | |
| from .visualizer import Visualizer | |
| def save_tracking_video( | |
| output_path: str | Path, | |
| video_input, | |
| inference, | |
| camera: Camera, | |
| visualizer: Visualizer, | |
| vertices: torch.Tensor, | |
| fallback_fps: float, | |
| ) -> None: | |
| fps = float(fallback_fps) if video_input.fps is None else float(video_input.fps) | |
| video_input.reset() | |
| frame_size = (int(inference.image_width * 2), int(inference.image_height * 2)) | |
| writer = cv2.VideoWriter( | |
| str(output_path), cv2.VideoWriter_fourcc(*"mp4v"), fps, frame_size | |
| ) | |
| if not writer.isOpened(): | |
| raise OSError(f"Failed to open video writer: {output_path}") | |
| try: | |
| num_frames = int(inference.vertex_coord.shape[0]) | |
| for frame_id in tqdm(range(num_frames), desc="visualize video"): | |
| frame_rgb = video_input.read() | |
| if frame_rgb is None: | |
| raise ValueError(f"Unexpected end of input at frame {frame_id}.") | |
| alignment = visualizer.vis_align(frame_rgb, inference.vertex_coord[frame_id]) | |
| alignment = visualizer.vis_contour(alignment, inference.vertex_coord[frame_id]) | |
| bbox = np.round(enlarge_bbox(inference.face_bbox[frame_id], 1.3)).astype(np.int32) | |
| cv2.rectangle(alignment, bbox[:2], bbox[2:], (255, 0, 0), 2) | |
| overlay, shading = visualizer.vis_flame_shading_overlay( | |
| camera, frame_rgb, vertices[frame_id], alpha=0.5, return_shading=True | |
| ) | |
| result = np.ascontiguousarray( | |
| np.concatenate( | |
| [ | |
| np.concatenate([frame_rgb, alignment], axis=1), | |
| np.concatenate([overlay, shading], axis=1), | |
| ], | |
| axis=0, | |
| ) | |
| ) | |
| writer.write(cv2.cvtColor(result, cv2.COLOR_RGB2BGR)) | |
| finally: | |
| writer.release() | |