Update demo.py

demo.py

@@ -8,31 +8,35 @@ import os
 import firebase_admin
 from firebase_admin import db, credentials
 
-
 ##################################################################
 # Constants
 ##################################################################
 
-
-NUMBER_OF_IMAGES_PER_ROW = 7
-NUMBER_OF_ROWS = 2
-
+# Constants for video display
+NUMBER_OF_VIDEOS_PER_ROW = 2  # Displaying two videos per row
+VIDEO_FOLDERS = {
+    "ours": "./videos/userstudy_200_ours/",
+    "omni": "./videos/userstudy_200_omni/"
+}
+NUMBER_OF_EXPERIMENTS = 200  # Assuming sample0.mp4 to sample199.mp4
 
 #################################################################################################################################################
 # Authentication
 #################################################################################################################################################
 
-
-# read secret api key
+# Read secret API key and other configurations from environment variables
 FIREBASE_API_KEY = os.environ['FirebaseSecret']
 FIREBASE_URL = os.environ['FirebaseURL']
 DATASET = os.environ['Dataset']
 
-# init firebase service
-firebase_creds = credentials.Certificate(json.loads(FIREBASE_API_KEY))
-firebase_app = firebase_admin.initialize_app(firebase_creds, {'databaseURL': FIREBASE_URL})
-firebase_data_ref = db.reference("data")
-
+# Initialize Firebase service
+try:
+    firebase_creds = credentials.Certificate(json.loads(FIREBASE_API_KEY))
+    firebase_app = firebase_admin.initialize_app(firebase_creds, {'databaseURL': FIREBASE_URL})
+    firebase_data_ref = db.reference("data")
+except Exception as e:
+    print(f"Failed to initialize Firebase: {e}")
+    raise e
 
 ##################################################################
 # Data Layer
@@ -40,89 +44,114 @@ firebase_data_ref = db.reference("data")
 
 
 class Experiment(dict):
