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()