"""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