Video recording utilities

requirements.txt

@@ -7,7 +7,7 @@ astunparse==1.6.3
 async-timeout==4.0.2
 asynctest==0.13.0
 atari-py==0.2.6
-attrs==21.4.0
+attrs==22.1.0
 Automat==20.2.0
 beautifulsoup4==4.11.1
 blessings==1.7
@@ -21,20 +21,20 @@ colorama==0.4.5
 colorful==0.5.4
 commonmark==0.9.1
 constantly==15.1.0
-coverage==6.4.2
+coverage==6.4.4
 cycler==0.11.0
 dm-tree==0.1.7
 eclipse-sumo==1.10.0
 execnet==1.9.0
-filelock==3.7.1
+filelock==3.8.0
 flatbuffers==1.12
-fonttools==4.34.4
-frozenlist==1.3.0
+fonttools==4.35.0
+frozenlist==1.3.1
 future==0.18.2
 gast==0.4.0
 google==3.0.0
 google-api-core==2.8.2
-google-auth==2.9.1
+google-auth==2.10.0
 google-auth-oauthlib==0.4.6
 google-pasta==0.2.0
 googleapis-common-protos==1.56.4
@@ -51,21 +51,22 @@ incremental==21.3.0
 iniconfig==1.1.1
 jsonpatch==1.32
 jsonpointer==2.3
-jsonschema==4.7.2
+jsonschema==4.10.0
 keras==2.9.0
 Keras-Preprocessing==1.1.2
 kiwisolver==1.4.4
-libclang==14.0.1
-lz4==4.0.1
+libclang==14.0.6
+lz4==4.0.2
 Markdown==3.4.1
-matplotlib==3.5.2
+MarkupSafe==2.1.1
+matplotlib==3.5.3
 msgpack==1.0.4
 multidict==6.0.2
 numpy==1.21.6
 nvidia-ml-py3==7.352.0
 oauthlib==3.2.0
-opencensus==0.10.0
-opencensus-context==0.1.2
+opencensus==0.11.0
+opencensus-context==0.1.3
 opencv-python==4.1.2.30
 opencv-python-headless==4.1.2.30
 opt-einsum==3.3.0
@@ -75,6 +76,7 @@ panda3d-gltf==0.13
 panda3d-simplepbr==0.10
 pandas==1.3.5
 Pillow==9.2.0
+pkgutil_resolve_name==1.3.10
 pluggy==1.0.0
 prometheus-client==0.14.1
 protobuf==3.19.4
@@ -85,7 +87,7 @@ py-spy==0.3.12
 pyasn1==0.4.8
 pyasn1-modules==0.2.8
 pybullet==3.0.6
-Pygments==2.12.0
+Pygments==2.13.0
 pyparsing==3.0.9
 pyrsistent==0.18.1
 pytest==7.1.2
@@ -94,9 +96,9 @@ pytest-cov==3.0.0
 pytest-forked==1.4.0
 pytest-xdist==2.5.0
 python-dateutil==2.8.2
-pytz==2022.1
+pytz==2022.2.1
 PyYAML==6.0
-pyzmq==23.2.0
+pyzmq==23.2.1
 ray==1.0.1.post1
 redis==3.4.1
 requests==2.28.1
@@ -127,13 +129,13 @@ tornado==6.2
 trimesh==3.9.29
 Twisted==22.4.0
 typing_extensions==4.3.0
-urllib3==1.26.10
+urllib3==1.26.11
 visdom==0.1.8.9
 wcwidth==0.2.5
 websocket-client==1.3.3
-Werkzeug==2.1.2
+Werkzeug==2.2.2
 wrapt==1.14.1
-yarl==1.7.2
+yarl==1.8.1
 yattag==1.14.0
 zipp==3.8.1
 zope.interface==5.4.0

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,76 @@
+# 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, 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,104 @@
+# 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()