Video recording utilities

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",

smarts/env/wrappers/gif_recorder.py

@@ -0,0 +1,77 @@
+# MIT License
+#
+# Copyright (C) 2022. Huawei Technologies Co., Ltd. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NON-INFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+# THE SOFTWARE.
+import os
+import sys
+
+try:
+    from moviepy.editor import ImageClip
+    from moviepy.editor import ImageSequenceClip
+except (ImportError, ModuleNotFoundError):
+    print(sys.exc_info())
+    print(
+        "You may not have installed the [gym] dependencies required to capture the video. Install them first using the command `pip install -e .[gym]` at the source directory."
+    )
+import shutil
+import time
+from pathlib import Path
+
+
+class GifRecorder:
+    """
+    Use images(rgb_array) to create a gif file.
+    """
+
+    def __init__(self, dir, env):
+        timestamp_str = time.strftime("%Y%m%d-%H%M%S")
+        self.dir = dir + "_" + timestamp_str
+        self.env = env
+
+        try:
+            os.mkdir(self.dir)
+        except:
+            pass
+
+        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:
+            pass

smarts/env/wrappers/recorder_wrapper.py

@@ -0,0 +1,103 @@
+# MIT License
+#
+# Copyright (C) 2022. Huawei Technologies Co., Ltd. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NON-INFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+# THE SOFTWARE.
+import os
+import gym
+import gym.envs
+
+from smarts.env.wrappers.gif_recorder import GifRecorder
+
+
+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:
+            os.mkdir("videos")
+        except:
+            pass
+
+        super().__init__(env)
+        # assert "rgb_array" in env.metadata.get("render_modes", [])
+        self.dir = "videos/" + dir
+        self.gif_recorder = None
+        self.recording = False
+        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()
+
+        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")
+        self.gif_recorder.capture_frame(self.next_frame_id(), image)
+        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")
+            self.gif_recorder.capture_frame(self.next_frame_id(), image)
+
+        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()
+            self.gif_recorder = None
+            self.recording = False
+
+    def __del__(self):
+        self.close()