optimize code

app.py

@@ -6,6 +6,7 @@ import gdown
 import base64
 from time import gmtime, strftime
 from csv import writer
+import json
 
 from datasets import load_dataset
 from hfserver import HuggingFaceDatasetSaver, HuggingFaceDatasetJSONSaver
@@ -45,6 +46,7 @@ if LOAD_DATA_GOOGLE_DRIVE:  # download data from google drive
 else:  # local data
     VIDEO_PATH = 'robotinder-data'
 
+VIDEO_INFO = os.path.join(VIDEO_PATH, 'video_info.json')
 
 def inference(video_path):
     with open(video_path, "rb") as f:
@@ -123,15 +125,8 @@ def update(user_choice, left, right, choose_env, data_folder=VIDEO_PATH, flag_to
     else:
         env_name = choose_env
     # choose video
-    videos = os.listdir(os.path.join(data_folder, env_name))
-    video_files = []
-    for f in videos:
-        if f.endswith(f'.{FORMAT}'):
-            video_files.append(os.path.join(data_folder, env_name, f))
-    # randomly choose two videos
-    selected_video_ids = np.random.choice(len(video_files), 2, replace=False)
-    left = video_files[selected_video_ids[0]]
-    right = video_files[selected_video_ids[1]]
+    left, right = randomly_select_videos(env_name)
+
     last_left_video_path = left
     last_right_video_path = right
     last_infer_left_video_path = inference(left)
@@ -143,11 +138,41 @@ def replay(left, right):
     return left, right
 
 def parse_envs(folder=VIDEO_PATH):
-    envs = []
-    for f in os.listdir(folder):
-        if os.path.isdir(os.path.join(folder, f)):
-            envs.append(f)
-    return envs
+    """
+    return a dict of env_name: video_paths
+    """
+    files = {}
+    for env_name in os.listdir(folder):
+        env_path = os.path.join(folder, env_name)
+        if os.path.isdir(env_path):
+            videos = os.listdir(env_path)
+            video_files = []
+            for video in videos:
+                if video.endswith(f'.{FORMAT}'):
+                    video_path = os.path.join(folder, env_name, video)
+                    video_files.append(video_path)
+            files[env_name] = video_files
+
+    with open(VIDEO_INFO, 'w') as fp:
+        json.dump(files, fp)
+
+    return files
+
+def get_env_names():
+    with open(VIDEO_INFO, 'r') as fp:
+        files = json.load(fp)
+    return list(files.keys())
+
+def randomly_select_videos(env_name):
+    # load the parsed video info
+    with open(VIDEO_INFO, 'r') as fp:
+        files = json.load(fp)
+    env_files = files[env_name]
+    # randomly choose two videos
+    selected_video_ids = np.random.choice(len(env_files), 2, replace=False)
+    left_video_path = env_files[selected_video_ids[0]]
+    right_video_path = env_files[selected_video_ids[1]]
+    return left_video_path, right_video_path
 
 def build_interface(iter=3, data_folder=VIDEO_PATH):
     import sys
@@ -162,26 +187,20 @@ def build_interface(iter=3, data_folder=VIDEO_PATH):
     # callback = gr.CSVLogger()
     callback = hf_writer
 
+    # parse the video folder 
+    files = parse_envs()   
+    
     # build gradio interface
     with gr.Blocks() as demo:
         gr.Markdown("## Here is <span style=color:cyan>RoboTinder</span>!")
         gr.Markdown("### Select the best robot behaviour in your choice!")
         # some initial values
-        envs = parse_envs()   
-        env_name = envs[random.randint(0, len(envs)-1)] # random pick an env 
+        env_name = list(files.keys())[random.randint(0, len(files)-1)] # random pick an env 
         with gr.Row():
             str_env_name = gr.Markdown(f"{env_name}")
 
         # choose video
-        videos = os.listdir(os.path.join(data_folder, env_name))
-        video_files = []
-        for f in videos:
-            if f.endswith(f'.{FORMAT}'):
-                video_files.append(os.path.join(data_folder, env_name, f))
-        # randomly choose two videos
-        selected_video_ids = np.random.choice(len(video_files), 2, replace=False)
-        left_video_path = video_files[selected_video_ids[0]]
-        right_video_path = video_files[selected_video_ids[1]]
+        left_video_path, right_video_path = randomly_select_videos(env_name)
         
         with gr.Row():
             if FORMAT == 'mp4':
@@ -219,7 +238,7 @@ def build_interface(iter=3, data_folder=VIDEO_PATH):
         btn2.click(fn=update, inputs=[user_choice, left, right, choose_env], outputs=[left, right, str_env_name])
 
         # We can choose which components to flag -- in this case, we'll flag all of them
-        btn2.click(lambda *args: callback.flag(args), [user_choice, left, right], None, preprocess=False)
+        # btn2.click(lambda *args: callback.flag(args), [user_choice, left, right], None, preprocess=False)  # not using the gradio flagging anymore
 
     return demo