xref: /aosp_15_r20/external/libopus/dnn/torch/osce/stndrd/presentation/playback.py (revision a58d3d2adb790c104798cd88c8a3aff4fa8b82cc)
1import matplotlib
2import matplotlib.pyplot as plt
3import matplotlib.animation
4
5def make_playback_animation(savepath, spec, duration_ms, vmin=20, vmax=90):
6    fig, axs = plt.subplots()
7    axs.set_axis_off()
8    fig.set_size_inches((duration_ms / 1000 * 5, 5))
9    frames = []
10    frame_duration=20
11    num_frames = int(duration_ms / frame_duration + .99)
12
13    spec_height, spec_width = spec.shape
14    for i in range(num_frames):
15        xpos = (i - 1) / (num_frames - 3) * (spec_width - 1)
16        new_frame = axs.imshow(spec, cmap='inferno', origin='lower', aspect='auto', vmin=vmin, vmax=vmax)
17        if i in {0, num_frames - 1}:
18            frames.append([new_frame])
19        else:
20            line = axs.plot([xpos, xpos], [0, spec_height-1], color='white', alpha=0.8)[0]
21            frames.append([new_frame, line])
22
23
24    ani = matplotlib.animation.ArtistAnimation(fig, frames, blit=True, interval=frame_duration)
25    ani.save(savepath, dpi=720)