Spaces:
Running on Zero
Running on Zero
| """Skeleton -> MP4 renderer for the FlashMotion music-to-dance Space. | |
| The joint layout, the parent table, the bone connectivity and the foot-contact | |
| colouring are taken from the authors' ``vis.py`` | |
| (``Sun-Happy-YKX/FlashMotion-release``, MIT). Three deliberate deviations, all | |
| forced by the Space environment: | |
| * The authors shell out to a system ``ffmpeg`` (``os.system("ffmpeg ...")``). | |
| Spaces have no system ffmpeg, so the frames go through ``imageio_ffmpeg`` | |
| instead, which ships a bundled binary and muxes the music in the same call. | |
| * The authors render to a GIF and re-encode it. Frames are drawn straight into | |
| the MP4 encoder here -- same picture, one fewer lossy round trip. | |
| * The authors' camera is a fixed ``xcenter, ycenter, zcenter = 0, 0, 2.5`` with | |
| ``axrange = 3``, i.e. z in [1, 4] -- but FlashMotion puts the dancer's feet on | |
| a shared floor at z ~ 0, so that framing crops the legs. The camera here is | |
| fitted to the clip with a fixed pose, so nothing jitters and nothing is cut. | |
| The drawing itself is a small software rasteriser over Pillow (an orthographic | |
| camera plus a per-bone depth sort) rather than matplotlib's 3-D axes: a | |
| matplotlib 3-D frame costs ~25 ms, which would dominate the whole pipeline at | |
| 30 FPS. | |
| """ | |
| from __future__ import annotations | |
| import math | |
| import os | |
| import numpy as np | |
| from PIL import Image, ImageDraw | |
| # --- the authors' skeleton definition (vis.py, verbatim) -------------------- # | |
| SMPL_PARENTS = [ | |
| -1, 0, 0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 9, 9, 12, 13, 14, 16, 17, 18, 19, | |
| 20, 21, | |
| ] | |
| FOOT_JOINTS = (7, 8, 10, 11) # lankle, rankle, ltoes, rtoes | |
| NUM_JOINTS = len(SMPL_PARENTS) | |
| HEAD_JOINT = 15 | |
| NECK_JOINT = 12 | |
| # Bones are drawn as a slightly tapered skeleton: the torso/limbs read better | |
| # than a uniform stick figure at 512 px. | |
| BONES = [(child, parent) for child, parent in enumerate(SMPL_PARENTS) if parent >= 0] | |
| SIZE = 512 | |
| FPS = 30 | |
| BG_TOP = (18, 22, 33) | |
| BG_BOTTOM = (46, 52, 70) | |
| FLOOR_LINE = (72, 82, 108) | |
| BONE_COLOR = (238, 240, 248) | |
| SPINE_COLOR = (176, 196, 255) | |
| HEAD_COLOR = (255, 214, 140) | |
| FOOT_DOWN = (255, 96, 96) | |
| FOOT_UP = (86, 224, 140) | |
| SHADOW = (12, 14, 20) | |
| # Fixed camera: a three-quarter front view, folded down a little so the floor | |
| # line is visible. Angles are in degrees, measured in the model's Z-up frame. | |
| CAM_AZIMUTH = 35.0 | |
| CAM_ELEVATION = 14.0 | |
| CAM_DISTANCE = 6.0 | |
| CAM_TARGET_Z = 0.85 | |
| ORTHO_MARGIN = 1.30 | |
| def _background(size): | |
| img = Image.new("RGB", (size, size), BG_TOP) | |
| draw = ImageDraw.Draw(img) | |
| for y in range(size): | |
| t = y / max(size - 1, 1) | |
| color = tuple( | |
| int(round(BG_TOP[c] + (BG_BOTTOM[c] - BG_TOP[c]) * t)) for c in range(3) | |
| ) | |
| draw.line([(0, y), (size, y)], fill=color) | |
| return img | |
| def _camera_basis(): | |
| az = math.radians(CAM_AZIMUTH) | |
| el = math.radians(CAM_ELEVATION) | |
| offset = np.array( | |
| [ | |
| math.cos(el) * math.cos(az), | |
| math.cos(el) * math.sin(az), | |
| math.sin(el), | |
| ] | |
| ) | |
| eye = np.array([0.0, 0.0, CAM_TARGET_Z]) + offset * CAM_DISTANCE | |
| forward = np.array([0.0, 0.0, CAM_TARGET_Z]) - eye | |
| forward = forward / np.linalg.norm(forward) | |
| world_up = np.array([0.0, 0.0, 1.0]) | |
| right = np.cross(forward, world_up) | |
| right = right / np.linalg.norm(right) | |
| up = np.cross(right, forward) | |
| return eye, right, up, forward | |
| def _fit_scale(joints, size): | |
| """Orthographic half-width that just fits the clip with a little margin.""" | |
| eye, right, up, _ = _camera_basis() | |
| pts = joints.reshape(-1, 3) - eye | |
| u = pts @ right | |
| v = pts @ up | |
| half = max( | |
| np.abs(u - u.mean()).max(), np.abs(v - v.mean()).max(), 1e-6 | |
| ) | |
| return float(half) * ORTHO_MARGIN, float(u.mean()), float(v.mean()) | |
| def _project(points, eye, right, up, scale, cu, cv, size): | |
| pts = points - eye | |
| u = pts @ right | |
| v = pts @ up | |
| px = (u - cu) / scale * (size / 2.0) + size / 2.0 | |
| py = size / 2.0 - (v - cv) / scale * (size / 2.0) | |
| depth = pts @ _project.forward | |
| return px, py, depth | |
| _project.forward = np.array([0.0, 0.0, 1.0]) | |
| def _draw_floor(draw, eye, right, up, scale, cu, cv, size, floor_z, joints): | |
| """A few concentric floor lines for depth cueing.""" | |
| lo = joints.reshape(-1, 3).min(axis=0) | |
| hi = joints.reshape(-1, 3).max(axis=0) | |
| span = float(max(hi[0] - lo[0], hi[1] - lo[1], 1.0)) * 1.8 | |
| cx = float((lo[0] + hi[0]) / 2.0) | |
| cy = float((lo[1] + hi[1]) / 2.0) | |
| step = span / 6.0 | |
| for k in range(-3, 4): | |
| x = cx + k * step | |
| pts = np.array([[x, cy - span, floor_z], [x, cy + span, floor_z]]) | |
| px, py, _ = _project(pts, eye, right, up, scale, cu, cv, size) | |
| draw.line([(px[0], py[0]), (px[1], py[1])], fill=FLOOR_LINE, width=1) | |
| y = cy + k * step | |
| pts = np.array([[cx - span, y, floor_z], [cx + span, y, floor_z]]) | |
| px, py, _ = _project(pts, eye, right, up, scale, cu, cv, size) | |
| draw.line([(px[0], py[0]), (px[1], py[1])], fill=FLOOR_LINE, width=1) | |
| def _draw_skeleton(draw, pose, eye, right, up, scale, cu, cv, size, planted): | |
| px, py, depth = _project(pose, eye, right, up, scale, cu, cv, size) | |
| # soft contact shadow under the feet | |
| foot_px = px[list(FOOT_JOINTS)] | |
| foot_py = py[list(FOOT_JOINTS)] | |
| sx = float(np.mean(foot_px)) | |
| sy = float(np.mean(foot_py)) + 26.0 | |
| draw.ellipse( | |
| [sx - 62, sy - 12, sx + 62, sy + 12], fill=SHADOW | |
| ) | |
| bones = sorted(BONES, key=lambda bp: -0.5 * (depth[bp[0]] + depth[bp[1]])) | |
| half = max(2, int(round(size / 190))) | |
| for child, parent in bones: | |
| spine = parent in (0, 3, 6, 9, 12) and child not in FOOT_JOINTS | |
| color = SPINE_COLOR if spine else BONE_COLOR | |
| draw.line( | |
| [(px[child], py[child]), (px[parent], py[parent])], | |
| fill=color, | |
| width=half * 2, | |
| ) | |
| draw.ellipse( | |
| [px[child] - half, py[child] - half, px[child] + half, py[child] + half], | |
| fill=color, | |
| ) | |
| draw.ellipse( | |
| [px[parent] - half, py[parent] - half, px[parent] + half, py[parent] + half], | |
| fill=color, | |
| ) | |
| head_r = size / 22.0 | |
| hx, hy = float(px[HEAD_JOINT]), float(py[HEAD_JOINT]) | |
| nx, ny = float(px[NECK_JOINT]), float(py[NECK_JOINT]) | |
| draw.ellipse([hx - head_r, hy - head_r, hx + head_r, hy + head_r], fill=HEAD_COLOR) | |
| # foot contact markers: red = planted, green = in the air | |
| r = size / 60.0 | |
| for k, idx in enumerate(FOOT_JOINTS): | |
| color = FOOT_DOWN if planted[k] else FOOT_UP | |
| draw.ellipse( | |
| [px[idx] - r, py[idx] - r, px[idx] + r, py[idx] + r], fill=color | |
| ) | |
| def write_audio_clip(src_path, start, duration, out_wav): | |
| """Trim ``[start, start+duration]`` out of ``src_path`` at its native rate.""" | |
| import soundfile as sf | |
| with sf.SoundFile(src_path) as handle: | |
| sr = handle.samplerate | |
| handle.seek(int(round(start * sr))) | |
| data = handle.read(int(round(duration * sr)), dtype="float32", always_2d=True) | |
| peak = float(np.abs(data).max()) if data.size else 0.0 | |
| if peak > 1e-6: | |
| data = data / peak * 0.95 | |
| sf.write(out_wav, data, sr) | |
| return out_wav | |
| def render_mp4( | |
| joints, | |
| contact, | |
| audio_wav_path, | |
| out_path, | |
| fps=FPS, | |
| size=SIZE, | |
| progress=None, | |
| progress_offset=0.0, | |
| progress_span=1.0, | |
| ): | |
| """Render ``joints`` (T, 24, 3, Z-up metres) to an MP4 with music muxed in. | |
| ``contact`` is a (T, 4) array in [0, 1] for the four foot joints (from the | |
| decoder's contact head); a value > 0.95 paints that foot red (planted), | |
| otherwise green (in the air). | |
| """ | |
| import imageio_ffmpeg | |
| joints = np.asarray(joints, dtype=np.float64) | |
| frames = joints.shape[0] | |
| contact = np.asarray(contact, dtype=np.float64) | |
| if contact.ndim != 2 or contact.shape[1] != len(FOOT_JOINTS): | |
| contact = np.resize(contact, (frames, len(FOOT_JOINTS))) | |
| if contact.shape[0] != frames: | |
| idx = np.linspace(0, contact.shape[0] - 1, frames).round().astype(int) | |
| contact = contact[idx] | |
| planted = contact > 0.95 | |
| eye, right, up, forward = _camera_basis() | |
| _project.forward = forward | |
| scale, cu, cv = _fit_scale(joints, size) | |
| floor_z = float(joints[..., 2].min()) | |
| kwargs = dict(size=(size, size), fps=fps, quality=8, pix_fmt_in="rgb24") | |
| if audio_wav_path is not None and os.path.exists(audio_wav_path): | |
| kwargs["audio_path"] = audio_wav_path | |
| kwargs["audio_codec"] = "aac" | |
| kwargs["output_params"] = ["-shortest"] | |
| bg = _background(size) | |
| writer = imageio_ffmpeg.write_frames(out_path, **kwargs) | |
| writer.send(None) | |
| try: | |
| for num in range(frames): | |
| img = bg.copy() | |
| draw = ImageDraw.Draw(img) | |
| _draw_floor(draw, eye, right, up, scale, cu, cv, size, floor_z, joints[num]) | |
| _draw_skeleton( | |
| draw, joints[num], eye, right, up, scale, cu, cv, size, planted[num] | |
| ) | |
| writer.send(np.asarray(img)) | |
| if progress is not None and num % 15 == 0: | |
| progress( | |
| progress_offset + progress_span * (num / max(frames, 1)), | |
| desc="Rendering animation", | |
| ) | |
| finally: | |
| writer.close() | |
| return out_path | |