import matplotlib

matplotlib.use('Qt4Agg');
import matplotlib.pyplot as plt
import numpy as np
import imageio
from userlib import dmd
from skimage.color import rgb2gray
from skimage.transform import downscale_local_mean
import matplotlib.animation as animation
# we want data in the form
# time ------>
# frame
# |
# |
# V
#print(data_gray.shape)

# load video
vid = imageio.get_reader('tripletap.mp4')
# Get video infromation
height= vid.get_data(0).shape[0]
width = vid.get_data(0).shape[1]
numFrames = len(vid)
print()
print("Video Information:")
print("Total frames = ", numFrames)
print("Frame size = ", width, " x ", height)
print()
# [frame, height, width, RGB]
video = np.array([frame for frame in vid])
# [row, col]
data_gray = np.zeros([width*height//100, numFrames])
for n in range(numFrames):
data = np.array(video[n,:,:,:])
img_gray = downscale_local_mean(rgb2gray(data), (10, 10))
data_gray[:, n] = img_gray.flatten();
# Plot initial
#plt.imshow(data_gray[:,18].reshape(height//10, width//10))
#plt.show()
# Perform dmd
threshold = 1
Xbg = dmd(data_gray, 1, threshold)
print("data_gray = ", data_gray.shape)
print("Xbg = ", Xbg.shape)
# Finding residual R
Xfg = data_gray - abs(Xbg)
R = np.copy(Xfg)
R[R > 0 ] = 0
Xfg[Xfg < 0] = 0
Xbg = abs(Xbg) + R
# Print video
fig = plt.figure()
ims = []
for n in range(numFrames):
im = plt.imshow(Xfg[:, n].reshape(height//10, width//10), animated=True)
ims.append([im])
ani = animation.ArtistAnimation(fig, ims, interval=50, blit=True,
repeat_delay=1000)
plt.show()
ani.save("foreground.mp4")
fig = plt.figure()
ims2 = []
for n in range(numFrames):
im = plt.imshow(Xbg[:, n].reshape(height//10, width//10), animated=True)
ims2.append([im])
ani = animation.ArtistAnimation(fig, ims2, interval=50, blit=True,
repeat_delay=1000)
plt.show()
ani.save("background.mp4")