diff --git a/setup.py b/setup.py index a55ab0485b..b6249d9a38 100644 --- a/setup.py +++ b/setup.py @@ -69,6 +69,7 @@ "sphinxcontrib-apidoc>=0.3.0", ], "extras": ["pynput>=1.7.4"], # Used by HumanKeyboardAgent + "gym": ["moviepy = 1.0.3"] "remote_agent": ["grpcio==1.32.0"], "rllib": [ "opencv-python==4.1.2.30", diff --git a/smarts/env/wrappers/gif_recorder.py b/smarts/env/wrappers/gif_recorder.py index ce71d1ba33..fab2fd2628 100644 --- a/smarts/env/wrappers/gif_recorder.py +++ b/smarts/env/wrappers/gif_recorder.py @@ -21,8 +21,8 @@ # THE SOFTWARE. import os -from moviepy.editor import * -import gym.envs +from moviepy.editor import ImageClip +from moviepy.editor import ImageSequenceClip import shutil import time from pathlib import Path @@ -30,6 +30,9 @@ class GifRecorder: def __init__(self, dir, env): + """ + Use images(rgb_array) to create a gif file. + """ timestamp_str = time.strftime("%Y%m%d-%H%M%S") self.dir = dir + "_" + timestamp_str self.env = env @@ -42,15 +45,24 @@ def __init__(self, dir, env): self._dir_name = str(Path(dir).name) def capture_frame(self, step_num, image): + """ + Create image according to the rgb_array and store it with step number in the destinated folder + """ with ImageClip(image) as image_clip: image_clip.save_frame(f"{self.dir}/{self._dir_name}_{step_num}.jpeg") def generate_gif(self): + """ + Use the images in the same folder to create a gif file. + """ with ImageSequenceClip(self.dir, fps=10) as clip: clip.write_gif(f"videos/{self._dir_name}.gif") clip.close() def close_recorder(self): + """ + close the recorder by deleting the image folder. + """ try: shutil.rmtree(self.dir) except: diff --git a/smarts/env/wrappers/recorder_wrapper.py b/smarts/env/wrappers/recorder_wrapper.py index cf78d7c6b3..b562c439ae 100644 --- a/smarts/env/wrappers/recorder_wrapper.py +++ b/smarts/env/wrappers/recorder_wrapper.py @@ -29,6 +29,9 @@ class RecorderWrapper(gym.Wrapper): + """ + A Wrapper that interacts the gym environment with the GifRecorder to record video step by step. + """ def __init__(self, dir, env): try: @@ -44,6 +47,9 @@ def __init__(self, dir, env): self.current_frame = -1 def reset(self, **kwargs): + """ + Reset the gym environment and restart recording. + """ observations = super().reset(**kwargs) if self.recording == False: self.start_recording() @@ -51,6 +57,9 @@ def reset(self, **kwargs): return observations def start_recording(self): + """ + Start the gif recorder and capture the first frame. + """ if self.gif_recorder is None: self.gif_recorder = GifRecorder(self.dir, self.env) image = super().render(mode="rgb_array") @@ -58,9 +67,15 @@ def start_recording(self): self.recording = True def stop_recording(self): + """ + Stop recording. + """ self.recording = False def step(self, action): + """ + Step the environment using the action and record the next frame. + """ observations, rewards, dones, infos = super().step(action) if self.recording == True: image = super().render(mode="rgb_array") @@ -69,10 +84,16 @@ def step(self, action): return observations, rewards, dones, infos def next_frame_id(self): + """ + Get the id for next frame. + """ self.current_frame += 1 return self.current_frame def close(self): + """ + Close the recorder by deleting the image folder and generate the gif file. + """ if self.gif_recorder is not None: self.gif_recorder.generate_gif() self.gif_recorder.close_recorder()