File size: 644 Bytes
393d3de |
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 |
import imageio
import os
class VideoRecorder(object):
def __init__(self, dir_name, height=256, width=256, fps=30):
self.dir_name = dir_name
self.height = height
self.width = width
self.fps = fps
self.frames = []
def init(self, enabled=True):
self.frames = []
self.enabled = self.dir_name is not None and enabled
def record(self, obs):
if self.enabled:
self.frames.append(obs)
def save(self, file_name):
if self.enabled:
path = os.path.join(self.dir_name, file_name)
imageio.mimsave(path, self.frames, fps=self.fps)
|