From f1034b82c5f485ae678d061eeb45556a40971f2d Mon Sep 17 00:00:00 2001
From: Julien Perez <perezjln@users.noreply.github.com>
Date: Mon, 20 May 2024 23:01:26 +0200
Subject: [PATCH] add argparse

---
 examples/trace_hdf5_img.py    | 36 +++++++++++++++++++----------------
 examples/trace_hdf5_mujoco.py | 11 +++++++----
 2 files changed, 27 insertions(+), 20 deletions(-)

diff --git a/examples/trace_hdf5_img.py b/examples/trace_hdf5_img.py
index 6fae6a0..0a2f225 100644
--- a/examples/trace_hdf5_img.py
+++ b/examples/trace_hdf5_img.py
@@ -1,28 +1,32 @@
 
 import h5py, tqdm
+import argparse
 import matplotlib.pyplot as plt
 import matplotlib.cm as cm
 import matplotlib.animation as animation
 
+if __name__ == "__main__":
 
-# Specify the path to your HDF5 file
-file_path = 'data/episode_46.hdf5'
+    # Parse command line arguments
+    parser = argparse.ArgumentParser(description='Trace video from HDF5 trace file')
+    parser.add_argument('--file_path', type=str, default='data/episode_46.hdf5', help='Path to HDF5 file')
+    parser.add_argument('--camera_id', type=str, default='observations/images/front', help='Camera ID')
+    parser.add_argument('--save_file', type=str, default='movie.mp4', help='Path to save the video file')
+    args = parser.parse_args()
 
-# Open the HDF5 file
-with h5py.File(file_path, 'r') as file:
+    # Access the specified HDF5 file
+    with h5py.File(args.file_path, 'r') as file:
 
-    # Access the datasets or attributes in the file
-    # For example, to access a dataset named 'data':
-    group_obs_img = file['observations/images/front']
+        # Access the datasets or attributes in the file
+        group_obs_img = file[args.camera_id]
 
-    img = [] # some array of images
-    frames = [] # for storing the generated images
+        img = [] # some array of images
+        frames = [] # for storing the generated images
 
-    # Write each image to the video
-    fig = plt.figure()
-    for i in tqdm.tqdm(range(group_obs_img.shape[0]), desc='Writing video'):
-        frames.append([plt.imshow(group_obs_img[i],animated=True)])
+        # Write each image to the video
+        fig = plt.figure()
+        for i in tqdm.tqdm(range(group_obs_img.shape[0]), desc='Writing video'):
+            frames.append([plt.imshow(group_obs_img[i],animated=True)])
 
-    ani = animation.ArtistAnimation(fig, frames, interval=50, blit=True,
-                                    repeat_delay=1000)
-    ani.save('movie.mp4')
+        ani = animation.ArtistAnimation(fig, frames, interval=50, blit=True, repeat_delay=1000)
+        ani.save(args.save_file)
diff --git a/examples/trace_hdf5_mujoco.py b/examples/trace_hdf5_mujoco.py
index 8f63717..2c55bf1 100644
--- a/examples/trace_hdf5_mujoco.py
+++ b/examples/trace_hdf5_mujoco.py
@@ -1,5 +1,6 @@
 import h5py
 import time
+import argparse
 
 import mujoco
 import mujoco.viewer
@@ -7,11 +8,10 @@
 
 from gym_lowcostrobot.simulated_robot import SimulatedRobot
 
-def do_replay_hdf5():
+def do_replay_hdf5(args):
 
     # Specify the path to your HDF5 file
-    file_path = 'data/episode_46.hdf5'
-    with h5py.File(file_path, 'r') as file:
+    with h5py.File(args.file_path, 'r') as file:
 
         m = mujoco.MjModel.from_xml_path("assets/scene_one_cube.xml")
         data = mujoco.MjData(m)
@@ -43,4 +43,7 @@ def do_replay_hdf5():
 
 
 if __name__ == "__main__":
-    do_replay_hdf5()
+    parser = argparse.ArgumentParser(description='Trace video from HDF5 trace file')
+    parser.add_argument('--file_path', type=str, default='data/episode_46.hdf5', help='Path to HDF5 file')
+    args = parser.parse_args()
+    do_replay_hdf5(args)