diff --git a/.github/workflows/ci-base-tests-linux.yml b/.github/workflows/ci-base-tests-linux.yml index f5b831a899..29d01f53f3 100644 --- a/.github/workflows/ci-base-tests-linux.yml +++ b/.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 diff --git a/CHANGELOG.md b/CHANGELOG.md index d323f76a8d..9ecbe6b9f9 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/requirements.txt b/requirements.txt index 34e5507080..f22b314408 100644 --- a/requirements.txt +++ b/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 diff --git a/smarts/env/wrappers/gif_recorder.py b/smarts/env/wrappers/gif_recorder.py new file mode 100644 index 0000000000..a39e0b3c70 --- /dev/null +++ b/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) diff --git a/smarts/env/wrappers/recorder_wrapper.py b/smarts/env/wrappers/recorder_wrapper.py new file mode 100644 index 0000000000..10b4eab5c1 --- /dev/null +++ b/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()