demo update

demo.py

@@ -18,14 +18,11 @@ from PIL import Image
 import httplib2
 
 
-
 #################################################################################################################################################
 # Authentication
 #################################################################################################################################################
 
 
-
-
 # read secret api key
 FIREBASE_API_KEY = os.environ['FirebaseSecret']
 FIREBASE_URL = os.environ['FirebaseURL']
@@ -40,20 +37,23 @@ SCOPES = ["https://www.googleapis.com/auth/drive"]
 
 
 class Experiment():
-    def __init__(self, name, corrupted, options, selected_image=None, initialized=False):
-        self.name = name
+    def __init__(self, corruption, image_id, corrupted, options, selected_image=None, initialized=False):
+        self.corruption = corruption
+        self.image_id = image_id
         self.corrupted = corrupted
         self.options = options
         self.selected_image = selected_image
         self.initialized = initialized
 
-    def to_dict(self):
+    def to_dict(self, custom_algo=None):
         return {
-            "experiment_name": self.name,
-            "corrupted": self.corrupted["name"],
+            "corruption": self.corruption,
+            "image_number": self.image_id,
+            "dataset": "CelebA",
+            "corrupted_filename": self.corrupted["name"],
             "options": [img["name"] for img in self.options],
-            "selected_image": self.selected_image,
-            "algo": self.options[self.selected_image]["algo"]
+            "selected_image": "None" if custom_algo else self.selected_image,
+            "algo": custom_algo if custom_algo is not None else self.options[self.selected_image]["algo"]
         }
     
     def to_pil(self):
@@ -75,8 +75,8 @@ class Experiment():
 
 class App():
 
-    NUM_THREADS = 8
-    NUM_TO_SCHEDULE = 8
+    NUM_THREADS = 16
+    NUM_TO_SCHEDULE = 10
 
     def __init__(self):
         self.init_remote()
@@ -132,16 +132,16 @@ class App():
         corrupted = self.current_experiment.corrupted
         images = self.current_experiment.options
 
-        self.corrupted_component = gr.Image(value=corrupted["pil"], label="corr", show_label=True, show_download_button=False, elem_id="padded")
+        self.corrupted_component = gr.Image(value=corrupted["pil"], label=self.current_experiment.corruption, show_label=True, show_download_button=False, elem_id="corrupted")
         self.img_components = [
-            gr.Image(value=img["pil"], label=f"{i}", show_label=True, show_download_button=False, elem_id="unsel")
+            gr.Image(value=img["pil"], label=f"{i}", show_label=False, show_download_button=False, elem_id="unsel")
             for i, img in enumerate(images)
         ]
 
         selected_index = self.current_experiment.selected_image
         if selected_index is not None:
             self.img_components[selected_index] = (
-                gr.Image(value=images[selected_index]["pil"], label=f"{selected_index}", show_label=True, show_download_button=False, elem_id="sel")
+                gr.Image(value=images[selected_index]["pil"], label=f"{selected_index}", show_label=False, show_download_button=False, elem_id="sel")
             )
 
         return [*self.img_components, self.corrupted_component]
@@ -150,8 +150,9 @@ class App():
         self.current_experiment.selected_image = int(evt.target.label)
         return self.build_components_from_experiment()
 
-    def save(self):
-        if save_to_firebase(self.current_experiment, self.firebase_data_ref):
+    def save(self, mode):
+        print(f"Saving experiment with mode {mode}")
+        if save_to_firebase(self.current_experiment, self.firebase_data_ref, mode=mode):
             self.next_experiment()
             self.build_components_from_experiment()
         return [*self.img_components, self.corrupted_component]
@@ -161,13 +162,15 @@ class App():
 # API calls
 #################################################################################################################################################
 
-def save_to_firebase(experiment, firebase_data_ref):
-        if experiment is None or experiment.selected_image is None:
+def save_to_firebase(experiment: Experiment, firebase_data_ref, mode):
+        if mode == "save" and (experiment is None or experiment.selected_image is None):
             gr.Warning("You must select an image before submitting")
             return False
+        if mode == "skip":
+            experiment.selected_image = None
 
         firebase_data_ref.push({
-            **experiment.to_dict(),
+            **experiment.to_dict(custom_algo=mode if mode == "skip" else None),
             "timestamp": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
         })
 
@@ -177,9 +180,27 @@ def save_to_firebase(experiment, firebase_data_ref):
 def list_folders(service):
     folders = []
 
+    corruptions = [
+        "brightness", 
+        "elastic_transform", 
+        "frost", 
+        "impulse_noise", 
+        "masking_random_color", 
+        "shot_noise", 
+        "speckle_noise", 
+        "contrast", 
+        "fog", 
+        "gaussian_noise", 
+        "latent", 
+        "masking_vline_random_color", 
+        "spatter"
+    ]
+
+    name_query = "(" + " or ".join([f"name contains '{c}'" for c in corruptions]) + ")"
+
     results = (
         service.files()
-        .list(pageSize=50, fields="nextPageToken, files(id, name)", orderBy="name", q="mimeType='application/vnd.google-apps.folder' and name contains 'Experiment'")
+        .list(pageSize=50, fields="nextPageToken, files(id, name)", orderBy="name", q=f"mimeType='application/vnd.google-apps.folder' and {name_query}")
         .execute()
     )
     folders.extend(results.get("files", []))
