Update app.py

app.py

@@ -55,7 +55,27 @@ def get_huggingface_dataset():
     log_file = os.path.join(dataset_dir, "flag_data.csv")
     return repo, log_file
 
-def update(user_choice, data_folder=VIDEO_PATH):
+def update(user_choice, left, right, data_folder=VIDEO_PATH, flag_to_huggingface=True):
+    global last_left_video_path 
+    global last_right_video_path 
+
+    if flag_to_huggingface: # log
+        env_name = str(last_left_video_path).split('/')[1]  # 'robotinder-data/ENV_NAME/'
+        current_time = strftime("%Y-%m-%d-%H-%M-%S", gmtime())
+        info = [env_name, user_choice, last_left_video_path, last_right_video_path, current_time]
+        print(info)
+        repo, log_file = get_huggingface_dataset()
+        with open(log_file, 'a') as file: # incremental change of the file
+            writer_object = writer(file)
+            writer_object.writerow(info)
+            file.close()
+        if int(current_time.split('-')[-2]) % 2 == 0:
+            try:
+                repo.push_to_hub(commit_message=f"Flagged sample at {current_time}")
+            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)]
     # choose video
@@ -64,23 +84,12 @@ def update(user_choice, data_folder=VIDEO_PATH):
     for f in videos:
         if f.endswith(f'.{FORMAT}'):
             video_files.append(os.path.join(data_folder, env_name, f))
-    # choose two videos
+    # 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]]
-
-    # log
-    current_time = strftime("%Y-%m-%d-%H-%M-%S", gmtime())
-    info = [env_name, user_choice, left, right, current_time]
-    print(info)
-
-    repo, log_file = get_huggingface_dataset() # flag without using gradio flagging
-    with open(log_file, 'a') as file: # incremental change of the file
-        writer_object = writer(file)
-        writer_object.writerow(info)
-        file.close()
-    repo.push_to_hub(commit_message=f"Flagged sample at {current_time}")
-
+    last_left_video_path = left
+    last_right_video_path = right
     return left, right
 
 def replay(left, right):  
@@ -107,23 +116,26 @@ def build_interface(iter=3, data_folder=VIDEO_PATH):
         gr.Markdown("Here is RoboTinder!")
         gr.Markdown("Select the best robot behaviour in your choice!")
         with gr.Row():
-            # some initial videos
+            # some initial values
+            envs = parse_envs()   
+            env_name = envs[random.randint(0, len(envs)-1)] # random pick an env 
+            left_video_path = os.path.join(os.path.dirname(__file__), 
+                                f"{VIDEO_PATH}/{env_name}/rl-video-episode-0.{FORMAT}")
+            right_video_path = os.path.join(os.path.dirname(__file__), 
+                                f"{VIDEO_PATH}/{env_name}/rl-video-episode-1.{FORMAT}")
             if FORMAT == 'mp4':
-                left_video_path = os.path.join(os.path.dirname(__file__), 
-                                    f"{VIDEO_PATH}/rl-video-episode-0.mp4")
-                right_video_path = os.path.join(os.path.dirname(__file__), 
-                                    f"{VIDEO_PATH}/rl-video-episode-1.mp4")
                 left = gr.PlayableVideo(left_video_path, label="left_video")
                 right = gr.PlayableVideo(right_video_path, label="right_video")
             else:
-                left_video_path = os.path.join(os.path.dirname(__file__), 
-                                    f"{VIDEO_PATH}/rl-video-episode-0.gif")
-                right_video_path = os.path.join(os.path.dirname(__file__), 
-                                    f"{VIDEO_PATH}/rl-video-episode-1.gif")  
                 left = gr.Image(left_video_path, shape=(1024, 768), label="left_video")
                 # right = gr.Image(right_video_path).style(height=768, width=1024)
                 right = gr.Image(right_video_path, label="right_video")
 
+        global last_left_video_path 
+        last_left_video_path = left_video_path
+        global last_right_video_path 
+        last_right_video_path = right_video_path
+
         btn1 = gr.Button("Replay")
         user_choice = gr.Radio(["Left", "Right", "Not Sure"], label="Which one is your favorite?")
         btn2 = gr.Button("Next")
@@ -140,6 +152,9 @@ def build_interface(iter=3, data_folder=VIDEO_PATH):
     return demo
 
 if __name__ == "__main__":
+    last_left_video_path = None
+    last_right_video_path = None
+
     demo = build_interface()
     # demo.launch(share=True)
     demo.launch(share=False)