Merge pull request #1499 from AisenGinn/develop

Video recording utilities

.github/workflows/ci-base-tests-linux.yml

@@ -39,7 +39,7 @@ jobs:
           . ${{env.venv_dir}}/bin/activate
           pip install --upgrade pip
           pip install --upgrade wheel
-          pip install -e .[camera-obs,opendrive,rllib,test,test-notebook,torch,train]
+          pip install -e .[camera-obs,opendrive,rllib,test,test-notebook,torch,train,gym]
       - name: Run smoke tests
         run: |
           . ${{env.venv_dir}}/bin/activate

CHANGELOG.md

@@ -11,6 +11,7 @@ Copy and pasting the git commit messages is __NOT__ enough.
 ## [Unreleased]
 ### Added
 - Added single vehicle `Trip` into type. 
+- Added new video record ultility using moviepy.
 ### Deprecated
 ### Changed
 ### Removed

requirements.txt

@@ -53,8 +53,7 @@ iniconfig==1.1.1
 jsonpatch==1.32
 jsonpointer==2.3
 jsonschema==4.17.0
-keras==2.10.0
-Keras-Preprocessing==1.1.2
+keras==2.11.0
 kiwisolver==1.4.4
 libclang==14.0.6
 lz4==4.0.2
@@ -114,14 +113,14 @@ six==1.16.0
 soupsieve==2.3.2.post1
 tableprint==0.9.1
 tabulate==0.9.0
-tensorboard==2.10.1
+tensorboard==2.11.0
 tensorboard-data-server==0.6.1
 tensorboard-plugin-wit==1.8.1
 tensorboardX==2.5.1
-tensorflow==2.10.1
-tensorflow-estimator==2.10.0
+tensorflow==2.11.0
+tensorflow-estimator==2.11.0
 tensorflow-io-gcs-filesystem==0.27.0
-termcolor==2.1.0
+termcolor==2.1.1
 tomli==2.0.1
 torch==1.4.0
 torchvision==0.5.0

smarts/env/wrappers/gif_recorder.py

@@ -0,0 +1,85 @@
+# 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 logging
+import sys
+
+import gym
+import numpy as np
+
+try:
+    from moviepy.editor import ImageClip, ImageSequenceClip
+except (ImportError, ModuleNotFoundError):
+    logging.warning(sys.exc_info())
+    logging.warning(
+        "You may not have installed the [gym] dependencies required to capture the video. Install them first with the `smarts[gym]` extras."
+    )
+
+    raise
+
+import shutil
+import time
+from pathlib import Path
+
+
+class GifRecorder:
+    """
+    Use images(rgb_array) to create a gif file.
+    """
+
+    def __init__(self, video_name_folder: str, env: gym.Env):
+        timestamp_str = time.strftime("%Y%m%d-%H%M%S")
+        self.frame_folder = (
+            video_name_folder + "_" + timestamp_str
+        )  # folder that uses to contain temporary frame images, will be deleted after the gif is created.
+        self.env = env
+
+        Path.mkdir(
+            Path(self.frame_folder), exist_ok=True
+        )  # create temporary frame images folder if not exists.
+
+        self._video_root_path = str(
+            Path(video_name_folder).parent
+        )  # path of the video file
+        self._video_name = str(Path(video_name_folder).name)  # name of the video
+
+    def capture_frame(self, step_num: int, image: np.ndarray):
+        """
+        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.frame_folder}/{self._video_name}_{step_num}.jpeg"
+            )
+
+    def generate_gif(self):
+        """
+        Use the images in the same folder to create a gif file.
+        """
+        with ImageSequenceClip(self.frame_folder, fps=10) as clip:
+            clip.write_gif(f"{self._video_root_path}/{self._video_name}.gif")
+        clip.close()
+
+    def close_recorder(self):
+        """
+        close the recorder by deleting the image folder.
+        """
+        shutil.rmtree(self.frame_folder, ignore_errors=True)

smarts/env/wrappers/recorder_wrapper.py

@@ -0,0 +1,110 @@
+# 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 typing
+from pathlib import Path
+
+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, video_name: str, env: gym.Env):
+
+        root_path = Path(__file__).parents[3]  # smarts main repo path
+        video_folder = os.path.join(
+            root_path, "videos"
+        )  # video folder for all video recording file (.gif)
+        Path.mkdir(
+            Path(video_folder), exist_ok=True
+        )  # create video folder if not exist
+
+        super().__init__(env)
+        self.video_name_folder = os.path.join(
+            video_folder, video_name
+        )  # frames folder that uses to contain temporary frame images, will be created using video name and current time stamp in gif_recorder when recording starts
+        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.video_name_folder, 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()