Spaces:
Running
Running
| # Copyright (c) Meta Platforms, Inc. and affiliates. | |
| import plotly.graph_objects as go | |
| from plotly.graph_objects import Figure | |
| import numpy as np | |
| import os | |
| import imageio | |
| from pathlib import Path | |
| from typing import List, Optional, Tuple, Union, Dict, Any | |
| from PIL import Image | |
| import io | |
| import numpy as np | |
| from tqdm import tqdm | |
| def img_bytes_to_np(img_bytes): | |
| return np.array(Image.open(io.BytesIO(img_bytes))) | |
| def make_video( | |
| scene: Figure, | |
| output_path: str = "scene_video.mp4", | |
| fps: int = 15, | |
| duration: int = 1, | |
| camera_trajectory: Optional[List[Dict[str, Any]]] = None, | |
| temp_dir: Optional[str] = None, | |
| trajectory_kwargs: Optional[Dict[str, Any]] = None, | |
| ) -> str: | |
| """ | |
| Creates a video by updating the camera view location and saving snapshots. | |
| Args: | |
| scene: A Plotly Figure object, typically created with plot_tdfy_scene | |
| output_path: Path to save the output video file | |
| fps: Frames per second for the output video | |
| duration: Duration of the video in seconds | |
| camera_trajectory: List of camera positions. If None, creates a default circular trajectory. | |
| Each item should be a dict with eye, center, and up keys as expected by Plotly's scene.camera. | |
| temp_dir: Directory to store temporary frame images. If None, uses ./tmp_frames | |
| Returns: | |
| Path to the saved video file | |
| """ | |
| if not scene._has_subplots(): | |
| raise ValueError("Scene must have subplots to create a video") | |
| num_frames = fps * duration | |
| if camera_trajectory is None: | |
| if trajectory_kwargs is None: | |
| trajectory_kwargs = {} | |
| camera_trajectory = _create_default_camera_trajectory( | |
| num_frames, **trajectory_kwargs | |
| ) | |
| frames = [] | |
| for i, camera_pos in tqdm(enumerate(camera_trajectory), total=num_frames): | |
| # update the camera position | |
| scene.update_scenes(camera=camera_pos) | |
| img_as_png = scene.to_image(engine="kaleido") | |
| frames.append(img_bytes_to_np(img_as_png)) | |
| return frames | |
| def _create_default_camera_trajectory( | |
| num_frames: int, | |
| axis: str = "y", | |
| elevation: float = 1.0, | |
| radius: float = 2.0, | |
| **kwargs, | |
| ) -> List[Dict[str, Any]]: | |
| """ | |
| Creates a default camera trajectory, rotating around the scene in a circle. | |
| Args: | |
| num_frames: Number of frames in the trajectory | |
| axis: Axis to rotate around ('x', 'y', or 'z') | |
| Returns: | |
| List of camera positions | |
| """ | |
| trajectory = [] | |
| # Create a circular path | |
| for i in range(num_frames): | |
| angle = (i / num_frames) * 2 * np.pi | |
| # Default position (all zeros) | |
| eye_x, eye_y, eye_z = 0.0, 0.0, 0.0 | |
| # Calculate camera position based on selected axis | |
| if axis.lower() == "z": | |
| # Rotate in the xy-plane (around z-axis) | |
| eye_x = radius * np.sin(angle) | |
| eye_y = radius * np.cos(angle) | |
| eye_z = elevation # Slightly above the scene | |
| up = {"x": 0, "y": 0, "z": 1} | |
| elif axis.lower() == "y": | |
| # Rotate in the xz-plane (around y-axis) | |
| eye_x = radius * np.sin(angle) | |
| eye_z = radius * np.cos(angle) | |
| eye_y = elevation # Slightly offset from y-axis | |
| up = {"x": 0, "y": 1, "z": 0} | |
| elif axis.lower() == "x": | |
| # Rotate in the yz-plane (around x-axis) | |
| eye_y = radius * np.sin(angle) | |
| eye_z = radius * np.cos(angle) | |
| eye_x = elevation # Slightly offset from x-axis | |
| up = {"x": 1, "y": 0, "z": 0} | |
| else: | |
| raise ValueError(f"Invalid axis: {axis}. Must be 'x', 'y', or 'z'") | |
| camera_pos = { | |
| "eye": {"x": eye_x, "y": eye_y, "z": eye_z}, | |
| "center": {"x": 0, "y": 0, "z": 0}, # Look at center | |
| "up": up, # Orientation based on rotation axis | |
| } | |
| trajectory.append(camera_pos) | |
| return trajectory | |