Update demo.py

demo.py

@@ -21,7 +21,7 @@ VIDEO_FOLDERS = {
     "ours_run_3": "./videos/samples_save_myomni_run_ours_3/",
     "sdf_walk_1": "./videos/samples_save_myomni_walk_sdf/",
     "sdf_walk_2": "./videos/samples_save_myomni_walk_sdf_2/",
-    "sdf_walk_2": "./videos/samples_save_myomni_walk_sdf_3/",
+    "sdf_walk_3": "./videos/samples_save_myomni_walk_sdf_3/",  # Fixed duplicate key
     "sdf_run_1": "./videos/samples_save_myomni_run_sdf/",
     "sdf_run_2": "./videos/samples_save_myomni_run_sdf_2/",
     "sdf_run_3": "./videos/samples_save_myomni_run_sdf_3/",
@@ -32,7 +32,7 @@ VIDEO_FOLDERS = {
     "sum_run_2": "./videos/samples_save_myomni_run_sum_2/",
     "sum_run_3": "./videos/samples_save_myomni_run_sum_3/",
 }
-NUMBER_OF_EXPERIMENTS = 20  # Assuming sample0.mp4 to sample199.mp4
+NUMBER_OF_EXPERIMENTS = 20  # Adjusted to match the user's comment
 
 #################################################################################################################################################
 # Authentication
@@ -59,22 +59,21 @@ except Exception as e:
 
 class Experiment(dict):
     """
-    Represents an experiment consisting of two videos and the user's selection.
+    Represents an experiment consisting of three videos and the user's selection.
     """
-    def __init__(self, dataset, video_left,video_middle, video_right, selected_video=None):
+    def __init__(self, dataset, video_left, video_middle, video_right, selected_video=None):
         super().__init__(
             dataset=dataset,
             video_left=video_left,
             video_middle=video_middle,
             video_right=video_right,
-            
             selected_video=selected_video,
         )
 
 def experiment_to_dict(experiment, skip=False):
     """
     Converts the Experiment object to a dictionary for Firebase storage.
-    Includes whether the selected video is from 'omni' or 'ours' folder.
+    Includes whether the selected video is from 'sdf', 'sum', or 'ours' folder.
     """
     info = {
         "dataset": experiment["dataset"],
@@ -92,10 +91,10 @@ def experiment_to_dict(experiment, skip=False):
         # Determine "Control" based on which video was selected
         if experiment["selected_video"] == "left":
             selected_video_path = Path(experiment["video_left"])
-        elif experiment["selected_video"] == "right":
-            selected_video_path = Path(experiment["video_right"])
         elif experiment["selected_video"] == "middle":
             selected_video_path = Path(experiment["video_middle"])
+        elif experiment["selected_video"] == "right":
+            selected_video_path = Path(experiment["video_right"])
         else:
             selected_video_path = None
 
@@ -106,7 +105,7 @@ def experiment_to_dict(experiment, skip=False):
             elif "sum" in selected_video_path.parts:
                 info["Control"] = "SUM"
             elif "ours" in selected_video_path.parts:
-                info["Control"] = "SUM"
+                info["Control"] = "Ours"  # Corrected from "SUM" to "Ours"
             else:
                 info["Control"] = "Unknown"
         else:
@@ -118,35 +117,37 @@ def experiment_to_dict(experiment, skip=False):
 
 def generate_new_experiment() -> Experiment:
     """
-    Generates a new experiment by randomly selecting a sample number and fetching corresponding videos from both folders.
-    The positions of the videos (left/right) are randomized.
+    Generates a new experiment by randomly selecting a sample number and fetching corresponding videos from multiple folders.
+    The positions of the videos (left/middle/right) are randomized.
     """
-    # Randomly select an integer between 0 and 199 inclusive
-    
+    # Randomly select an action and folder
     random_action = random.choice(["walk", "run"])
-    random_folder = random.randint(1,4)
-    if random_folder == 1:
-        random_int = random.randint(0, 10)
-    else:
-        random_int = random.randint(0, NUMBER_OF_EXPERIMENTS - 1)
+    random_folder = random.randint(1, 3)  # Adjusted to match folder numbering
+
+    # Generate a two-digit sample number as a string
+    random_int = random.randint(0, NUMBER_OF_EXPERIMENTS - 1)
     video_name = f"sample00_rep{random_int:02d}.mp4"
 
-    # Paths to the two videos
-    video_ours_path = Path(VIDEO_FOLDERS[f"ours_{random_action}_{random_folder}"]) / video_name
-    video_sdf_path = Path(VIDEO_FOLDERS[f"sdf_{random_action}_{random_folder}"]) / video_name
-    video_sum_path = Path(VIDEO_FOLDERS[f"sum_{random_action}_{random_folder}"]) / video_name
+    # Paths to the three videos
+    try:
+        video_ours_path = Path(VIDEO_FOLDERS[f"ours_{random_action}_{random_folder}"]) / video_name
+        video_sdf_path = Path(VIDEO_FOLDERS[f"sdf_{random_action}_{random_folder}"]) / video_name
+        video_sum_path = Path(VIDEO_FOLDERS[f"sum_{random_action}_{random_folder}"]) / video_name
+    except KeyError as e:
+        raise KeyError(f"Invalid folder key: {e}")
+
     videos = [video_ours_path, video_sdf_path, video_sum_path]
-    # Ensure that both videos exist
-    if not video_ours_path.exists() or not video_sdf_path.exists() or not video_sum_path.exists():
-        raise FileNotFoundError(f"One or both video files not found: {video_ours_path}, {video_sdf_path}, {video_sum_path}")
-
-    # Randomly decide the order (left/right/middle)
-    # if random.choice([True, False]):
-    #     video_left, video_right = video_ours_path, video_omni_path
-    # else:
-    #     video_left, video_right = video_omni_path, video_ours_path
-    video_left, video_middle, video_right = random.shuffle(videos)
-    print(f"Generated new experiment with left video: {video_left}, right video: {video_right}")
+
+    # Ensure that all videos exist
+    missing_videos = [str(v) for v in videos if not v.exists()]
+    if missing_videos:
+        raise FileNotFoundError(f"The following video files were not found: {', '.join(missing_videos)}")
+
+    # Shuffle the videos to assign to left, middle, right
+    random.shuffle(videos)
+    video_left, video_middle, video_right = videos  # Correct assignment after shuffling
+
+    print(f"Generated new experiment with left video: {video_left}, middle video: {video_middle}, right video: {video_right}")
     return Experiment(
         DATASET,
         str(video_left),
@@ -317,7 +318,7 @@ with gr.Blocks(title="Unsupervised Video Editing", css=css) as demo:
             message_component = gr.Markdown("")  # For displaying messages
 
     with gr.Row():
-        # Define both video components first
+        # Define three video components first
         with gr.Column(scale=1):
             video_left_component = gr.Video(
                 label="left",
@@ -355,15 +356,14 @@ with gr.Blocks(title="Unsupervised Video Editing", css=css) as demo:
         with gr.Column(scale=1):
             select_right = gr.Button("Select Right Video", elem_id="select-button-right", variant="primary")
 
-    with gr.Row():
-        # Configure the skip button
-        btn_skip.click(
-            fn=skip_experiment,
-            inputs=[experiment, message_component],
-            outputs=[experiment, video_left_component, video_middle_component, video_right_component, message_component],
-            queue=True,
-            show_progress=True
-        )
+    # Configure the skip button
+    btn_skip.click(
+        fn=skip_experiment,
+        inputs=[experiment, message_component],
+        outputs=[experiment, video_left_component, video_middle_component, video_right_component, message_component],
+        queue=True,
+        show_progress=True
+    )
 
     # Configure the select buttons to handle selection and load next experiment
     select_left.click(
@@ -374,7 +374,7 @@ with gr.Blocks(title="Unsupervised Video Editing", css=css) as demo:
         show_progress=True
     )
     select_middle.click(
-        fn=lambda exp, msg: select_and_load("left", exp, msg),
+        fn=lambda exp, msg: select_and_load("middle", exp, msg),  # Corrected to "middle"
         inputs=[experiment, message_component],
         outputs=[experiment, video_left_component, video_middle_component, video_right_component, message_component],
         queue=True,
@@ -392,8 +392,8 @@ with gr.Blocks(title="Unsupervised Video Editing", css=css) as demo:
     demo.load(
         fn=load_initial,
         inputs=None,
-        outputs=[experiment, video_left_component,video_middle_component, video_right_component, message_component]
+        outputs=[experiment, video_left_component, video_middle_component, video_right_component, message_component]
     )
 
-# Launch the app with share=True to create a public link
-demo.launch(share=True)
+# Launch the app without share=True since it's not supported on Hugging Face Spaces
+demo.launch()