Update app.py

app.py

@@ -10,6 +10,12 @@ from csv import writer
 from datasets import load_dataset
 from hfserver import HuggingFaceDatasetSaver, HuggingFaceDatasetJSONSaver
 
+ENVS = ['ShadowHand', 'ShadowHandCatchAbreast', 'ShadowHandOver', 'ShadowHandBlockStack', 'ShadowHandCatchUnderarm',
+'ShadowHandCatchOver2Underarm', 'ShadowHandBottleCap', 'ShadowHandLiftUnderarm', 'ShadowHandTwoCatchUnderarm',
+'ShadowHandDoorOpenInward', 'ShadowHandDoorOpenOutward', 'ShadowHandDoorCloseInward', 'ShadowHandDoorCloseOutward',
+'ShadowHandPushBlock', 'ShadowHandKettle', 
+'ShadowHandScissors', 'ShadowHandPen', 'ShadowHandSwingCup', 'ShadowHandGraspAndPlace', 'ShadowHandSwitch']
+
 # download data from huggingface dataset
 # dataset = load_dataset("quantumiracle-git/robotinder-data")
 
@@ -89,7 +95,7 @@ def get_huggingface_dataset():
     log_file = os.path.join(dataset_dir, "flag_data.csv")
     return repo, log_file
 
-def update(user_choice, left, right, data_folder=VIDEO_PATH, flag_to_huggingface=True):
+def update(user_choice, left, right, choose_env, data_folder=VIDEO_PATH, flag_to_huggingface=True):
     global last_left_video_path 
     global last_right_video_path 
     global last_infer_left_video_path
@@ -111,9 +117,11 @@ def update(user_choice, left, right, data_folder=VIDEO_PATH, flag_to_huggingface
             except:
                 repo.git_pull(lfs=True)  # sync with remote first
                 repo.push_to_hub(commit_message=f"Flagged sample at {current_time}")
-
-    envs = parse_envs()   
-    env_name = envs[random.randint(0, len(envs)-1)]
+    if choose_env == 'Random':
+        envs = parse_envs()   
+        env_name = envs[random.randint(0, len(envs)-1)]
+    else:
+        env_name = choose_env
     # choose video
     videos = os.listdir(os.path.join(data_folder, env_name))
     video_files = []
@@ -156,22 +164,26 @@ def build_interface(iter=3, data_folder=VIDEO_PATH):
 
     # build gradio interface
     with gr.Blocks() as demo:
-        gr.Markdown("Here is RoboTinder!")
+        gr.Markdown("Here is **RoboTinder**!")
         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 
+        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]]
+        
         with gr.Row():
-            # some initial values
-            envs = parse_envs()   
-            env_name = envs[random.randint(0, len(envs)-1)] # random pick an 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_path = video_files[selected_video_ids[0]]
-            right_video_path = video_files[selected_video_ids[1]]
             if FORMAT == 'mp4':
                 # left = gr.PlayableVideo(left_video_path, label="left_video")
                 # right = gr.PlayableVideo(right_video_path, label="right_video")
@@ -197,13 +209,14 @@ def build_interface(iter=3, data_folder=VIDEO_PATH):
 
         # btn1 = gr.Button("Replay")
         user_choice = gr.Radio(["Left", "Right", "Not Sure"], label="Which one is your favorite?")
+        choose_env = gr.Radio(["Random"]+ENVS, label="Choose the next task:")
         btn2 = gr.Button("Next")
 
         # This needs to be called at some point prior to the first call to callback.flag()
         callback.setup([user_choice, left, right], "flagged_data_points")
         
         # btn1.click(fn=replay, inputs=[left, right], outputs=[left, right])
-        btn2.click(fn=update, inputs=[user_choice, left, right], outputs=[left, right])
+        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)