add header to gif_recorder.py and recorder_wrapper.py.

smarts/env/wrappers/gif_recorder.py

@@ -1,36 +1,53 @@
+# 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
 
 from moviepy.editor import *
 import gym.envs
 import shutil
+import time
+from pathlib import Path
 
 
 class GifRecorder:
     def __init__(self, dir, env):
-        self.dir = dir
+        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 = None
-        if "/" not in dir:
-            self.dir_name = dir
-        else:
-            last_index = dir.rindex("/")
-            self.dir_name = dir[-(len(dir) - last_index - 1) :]
-
-    def capture_frame(self, step_num):
-        image = self.env.render(mode="rgb_array")
+        self._dir_name = str(Path(dir).name)
 
+    def capture_frame(self, step_num, image):
         with ImageClip(image) as image_clip:
-            image_clip.save_frame(f"{self.dir}/{self.dir_name}_{step_num}.jpeg")
+            image_clip.save_frame(f"{self.dir}/{self._dir_name}_{step_num}.jpeg")
 
     def generate_gif(self):
         with ImageSequenceClip(self.dir, fps=10) as clip:
-            clip.write_gif(f"videos/{self.dir_name}.gif")
+            clip.write_gif(f"videos/{self._dir_name}.gif")
         clip.close()
 
     def close_recorder(self):

smarts/env/wrappers/recorder_wrapper.py

@@ -1,3 +1,24 @@
+# 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
 
 from moviepy.editor import *
@@ -32,7 +53,8 @@ def reset(self, **kwargs):
     def start_recording(self):
         if self.gif_recorder is None:
             self.gif_recorder = GifRecorder(self.dir, self.env)
-        self.gif_recorder.capture_frame(self.next_frame_id())
+        image = super().render(mode="rgb_array")
+        self.gif_recorder.capture_frame(self.next_frame_id(),image)
         self.recording = True
 
     def stop_recording(self):
@@ -41,7 +63,8 @@ def stop_recording(self):
     def step(self, action):
         observations, rewards, dones, infos = super().step(action)
         if self.recording == True:
-            self.gif_recorder.capture_frame(self.next_frame_id())
+            image = super().render(mode="rgb_array")
+            self.gif_recorder.capture_frame(self.next_frame_id(),image)
 
         return observations, rewards, dones, infos
 
@@ -50,11 +73,11 @@ def next_frame_id(self):
         return self.current_frame
 
     def close(self):
-        self.gif_recorder.close_recorder()
+        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.gif_recorder.close_recorder()
-
-    def close_recorder(self):
-        self.gif_recorder.generate_gif()
-        self.gif_recorder.close_recorder()
+        self.close()