Spaces:
Running on Zero
Running on Zero
| """World-anchored checkerboard with a camera that translates with the root.""" | |
| from pathlib import Path | |
| import shutil | |
| import matplotlib | |
| matplotlib.use("Agg") | |
| import matplotlib.pyplot as plt | |
| from matplotlib.animation import FFMpegWriter | |
| from mpl_toolkits.mplot3d.art3d import Poly3DCollection | |
| import numpy as np | |
| CHAINS = [ | |
| [0, 2, 5, 8, 11], [0, 1, 4, 7, 10], | |
| [0, 3, 6, 9, 12, 15], [9, 14, 17, 19, 21], [9, 13, 16, 18, 20], | |
| ] | |
| COLORS = ["#43d5ca", "#3899ed", "#f3c46c", "#e586ad", "#b7a0f5"] | |
| FLOOR_COLORS = ["#e4d8bd", "#627988"] | |
| def world_to_camera(points, root): | |
| """Apply the identical horizontal translation to body, floor and landmarks.""" | |
| result = np.array(points, dtype=np.float64, copy=True) | |
| result[..., 0] -= root[0] | |
| result[..., 2] -= root[2] | |
| return result | |
| def checkerboard(root, radius, tile_size=0.5): | |
| """Tiles retain world coordinates and parity when the camera crosses a cell.""" | |
| x0, z0 = np.floor(np.asarray(root)[[0, 2]] / tile_size).astype(int) | |
| extent = int(np.ceil(3 * radius / tile_size)) + 2 | |
| tiles, colors = [], [] | |
| for ix in range(x0 - extent, x0 + extent + 1): | |
| for iz in range(z0 - extent, z0 + extent + 1): | |
| x, z = ix * tile_size, iz * tile_size | |
| world = [[x, 0, z], [x + tile_size, 0, z], | |
| [x + tile_size, 0, z + tile_size], [x, 0, z + tile_size]] | |
| camera = world_to_camera(world, root) | |
| tiles.append(camera[:, [0, 2, 1]]) | |
| colors.append(FLOOR_COLORS[(ix + iz) % 2]) | |
| return tiles, colors | |
| def render_mp4(joints, destination, fps=20): | |
| joints = np.asarray(joints, dtype=np.float64) | |
| if joints.ndim != 3 or joints.shape[1:] != (22, 3) or len(joints) < 2: | |
| raise ValueError("Rendering expects positions with shape (T,22,3).") | |
| if not np.isfinite(joints).all(): | |
| raise ValueError("Invalid joint positions for rendering.") | |
| if not np.isfinite(fps) or not 1 <= float(fps) <= 120: | |
| raise ValueError("Rendering requires a frame rate between 1 and 120 fps.") | |
| ffmpeg = shutil.which("ffmpeg") | |
| if not ffmpeg: | |
| try: | |
| import imageio_ffmpeg | |
| ffmpeg = imageio_ffmpeg.get_ffmpeg_exe() | |
| except (ImportError, RuntimeError) as error: | |
| raise RuntimeError("FFmpeg was not found. Install it before rendering MP4 files.") from error | |
| matplotlib.rcParams["animation.ffmpeg_path"] = ffmpeg | |
| destination = Path(destination) | |
| destination.parent.mkdir(parents=True, exist_ok=True) | |
| positions = joints.copy() | |
| positions[:, :, 1] -= positions[:, :, 1].min() | |
| local = positions.copy() | |
| local[:, :, 0] -= positions[:, :1, 0] | |
| local[:, :, 2] -= positions[:, :1, 2] | |
| height = max(float(positions[:, :, 1].max()), 1.0) | |
| radius = max(float(np.abs(local[:, :, [0, 2]]).max()) * 1.3, height * 0.85) | |
| fig = plt.figure(figsize=(6, 6), dpi=96, facecolor="#111c2d") | |
| ax = fig.add_subplot(111, projection="3d", facecolor="#111c2d", computed_zorder=False) | |
| ax.set_xlim(-radius, radius) | |
| ax.set_ylim(-radius, radius) | |
| ax.set_zlim(-0.02, height * 1.15) | |
| ax.set_box_aspect((2 * radius, 2 * radius, height * 1.15)) | |
| ax.view_init(elev=20, azim=-65) | |
| ax.set_axis_off() | |
| fig.subplots_adjust(0, 0.05, 1, 1) | |
| floor = Poly3DCollection([], edgecolors="#344353", linewidths=0.3, zorder=0) | |
| ax.add_collection3d(floor) | |
| trail, = ax.plot([], [], [], color="#eeae56", lw=1.5, alpha=0.8, zorder=2) | |
| start, = ax.plot([], [], [], marker="o", markersize=6, color="#f07d55", zorder=2) | |
| lines = [ax.plot([], [], [], color=color, lw=3.2, solid_capstyle="round", zorder=3)[0] | |
| for color in COLORS] | |
| fig.text(0.5, 0.035, "World-fixed floor 路 0.5 m tiles 路 camera follows the person", | |
| ha="center", color="#bdc9d8", fontsize=8, | |
| bbox={"facecolor": "#111c2d", "edgecolor": "none", "pad": 4}) | |
| distance = fig.text(0.5, 0.085, "", ha="center", color="#e4d8bd", fontsize=9, | |
| bbox={"facecolor": "#111c2d", "edgecolor": "none", "pad": 4}) | |
| roots = positions[:, 0].copy() | |
| roots[:, 1] = 0.006 | |
| writer = FFMpegWriter(fps=fps, codec="libx264", bitrate=1800, | |
| extra_args=["-pix_fmt", "yuv420p", "-movflags", "+faststart", "-threads", "2"]) | |
| try: | |
| with writer.saving(fig, str(destination), dpi=96): | |
| for index, frame in enumerate(positions): | |
| root = frame[0] | |
| tiles, colors = checkerboard(root, radius) | |
| floor.set_verts(tiles) | |
| floor.set_facecolor(colors) | |
| for line, chain in zip(lines, CHAINS): | |
| points = local[index, chain] | |
| line.set_data(points[:, 0], points[:, 2]) | |
| line.set_3d_properties(points[:, 1]) | |
| path = world_to_camera(roots[:index + 1], root) | |
| outside = (np.abs(path[:, 0]) > radius) | (np.abs(path[:, 2]) > radius) | |
| path[outside] = np.nan | |
| trail.set_data(path[:, 0], path[:, 2]) | |
| trail.set_3d_properties(path[:, 1]) | |
| origin = world_to_camera(roots[0], root) | |
| visible_origin = abs(origin[0]) <= radius and abs(origin[2]) <= radius | |
| start.set_data([origin[0]] if visible_origin else [], [origin[2]] if visible_origin else []) | |
| start.set_3d_properties([origin[1]] if visible_origin else []) | |
| traveled = np.linalg.norm(frame[0, [0, 2]] - positions[0, 0, [0, 2]]) | |
| distance.set_text(f"{index / fps:.2f} s 路 displacement from start: {traveled:.2f} m") | |
| writer.grab_frame(facecolor=fig.get_facecolor()) | |
| finally: | |
| plt.close(fig) | |
| if not destination.is_file() or destination.stat().st_size == 0: | |
| raise RuntimeError("FFmpeg did not produce the video.") | |
| return destination | |