Update demo.py

Modify code for showing only two videos, from the new folders.

demo.py

@@ -13,26 +13,10 @@ from firebase_admin import db, credentials
 
 # Constants for video display
 VIDEO_FOLDERS = {
-    "ours_walk_1": "./videos/samples_save_myomni_walk_ours/",
-    "ours_walk_2": "./videos/samples_save_myomni_walk_ours_2/",
-    "ours_walk_3": "./videos/samples_save_myomni_walk_ours_3/",
-    "ours_run_1": "./videos/samples_save_myomni_run_ours/",
-    "ours_run_2": "./videos/samples_save_myomni_run_ours_2/",
-    "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_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/",
-    "sum_walk_1": "./videos/samples_save_myomni_walk_sum/",
-    "sum_walk_2": "./videos/samples_save_myomni_walk_sum_2/",
-    "sum_walk_3": "./videos/samples_save_myomni_walk_sum_3/",
-    "sum_run_1": "./videos/samples_save_myomni_run_sum/",
-    "sum_run_2": "./videos/samples_save_myomni_run_sum_2/",
-    "sum_run_3": "./videos/samples_save_myomni_run_sum_3/",
+    "ours": "./ours/",
+    "theirs": "./3d/",
 }
-NUMBER_OF_EXPERIMENTS = 20  # Adjusted to match the user's comment
+NUMBER_OF_EXPERIMENTS = 16  # Adjusted to match the user's comment
 
 #################################################################################################################################################
 # Authentication
@@ -61,11 +45,10 @@ class Experiment(dict):
     """
     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_right, selected_video=None):
         super().__init__(
             dataset=dataset,
             video_left=video_left,
-            video_middle=video_middle,
             video_right=video_right,
             selected_video=selected_video,
         )
@@ -78,7 +61,6 @@ def experiment_to_dict(experiment, skip=False):
     info = {
         "dataset": experiment["dataset"],
         "video_left": experiment["video_left"],
-        "video_middle": experiment["video_middle"],
         "video_right": experiment["video_right"],
     }
 
@@ -91,8 +73,6 @@ 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"] == "middle":
-            selected_video_path = Path(experiment["video_middle"])
         elif experiment["selected_video"] == "right":
             selected_video_path = Path(experiment["video_right"])
         else:
@@ -100,12 +80,10 @@ def experiment_to_dict(experiment, skip=False):
 
         if selected_video_path:
             # Check the parent folder name to determine the source
-            if "sdf" in selected_video_path.parts:
-                info["Control"] = "SDF"
-            elif "sum" in selected_video_path.parts:
-                info["Control"] = "SUM"
-            elif "ours" in selected_video_path.parts:
-                info["Control"] = "Ours"  # Corrected from "SUM" to "Ours"
+            if "ours" in selected_video_path.parts:
+                info["Control"] = "Ours"
+            elif "3d" in selected_video_path.parts:
+                info["Control"] = "Theirs"
             else:
                 info["Control"] = "Unknown"
         else:
@@ -121,40 +99,39 @@ def generate_new_experiment() -> Experiment:
     The positions of the videos (left/middle/right) are randomized.
     """
     # Randomly select an action and folder
-    random_action = random.choice(["walk", "run"])
-    random_folder = random.randint(1, 3)  # Adjusted to match folder numbering
+
+    random_videos = random.randint(1, 16)  # Adjusted to match folder numbering
 
     # Generate a two-digit sample number as a string
-    if random_folder > 1:
-        random_int = random.randint(0, NUMBER_OF_EXPERIMENTS - 1)
-    else:
-        random_int = random.randint(0, 10)
-    video_name = f"sample00_rep{random_int:02d}.mp4"
+    # if random_folder > 1:
+    #     random_int = random.randint(0, NUMBER_OF_EXPERIMENTS - 1)
+    # else:
+    #     random_int = random.randint(0, 10)
+    video_name = f"{random_videos:02d}.mov"
 
     # 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
+        video_ours_path = Path(VIDEO_FOLDERS["ours"]) / video_name
+        video_theirs_path = Path(VIDEO_FOLDERS["theirs"]) / 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]
+    videos = [video_ours_path, video_theirs_path]
 
     # 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