@@ -188,7 +209,7 @@ def list_folders(service):
         page_token = results["nextPageToken"]
         results = (
             service.files()
-            .list(pageSize=50, fields="nextPageToken, files(id, name)", orderBy="name", q="mimeType='application/vnd.google-apps.folder' and name contains 'Experiment'", pageToken=page_token)
+            .list(pageSize=50, fields="nextPageToken, files(id, name)", orderBy="name", q=f"mimeType='application/vnd.google-apps.folder' and {name_query}", pageToken=page_token)
             .execute()
         )
         folders.extend(results.get("files", []))
@@ -235,17 +256,28 @@ def schedule_downloads(service, q_to_download, q_requested):
             break
 
         print("Scheduling new experiment")
+        # this should give a list of top-level folders of corruption types
         folders = list_folders(service)
+        # print(f"Found {len(folders)} folders")
+        # print(folders)
 
-        # sample a random folder from folders
+        # sample a random corruption from folders
         folder = random.choice(folders)
 
-        # list subfolders in the folder
+        # inside the corruption folders, there should be one folder per image id
         subfolders = list_files_in_folder(service, folder, filter_="mimeType='application/vnd.google-apps.folder'")
+        # print(f"Found {len(subfolders)} subfolders")
+        # print(subfolders)
+
+        # sample a random image from subfolders
+        subfolder = random.choice(subfolders)
+
+        # list subsubfolders in the subfolder
+        subsubfolders = list_files_in_folder(service, subfolder, filter_="mimeType='application/vnd.google-apps.folder'")
 
         # the results should be 2 subfolders: SDEdit and ODEdit
-        odedit_subfolder = [subfolder for subfolder in subfolders if "ODEdit" in subfolder["name"]][0]
-        sdedit_subfolder = [subfolder for subfolder in subfolders if "SDEdit" in subfolder["name"]][0]
+        odedit_subfolder = [f for f in subsubfolders if "ode" in f["name"]][0]
+        sdedit_subfolder = [f for f in subsubfolders if "sde" in f["name"]][0]
 
         odedit_files = list_files_in_folder(service, odedit_subfolder)
         sdedit_files = list_files_in_folder(service, sdedit_subfolder)    
@@ -256,11 +288,11 @@ def schedule_downloads(service, q_to_download, q_requested):
         selected_sdedit_files = random.sample(sdedit_files, k=5)
         selected_sdedit_files = [{**file, "algo": "SDEdit"} for file in selected_sdedit_files]
 
-        corrupted_file = list_files_in_folder(service, folder, filter_="mimeType contains 'image/'")[0]
+        corrupted_file = list_files_in_folder(service, subfolder, filter_="mimeType contains 'image/'")[0]
 
         selected_files = [*selected_odedit_files, *selected_sdedit_files]
 
-        experiment = Experiment(folder["name"], corrupted_file, selected_files)
+        experiment = Experiment(folder["name"], subfolder["name"], corrupted_file, selected_files)
 
         q_to_download.put(experiment)
         q_requested.task_done()
@@ -298,8 +330,9 @@ def download_thread(service, q_to_download, q_processed):
 css = """
 #unsel {border: solid 5px transparent !important; border-radius: 15px !important}
 #sel {border: solid 5px #00c0ff !important; border-radius: 15px !important}
-#padded {margin-left: 25% !important; margin-right: 5% !important}
-#paddedRight {margin-right: 5% !important}
+#corrupted {margin-left: 5%; margin-right: 5%; padding: 0 !important}
+#reducedHeight {height: 10px !important}
+#padded {padding-left: 2%; padding-right: 2%}
 """
 
 def build_demo():
@@ -307,24 +340,25 @@ def build_demo():
 
     with gr.Blocks(title="Unsupervised Image Editing", css=css) as demo:
 
-        with gr.Row():
-            corrupted_component = gr.Image(label="corr", elem_id="padded")
-            with gr.Column(scale=3):
-                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>", elem_id="paddedRight")
-                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</h3></div>", elem_id="paddedRight")
-                btn = gr.Button("Submit")
-        gr.Markdown("<hr>")
+        with gr.Row(elem_id="padded"):
+            corrupted_component = gr.Image(label="corr", elem_id="corrupted", show_label=False, show_download_button=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</h3></div>")
+                btn_skip = gr.Button("I have no preference")
+                btn_submit = gr.Button("Submit preference")
 
         img_components = []
         with gr.Row():
             for i, img in enumerate(app.img_components[:5]):
-                img_components.append(gr.Image(label=f"{i}", elem_id="unsel"))
+                img_components.append(gr.Image(label=f"{i}", elem_id="unsel", show_label=False, show_download_button=False))
         
         with gr.Row():
             for i, img in enumerate(app.img_components[5:]):
-                img_components.append(gr.Image(label=f"{i+5}", elem_id="unsel"))
+                img_components.append(gr.Image(label=f"{i+5}", elem_id="unsel", show_label=False, show_download_button=False))
 
-        btn.click(app.save, None, [*img_components, corrupted_component])
+        btn_skip.click(partial(app.save, mode="skip"), None, [*img_components, corrupted_component])
+        btn_submit.click(partial(app.save, mode="save"), None, [*img_components, corrupted_component])
         for img in img_components:
             img.select(app.on_select, None, img_components, show_progress="hidden")