| import numpy as np | |
| import imageio | |
| def load_video(filename): | |
| vid = imageio.get_reader(filename, 'ffmpeg') | |
| fps = vid.get_meta_data()['fps'] | |
| nframes = vid.count_frames() | |
| return vid, fps, nframes | |
| class SaveVideo: | |
| def __init__(self, outname, fps): | |
| self.outname = outname | |
| self.fps = fps | |
| def __enter__(self): | |
| self.writter = imageio.get_writer(self.outname, | |
| format='FFMPEG', | |
| fps=self.fps) | |
| return self | |
| def __exit__(self, exc_type, exc_value, exc_traceback): | |
| self.writter.close() | |
| def __iadd__(self, data): | |
| if np.max(data) <= 1: | |
| data = np.array(255*data, dtype=np.uint8) | |
| else: | |
| data = np.array(data, dtype=np.uint8) | |
| self.writter.append_data(data) | |
| return self | |