+    # Shuffle the videos to assign to left, right
     random.shuffle(videos)
-    video_left, video_middle, video_right = videos  # Correct assignment after shuffling
+    video_left, 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}")
+    print(f"Generated new experiment with left video: {video_left}, right video: {video_right}")
     return Experiment(
         DATASET,
         str(video_left),
-        str(video_middle),
         str(video_right),
     )
 
@@ -168,7 +145,6 @@ def load_initial():
         return [
             new_experiment,
             gr.update(value=new_experiment["video_left"]),
-            gr.update(value=new_experiment["video_middle"]),
             gr.update(value=new_experiment["video_right"]),
             ""  # Empty message
         ]
@@ -185,8 +161,6 @@ def select_and_load(selected_label, experiment, message_component):
         # Update the experiment's selected video
         if selected_label.lower() == "left":
             experiment["selected_video"] = "left"
-        elif selected_label.lower() == "middle":
-            experiment["selected_video"] = "middle"
         elif selected_label.lower() == "right":
             experiment["selected_video"] = "right"
         else:
@@ -215,7 +189,6 @@ def select_and_load(selected_label, experiment, message_component):
         return [
             new_experiment,
             gr.update(value=new_experiment["video_left"]),
-            gr.update(value=new_experiment["video_middle"]),
             gr.update(value=new_experiment["video_right"]),
             new_message
         ]
@@ -227,7 +200,6 @@ def select_and_load(selected_label, experiment, message_component):
         return [
             experiment,
             gr.update(value=experiment["video_left"]),
-            gr.update(value=experiment["video_middle"]),
             gr.update(value=experiment["video_right"]),
             new_message
         ]
@@ -264,7 +236,6 @@ def skip_experiment(experiment, message_component):
         return [
             new_experiment,
             gr.update(value=new_experiment["video_left"]),
-            gr.update(value=new_experiment["video_middle"]),
             gr.update(value=new_experiment["video_right"]),
             new_message
         ]
@@ -276,7 +247,6 @@ def skip_experiment(experiment, message_component):
         return [
             experiment,
             gr.update(value=experiment["video_left"]),
-            gr.update(value=experiment["video_middle"]),
             gr.update(value=experiment["video_right"]),
             new_message
         ]
@@ -331,15 +301,6 @@ with gr.Blocks(title="Unsupervised Video Editing", css=css) as demo:
                 show_share_button=False,
                 interactive=True
             )
-        with gr.Column(scale=1):
-            video_middle_component = gr.Video(
-                label="middle",
-                elem_id="video-middle",
-                show_label=False,
-                show_download_button=False,
-                show_share_button=False,
-                interactive=True
-            )
         with gr.Column(scale=1):
             video_right_component = gr.Video(
                 label="right",
@@ -354,8 +315,6 @@ with gr.Blocks(title="Unsupervised Video Editing", css=css) as demo:
         # Add Select buttons below each video
         with gr.Column(scale=1):
             select_left = gr.Button("Select Left Video", elem_id="select-button-left", variant="primary")
-        with gr.Column(scale=1):
-            select_middle = gr.Button("Select Middle Video", elem_id="select-button-middle", variant="primary")
         with gr.Column(scale=1):
             select_right = gr.Button("Select Right Video", elem_id="select-button-right", variant="primary")
 
@@ -363,7 +322,7 @@ with gr.Blocks(title="Unsupervised Video Editing", css=css) as demo:
     btn_skip.click(
         fn=skip_experiment,
         inputs=[experiment, message_component],
-        outputs=[experiment, video_left_component, video_middle_component, video_right_component, message_component],
+        outputs=[experiment, video_left_component, video_right_component, message_component],
         queue=True,
         show_progress=True
     )
@@ -376,13 +335,7 @@ with gr.Blocks(title="Unsupervised Video Editing", css=css) as demo:
         queue=True,
         show_progress=True
     )
-    select_middle.click(
-        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,
-        show_progress=True
-    )
+
     select_right.click(
         fn=lambda exp, msg: select_and_load("right", exp, msg),
         inputs=[experiment, message_component],
@@ -395,7 +348,7 @@ 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_right_component, message_component]
     )
 
 # Launch the app without share=True since it's not supported on Hugging Face Spaces