-    def __init__(self, dataset, corruption, image_id, corrupted, options, selected_image=None):
+    """
+    Represents an experiment consisting of two videos and the user's selection.
+    """
+    def __init__(self, dataset, video_left, video_right, selected_video=None):
         super().__init__(
             dataset=dataset,
-            corruption=corruption,
-            image_id=image_id,
-            corrupted=corrupted,
-            options=options,
-            selected_image=selected_image,
+            video_left=video_left,
+            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.
+    """
     info = {
-        # experiment info
         "dataset": experiment["dataset"],
-        "corruption": experiment["corruption"],
-        "image_number": experiment["image_id"],
-
-        # chosen image set info
-        "corrupted_filename": experiment["corrupted"]["name"],
-        "options": [img["name"] for img in experiment["options"]],
+        "video_left": experiment["video_left"],
+        "video_right": experiment["video_right"],
     }
 
     if skip:
-        info = {
-            **info,
-            # selected image info
-            "selected_image": "None",
-            "selected_algo": "None",
-        }
+        info["selected_video"] = "None"
+        info["Control"] = "None"
     else:
-        info = {
-            **info,
-            # selected image info
-            "selected_image": experiment["options"][experiment["selected_image"]]["name"],
-            "selected_algo": experiment["options"][experiment["selected_image"]]["algo"],
-        }
-    
+        info["selected_video"] = experiment["selected_video"]
+
+        # 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"])
+        else:
+            selected_video_path = None
+
+        if selected_video_path:
+            # Check the parent folder name to determine the source
+            if "userstudy_200_omni" in selected_video_path.parts:
+                info["Control"] = "OmniControl"
+            elif "userstudy_200_ours" in selected_video_path.parts:
+                info["Control"] = "Ours"
+            else:
+                info["Control"] = "Unknown"
+        else:
+            info["Control"] = "Unknown"
+
+    info["timestamp"] = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+
     return info
-    
+
+
 def generate_new_experiment() -> Experiment:
-    wanted_corruptions = ["spatter", "impulse_noise", "speckle_noise", "gaussian_noise", "pixelate", "jpeg_compression", "elastic_transform"]
-    corruption = random.choice([f for f in list(Path(f"./images/{DATASET}").glob("*/*")) if f.is_dir() and f.name in wanted_corruptions])
-    image_id = random.choice(list(corruption.glob("*")))
-    imgs_to_sample = (NUMBER_OF_IMAGES_PER_ROW * NUMBER_OF_ROWS) // 2
-
-    corrupted_image = {"name": str(random.choice(list(image_id.glob("*corrupted*"))))}
-    sdedit_images = [
-        {"name": str(img), "algo": "SDEdit"}
-        for img in random.sample(list((image_id / "sde").glob(f"*")), imgs_to_sample)
-    ]
-    odedit_images = [
-        {"name": str(img), "algo": "ODEdit"}
-        for img in random.sample(list((image_id / "ode").glob(f"*")), imgs_to_sample)
-    ]
-    total_images = sdedit_images + odedit_images
-    random.shuffle(total_images)
+    """
+    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.
+    """
+    # Randomly select an integer between 0 and 199 inclusive
+    random_int = random.randint(0, NUMBER_OF_EXPERIMENTS - 1)
+    video_name = f"sample{random_int}.mp4"
+
+    # Paths to the two videos
+    video_ours_path = Path(VIDEO_FOLDERS["ours"]) / video_name
+    video_omni_path = Path(VIDEO_FOLDERS["omni"]) / video_name
+
+    # Ensure that both videos exist
+    if not video_ours_path.exists() or not video_omni_path.exists():
+        raise FileNotFoundError(f"One or both video files not found: {video_ours_path}, {video_omni_path}")
+
+    # Randomly decide the order (left/right)
+    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
 
     return Experiment(
         DATASET,
-        corruption.name,
-        image_id.name,
-        corrupted_image,
-        total_images,
+        str(video_left),
+        str(video_right),
     )
 
-def save(experiment, corrupted_component, *img_components, mode):
-    if mode == "save" and (experiment is None or experiment["selected_image"] is None):
-        gr.Warning("You must select an image before submitting")
-        return [experiment, corrupted_component, *img_components]
-    if mode == "skip":
-        experiment["selected_image"] = None
 
-    dict_to_save = {
-        **experiment_to_dict(experiment, skip=(mode=="skip")),
-        "timestamp": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
-    }
-    firebase_data_ref.push(dict_to_save)
+def save(mode, experiment, video_left_component, video_right_component):
+    """
+    Saves the user's selection to Firebase.
+    - mode: "save" or "skip"
+    """
+    if mode == "save":
+        if experiment is None or experiment["selected_video"] is None:
+            gr.Warning("You must select a video before submitting")
+            return [experiment, video_left_component, video_right_component]
+    elif mode == "skip":
+        experiment["selected_video"] = "None"
+
+    try:
+        dict_to_save = experiment_to_dict(experiment, skip=(mode == "skip"))
+        firebase_data_ref.push(dict_to_save)
+
+        print("=====================")
+        print(dict_to_save)
+        print("=====================")
 
-    print("=====================")
-    print(dict_to_save)
-    print("=====================")
+        gr.Info("Your choice has been saved to Firebase")
+    except Exception as e:
+        gr.Warning(f"Failed to save your choice: {e}")
+        print(f"Error saving to Firebase: {e}")
 
-    gr.Info("Your choice has been saved to Firebase")
     return next()
 
 
@@ -132,61 +161,138 @@ def save(experiment, corrupted_component, *img_components, mode):
 
 
 def next():
+    """
+    Prepares the next set of videos for the user to evaluate.
+    """
     new_experiment = generate_new_experiment()
 
-    new_img_components = [
-        gr.Image(value=img["name"], label=f"{i}", elem_id="unsel", show_label=False, show_download_button=False, show_share_button=False, interactive=False)
-        for i, img in enumerate(new_experiment["options"])
-    ]
-    new_corrupted_component = gr.Image(value=new_experiment["corrupted"]["name"], label="corr", elem_id="corrupted", show_label=False, show_download_button=False, show_share_button=False, interactive=False)
+    # Create video components with initial state as unselected
+    video_left_component = gr.Video(
+        value=new_experiment["video_left"],
+        label="left",
+        elem_id="unsel",
+        show_label=False,
+        show_download_button=False,
+        show_share_button=False,
+        interactive=True
+    )
+    video_right_component = gr.Video(
+        value=new_experiment["video_right"],
+        label="right",
+        elem_id="unsel",
+        show_label=False,
+        show_download_button=False,
+        show_share_button=False,
+        interactive=True
+    )
 
-    return [new_experiment, new_corrupted_component, *new_img_components]
+    return [new_experiment, video_left_component, video_right_component]
 
-def on_select(evt: gr.SelectData, experiment, *img_components):  # SelectData is a subclass of EventData
-    new_selected = int(evt.target.label)
 
-    new_img_components = [
-        gr.Image(value=img["name"], label=f"{i}", elem_id="unsel", show_label=False, show_download_button=False, show_share_button=False, interactive=False)
-        for i, img in enumerate(experiment["options"])
-    ]
-    new_img_components[new_selected] = (
-        gr.Image(value=experiment["options"][new_selected]["name"], label=f"{new_selected}", elem_id="sel", show_label=False, show_download_button=False, show_share_button=False, interactive=False)
-    )
+def on_select(evt: gr.SelectData, experiment, video_left_component, video_right_component):
+    """
+    Handles the selection of a video by the user.
+    """
+    selected_label = evt.value  # Assuming evt.value contains the label ("left" or "right")
 
-    experiment["selected_image"] = int(evt.target.label)
+    if selected_label.lower() == "left":
+        experiment["selected_video"] = "left"
+    elif selected_label.lower() == "right":
+        experiment["selected_video"] = "right"
+    else:
+        experiment["selected_video"] = None  # For safety
+
+    # Update the video components to reflect selection
+    if experiment["selected_video"] == "left":
+        video_left_component.update(elem_id="sel")
+        video_right_component.update(elem_id="unsel")
+    elif experiment["selected_video"] == "right":
+        video_left_component.update(elem_id="unsel")
+        video_right_component.update(elem_id="sel")
+    else:
+        video_left_component.update(elem_id="unsel")
+        video_right_component.update(elem_id="unsel")
+
+    return [experiment, video_left_component, video_right_component]
 
-    return [experiment, *new_img_components]
 
 css = """
 #unsel {border: solid 5px transparent !important; border-radius: 15px !important; draggable: false}
 #sel {border: solid 5px #00c0ff !important; border-radius: 15px !important; draggable: false}
-#corrupted {margin-left: 5%; margin-right: 5%; padding: 0 !important; draggable: false}
 #reducedHeight {height: 10px !important}
 #padded {padding-left: 2%; padding-right: 2%}
 """
 
-with gr.Blocks(title="Unsupervised Image Editing", css=css) as demo:
-    experiment = gr.State(generate_new_experiment())
+with gr.Blocks(title="Unsupervised Video Editing", css=css) as demo:
+    experiment = gr.State()
 
     with gr.Row(elem_id="padded"):
-        corrupted_component = gr.Image(label="corr", elem_id="corrupted", show_label=False, show_download_button=False, show_share_button=False, interactive=False)
         with gr.Column(scale=3, elem_id="padded"):
-            gr.Markdown("<div style='width: 100%'><h1 style='text-align: center; display: inline-block; width: 100%'>The sample on the left is a corrupted image</h1></div>")
-            gr.Markdown("<div style='width: 100%'><h3 style='text-align: center; display: inline-block; width: 100%'>Below are decorrupted versions sampled from various models. Click on the picture you like best.<br/>⚠️Do not pay attention to the background. Consider first fidelity, then quality⚠️</h3></div>")
+            gr.Markdown("<div style='width: 100%'><h1 style='text-align: center; display: inline-block; width: 100%'>Choose Your Preferred Video</h1></div>")
+            gr.Markdown("<div style='width: 100%'><h3 style='text-align: center; display: inline-block; width: 100%'>Select the video you prefer.<br/>⚠️Consider fidelity and quality⚠️</h3></div>")
             btn_skip = gr.Button("I have no preference")
             btn_submit = gr.Button("Submit preference")
 
-    img_components = []
-    for row in range(NUMBER_OF_ROWS):
-        with gr.Row():
-            for col in range(NUMBER_OF_IMAGES_PER_ROW):
-                img_components.append(gr.Image(label=f"{row * NUMBER_OF_IMAGES_PER_ROW + col}", elem_id="unsel", show_label=False, show_download_button=False, show_share_button=False, interactive=False))
+    with gr.Row():
+        # Initialize video components; values will be set on load
+        video_left_component = gr.Video(
+            label="left",
+            elem_id="unsel",
+            show_label=False,
+            show_download_button=False,
+            show_share_button=False,
+            interactive=True
+        )
+        video_right_component = gr.Video(
+            label="right",
+            elem_id="unsel",
+            show_label=False,
+            show_download_button=False,
+            show_share_button=False,
+            interactive=True
+        )
+
+    # Configure the skip and submit buttons
+    btn_skip.click(
+        fn=save,
+        inputs=[gr.State(), video_left_component, video_right_component],
+        outputs=[gr.State(), video_left_component, video_right_component],
+        _js=None,
+        queue=False,
+        show_progress=False,
+        api_name=None,
+        mode="skip"  # Pass mode as "skip"
+    )
+    btn_submit.click(
+        fn=save,
+        inputs=[gr.State(), video_left_component, video_right_component],
+        outputs=[gr.State(), video_left_component, video_right_component],
+        _js=None,
+        queue=False,
+        show_progress=False,
+        api_name=None,
+        mode="save"  # Pass mode as "save"
+    )
 
-    btn_skip.click(partial(save, mode="skip"), [experiment, corrupted_component, *img_components], [experiment, corrupted_component, *img_components])
-    btn_submit.click(partial(save, mode="save"), [experiment, corrupted_component, *img_components], [experiment, corrupted_component, *img_components])
-    for img in img_components:
-        img.select(on_select, [experiment, *img_components], [experiment, *img_components], show_progress="hidden")
+    # Attach select events to video components
+    video_left_component.select(
+        fn=on_select,
+        inputs=[gr.SelectData(), experiment, video_left_component, video_right_component],
+        outputs=[experiment, video_left_component, video_right_component],
+        show_progress="hidden"
+    )
+    video_right_component.select(
+        fn=on_select,
+        inputs=[gr.SelectData(), experiment, video_left_component, video_right_component],
+        outputs=[experiment, video_left_component, video_right_component],
+        show_progress="hidden"
+    )
 
-    demo.load(next, None, [experiment, corrupted_component, *img_components])
+    # Load the first experiment on app startup
+    demo.load(
+        fn=next,
+        inputs=None,
+        outputs=[experiment, video_left_component, video_right_component]
+    )
 
-demo.launch()
+demo.launch()