File size: 6,146 Bytes
8e5456b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 | #!/usr/bin/env python3
"""Render one TriVis sentence across every system, straight from the scored dumps.
Reads the same .npz dumps `eval_trivis_metrics.py` scores, so the picture is
guaranteed to show exactly what the numbers describe -- no second generation pass,
no risk of a render that disagrees with the table.
All dumps are 50 joints in TriVis frame coordinates: body 0:8 (OpenPose 0-7),
left hand 8:29, right hand 29:50. Drawn in pixels so proportions are correct.
"""
import argparse
import glob
import os
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.animation import FFMpegWriter, FuncAnimation
# OpenPose 0-7: 0 nose, 1 neck, 2 Rsho, 3 Relb, 4 Rwri, 5 Lsho, 6 Lelb, 7 Lwri
BODY_EDGES = [(1, 0), (1, 2), (2, 3), (3, 4), (1, 5), (5, 6), (6, 7)]
HAND_EDGES = [(0, 1), (1, 2), (2, 3), (3, 4), (0, 5), (5, 6), (6, 7), (7, 8), (0, 9),
(9, 10), (10, 11), (11, 12), (0, 13), (13, 14), (14, 15), (15, 16),
(0, 17), (17, 18), (18, 19), (19, 20), (5, 9), (9, 13), (13, 17)]
LH, RH = 8, 29
WRIST_L, WRIST_R = 7, 4 # body wrists, to attach each hand root
def draw(ax, xy, color, hand_color, W, H, lo, hi, lw=1.6):
ax.clear()
px, py = xy[:, 0] * W, xy[:, 1] * H
for a, b in BODY_EDGES:
ax.plot([px[a], px[b]], [py[a], py[b]], color=color, lw=lw, solid_capstyle='round')
for off, wr in ((LH, WRIST_L), (RH, WRIST_R)):
ax.plot([px[wr], px[off]], [py[wr], py[off]], color=color, lw=lw * 0.8)
for a, b in HAND_EDGES:
ax.plot([px[off + a], px[off + b]], [py[off + a], py[off + b]],
color=hand_color, lw=lw * 1.1, solid_capstyle='round')
ax.set_xlim(lo[0], hi[0]); ax.set_ylim(hi[1], lo[1])
ax.set_aspect('equal'); ax.axis('off')
def load(p):
d = np.load(p, allow_pickle=True)
return {str(n): q for n, q in zip(d['names'], d['poses'])}
def main():
ap = argparse.ArgumentParser()
ap.add_argument('--panels', nargs='+', required=True,
help='label=path/to.npz, in display order')
ap.add_argument('--clip', default=None, help='clip name (default: first shared)')
ap.add_argument('--frames', type=int, default=6)
ap.add_argument('--fps', type=int, default=30)
ap.add_argument('--frame-w', type=float, default=1176.0)
ap.add_argument('--frame-h', type=float, default=1288.0)
ap.add_argument('--gloss', default=None, help='caption text')
ap.add_argument('--out', default='qual_trivis/strip')
ap.add_argument('--mp4', action='store_true')
ap.add_argument('--gif', action='store_true',
help='also write a gif (ffmpeg palettegen: far smaller and cleaner '
'than a direct matplotlib gif)')
ap.add_argument('--gif-fps', type=int, default=15)
ap.add_argument('--gif-width', type=int, default=900)
args = ap.parse_args()
panels = []
for spec in args.panels:
lab, path = spec.split('=', 1)
panels.append((lab, load(path)))
shared = set(panels[0][1])
for _, m in panels[1:]:
shared &= set(m)
if not shared:
raise SystemExit('no clip shared by all panels')
clip = args.clip or sorted(shared)[0]
if clip not in shared:
raise SystemExit(f'{clip} not in all panels')
print(f'clip {clip} ({len(shared)} shared)')
seqs = [(lab, np.asarray(m[clip], np.float32)) for lab, m in panels]
W, H = args.frame_w, args.frame_h
ref = seqs[0][1]
lo = (ref.reshape(-1, 2).min(0) * np.array([W, H])) - 60
hi = (ref.reshape(-1, 2).max(0) * np.array([W, H])) + 60
cols = [('#111111', '#c0392b'), ('#1f6f3f', '#27ae60'), ('#1a4f8a', '#2980b9'),
('#8e44ad', '#c39bd3'), ('#b9770e', '#e59866'), ('#117a65', '#45b39d')]
os.makedirs(os.path.dirname(args.out) or '.', exist_ok=True)
title = f"{args.gloss or clip}"
# ---- keyframe grid: one row per system ----
nf = args.frames
fig, axes = plt.subplots(len(seqs), nf, figsize=(1.55 * nf, 1.9 * len(seqs)))
if len(seqs) == 1:
axes = axes[None, :]
fig.suptitle(title, fontsize=9)
for r, (lab, seq) in enumerate(seqs):
ts = np.linspace(0, len(seq) - 1, nf).astype(int)
c, hc = cols[r % len(cols)]
for ci, t in enumerate(ts):
draw(axes[r, ci], seq[t], c, hc, W, H, lo, hi, lw=1.2)
if ci == 0:
axes[r, ci].set_ylabel(lab, fontsize=7)
axes[r, ci].set_title(f'{t+1}', fontsize=6)
plt.tight_layout(rect=[0, 0, 1, 0.94])
png = args.out + '.png'
fig.savefig(png, dpi=140)
plt.close(fig)
print(f'wrote {png}')
if args.mp4 or args.gif:
nT = max(len(s) for _, s in seqs)
fig, axes = plt.subplots(1, len(seqs), figsize=(2.1 * len(seqs), 3.4))
fig.suptitle(title, fontsize=9)
def frame(t):
for k, (lab, seq) in enumerate(seqs):
tt = min(t, len(seq) - 1)
c, hc = cols[k % len(cols)]
draw(axes[k], seq[tt], c, hc, W, H, lo, hi)
axes[k].set_title(f'{lab}\n{tt+1}/{len(seq)}', fontsize=7)
return []
anim = FuncAnimation(fig, frame, frames=nT, interval=1000 / args.fps)
mp4 = args.out + '.mp4'
anim.save(mp4, writer=FFMpegWriter(fps=args.fps, bitrate=2400))
plt.close(fig)
print(f'wrote {mp4}')
if args.gif:
import subprocess
gif = args.out + '.gif'
pal = args.out + '.pal.png'
vf = f'fps={args.gif_fps},scale={args.gif_width}:-1:flags=lanczos'
subprocess.run(['ffmpeg', '-y', '-loglevel', 'error', '-i', mp4,
'-vf', vf + ',palettegen=stats_mode=diff', pal], check=True)
subprocess.run(['ffmpeg', '-y', '-loglevel', 'error', '-i', mp4, '-i', pal,
'-lavfi', vf + '[x];[x][1:v]paletteuse=dither=bayer:bayer_scale=3',
gif], check=True)
os.remove(pal)
print(f'wrote {gif} ({os.path.getsize(gif)/1e6:.2f} MB)')
if __name__ == '__main__':
main()
|