def test_patch(wandb_init_run, git_repo):
wandb.gym.monitor()
with open("test.gif", "w") as f:
f.write("test")
ir = ImageEncoder("test.gif", (28, 28, 3), 10)
ir.close()
assert wandb_init_run.summary["videos"]['_type'] == "video-file"
class VideoRecorder():
"""
Use ImageEncoder gym.wrappers.monitoring.video_recorder because it make really nice videos using .mp4 and ffmpeg
"""
def start(self, output_file, observation_shape, fps=30):
"""
:param output_file:
:param observation_shape:
:param fps:
:return:
"""
height = observation_shape[2]
width = observation_shape[1]
pixFmt = observation_shape[0]
frame_shape = (height, width, pixFmt)
self._image_encoder = ImageEncoder(output_file, frame_shape, fps, fps)
def add_frame(self, observation):
"""
:param observation:
:return:
"""
self._image_encoder.capture_frame(observation.swapaxes(0, 2))
def close(self):
self._image_encoder.close()
def __del__(self):
# Release everything if job is finished
self.close()
def record_video(buffer, file_path, fps):
"""
Records the given image buffer to a video file.
:param array_like buffer: the list containing the sequential image frames to be recorded.
:param str file_path: the path to the video file.
:param int fps: the video frames-per-second.
:return:
"""
# creates video encoder and adds each frame in the buffer
video_recorder = ImageEncoder(file_path, buffer[0].shape, fps)
for frame in buffer:
video_recorder.capture_frame(frame)
video_recorder.close()
class EpisodeRecorder(gym.Wrapper):
"""Wrapper for recording episodes to MP4 or jpeg-frames"""
def __init__(self, env, path, fps=60, header=None):
"""Initialize the wrapper
Args:
env: The env to wrap and record
path: The location where to place recordings
fps: The FPS to encode the MP4 at
header: Optional prefix for the output file name
"""
super(EpisodeRecorder, self).__init__(env)
self.fps = fps
self.path = path
self.header = header
self.encoder = None
os.makedirs(path, exist_ok=True)
def _get_output_name(self):
"""Gets a unique name for the output video, starting from 00000.mp4
(Handles collisions with parallel ENVs)
"""
base_name = self.header + "_" if self.header is not None else ""
i = 0
while True:
try:
path = os.path.join(self.path,
f"{base_name}{str(i).zfill(5)}.mp4")
pathlib.Path(path).touch(exist_ok=False)
return path
except FileExistsError:
i += 1
continue
def _init_new(self, shape):
self._close_last()
self.cur_path = self._get_output_name()
self.encoder = ImageEncoder(self.cur_path, shape, self.fps)
def _dump_frame(self, frame):
self.encoder.capture_frame(frame)
def _close_last(self):
if self.encoder is not None:
self.encoder.close()
self.encoder = None
def close(self):
self._close_last()
self.env.close()
def reset(self):
obs = self.env.reset()
# Create a new output
self._init_new(obs.shape)
self._dump_frame(obs)
return obs
def step(self, action):
obs, reward, done, info = self.env.step(action)
self._dump_frame(obs)
return obs, reward, done, info