File size: 2,099 Bytes
3be620b |
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 |
import matplotlib.pyplot as plt
import numpy as np
from IPython.display import HTML
from matplotlib import animation
from ganime.visualization.images import unnormalize_if_necessary
def display_videos(data, ground_truth=None, n_rows=3, n_cols=3):
if ground_truth is not None:
data = np.concatenate((data, ground_truth), axis=2)
fig, axs = plt.subplots(nrows=n_rows, ncols=n_cols, squeeze=False, figsize=(16, 9))
# remove grid and ticks
plt.setp(axs, xticks=[], yticks=[])
plt.subplots_adjust(wspace=0, hspace=0)
ims = []
for i in range(n_rows):
for j in range(n_cols):
idx = i * n_cols + j
video = data[idx]
frame = unnormalize_if_necessary(video[0])
im = axs[i][j].imshow(frame, animated=True)
ims.append(im)
plt.close() # this is required to not display the generated image
def init():
for i in range(n_rows):
for j in range(n_cols):
idx = i * n_cols + j
video = data[idx]
im = ims[idx]
frame = unnormalize_if_necessary(video[0])
im.set_data(frame)
return ims
def animate(frame_id):
for i in range(n_rows):
for j in range(n_cols):
idx = i * n_cols + j
video = data[idx]
frame = video[frame_id, :, :, :]
frame = unnormalize_if_necessary(frame)
# if frame_id % 2 == 0:
# d[0:2, :, 0] = 255
# d[0:2, :, 1] = 0
# d[0:2, :, 2] = 0
# d[-2:, :, 0] = 255
# d[-2:, :, 1] = 0
# d[-2:, :, 2] = 0
ims[idx].set_data(frame)
return ims
anim = animation.FuncAnimation(
fig, animate, init_func=init, frames=data.shape[1], blit=True, interval=200
)
# FFwriter = animation.FFMpegWriter(fps=10, codec="libx264")
# anim.save("basic_animation1.mp4", writer=FFwriter)
return HTML(anim.to_html5_video())
|