Update app.py

app.py

@@ -49,6 +49,7 @@ else:  # local data
 VIDEO_INFO = os.path.join(VIDEO_PATH, 'video_info.json')
 
 def inference(video_path):
+    # for displaying mp4 with autoplay on Gradio
     with open(video_path, "rb") as f:
         data = f.read()
         b64 = base64.b64encode(data).decode()
@@ -137,20 +138,49 @@ def update(user_choice, left, right, choose_env, data_folder=VIDEO_PATH, flag_to
 def replay(left, right):  
     return left, right
 
-def parse_envs(folder=VIDEO_PATH):
+def parse_envs(folder=VIDEO_PATH, filter=True, MAX_ITER=20000, DEFAULT_ITER=20000):
     """
     return a dict of env_name: video_paths
     """
     files = {}
+    if filter:
+        df = pd.read_csv('Bidexhands_Video.csv')
+        # print(df)
     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:
+            for video in videos:  # video name rule: EnvName_Alg_Seed_Timestamp_Checkpoint_video-episode_EpisodeID
                 if video.endswith(f'.{FORMAT}'):
+                    if filter:
+                        seed = video.split('_')[2]
+                        checkpoint = video.split('_')[4]
+                        try:
+                            succeed_iteration = df.loc[(df['seed'] == int(seed)) & (df['env_name'] == str(env_name))]['succeed_iteration'].iloc[0]
+                        except:
+                            print(f'Env {env_name} with seed {seed} not found in Bidexhands_Video.csv')
+                            
+                        if 'unsolved' in succeed_iteration:
+                            continue
+                        elif pd.isnull(succeed_iteration):
+                            min_iter = DEFAULT_ITER
+                            max_iter = MAX_ITER
+                        elif '-' in succeed_iteration:
+                            [min_iter, max_iter] = succeed_iteration.split('-')
+                        else:
+                            min_iter = succeed_iteration
+                            max_iter = MAX_ITER
+
+                        # check if the checkpoint is in the valid range
+                        valid_checkpoints = np.arange(int(min_iter), int(max_iter)+1000, 1000)
+                        if int(checkpoint) not in valid_checkpoints:
+                            continue
+                    
                     video_path = os.path.join(folder, env_name, video)
                     video_files.append(video_path)
+                    # print(video_path)
+
             files[env_name] = video_files
 
     with open(VIDEO_INFO, 'w') as fp: