smog / src /utils /video.py
vonexel's picture
add: src
fe64bad verified
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