From ab9a08df14f9d52153fda56e5bf39c8ee62f6d79 Mon Sep 17 00:00:00 2001
From: Tucker <montgomery.alban@huawei.com>
Date: Fri, 22 Jul 2022 19:13:04 +0000
Subject: [PATCH] add docstring the gif_recorder.py and recorder_wrapper.py.

---
 setup.py                                |  1 +
 smarts/env/wrappers/gif_recorder.py     | 16 ++++++++++++++--
 smarts/env/wrappers/recorder_wrapper.py | 21 +++++++++++++++++++++
 3 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/setup.py b/setup.py
index a55ab0485b..b6249d9a38 100644
--- a/setup.py
+++ b/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",
diff --git a/smarts/env/wrappers/gif_recorder.py b/smarts/env/wrappers/gif_recorder.py
index ce71d1ba33..fab2fd2628 100644
--- a/smarts/env/wrappers/gif_recorder.py
+++ b/smarts/env/wrappers/gif_recorder.py
@@ -21,8 +21,8 @@
 # THE SOFTWARE.
 import os
 
-from moviepy.editor import *
-import gym.envs
+from moviepy.editor import ImageClip
+from moviepy.editor import ImageSequenceClip
 import shutil
 import time
 from pathlib import Path
@@ -30,6 +30,9 @@
 
 class GifRecorder:
     def __init__(self, dir, env):
+        """
+        Use images(rgb_array) to create a gif file.
+        """
         timestamp_str = time.strftime("%Y%m%d-%H%M%S")
         self.dir = dir + "_" + timestamp_str
         self.env = env
@@ -42,15 +45,24 @@ def __init__(self, dir, env):
         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:
diff --git a/smarts/env/wrappers/recorder_wrapper.py b/smarts/env/wrappers/recorder_wrapper.py
index cf78d7c6b3..b562c439ae 100644
--- a/smarts/env/wrappers/recorder_wrapper.py
+++ b/smarts/env/wrappers/recorder_wrapper.py
@@ -29,6 +29,9 @@
 
 
 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:
@@ -44,6 +47,9 @@ def __init__(self, dir, env):
         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()
@@ -51,6 +57,9 @@ def reset(self, **kwargs):
         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")
@@ -58,9 +67,15 @@ def start_recording(self):
         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")
@@ -69,10 +84,16 @@ def step(self, action):
         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()