| 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) | |