label interface

src/explore_interface.py

@@ -18,10 +18,13 @@ def get_next_image(
         for split, values in global_variables.all_metadata.items():
             current_metadata[split] = [x for x in values if all([c in x["concepts"] for c in concepts])]
         selected_concepts = concepts
-    if split == "all":
-        split = random.choice(["train", "validation", "test"])
     try:
-        sample = random.choice(current_metadata[split])
+        if split == "all":
+            sample = random.choice(
+                current_metadata["train"] + current_metadata["validation"] + current_metadata["test"]
+            )
+        else:
+            sample = random.choice(current_metadata[split])
         image_path = f"{ASSETS_FOLDER}/{DATASET_NAME}/data/{split}/{sample['file_name']}"
         return image_path, sample["class"], sample["concepts"], current_metadata, selected_concepts
     except IndexError:
@@ -40,8 +43,9 @@ with gr.Blocks() as interface:
                     choices=["train", "validation", "test", "all"],
                     value="train",
                 )
-                concepts = gr.CheckboxGroup(
+                concepts = gr.Dropdown(
                     label="Concepts",
+                    multiselect=True,
                     choices=CONCEPTS,
                 )
             button = gr.Button(

src/global_variables.py

@@ -28,17 +28,32 @@ def setup():
             for row in reader:
                 all_metadata[split].append(row)
 
-def save_metadata():
+def get_metadata(split):
     global all_metadata
-    for split, values in all_metadata.items():
-        with jsonlines.open(f"{ASSETS_FOLDER}/{DATASET_NAME}/data/{split}/metadata.jsonl", mode='w') as writer:
-            writer.write_all(values)
-        hf_api.upload_file(
-            path_or_fileobj=f"{ASSETS_FOLDER}/{DATASET_NAME}/data/{split}/metadata.jsonl",
-            path_in_repo=f"data/{split}/metadata.jsonl",
-            repo_id=DATASET_NAME,
-            repo_type="dataset",
-        )
+    global hf_api
+    hf_api.hf_hub_download(
+        repo_id=DATASET_NAME,
+        filename="metadata.jsonl",
+        subfolder=f"data/{split}",
+        repo_type="dataset",
+        local_dir=f"{ASSETS_FOLDER}/{DATASET_NAME}",
+    )
+    all_metadata[split] = []
+    with jsonlines.open(f"{ASSETS_FOLDER}/{DATASET_NAME}/data/{split}/metadata.jsonl") as reader:
+        for row in reader:
+            all_metadata[split].append(row)
+
+def save_metadata(split):
+    global all_metadata
+    values = all_metadata[split]
+    with jsonlines.open(f"{ASSETS_FOLDER}/{DATASET_NAME}/data/{split}/metadata.jsonl", mode='w') as writer:
+        writer.write_all(values)
+    hf_api.upload_file(
+        path_or_fileobj=f"{ASSETS_FOLDER}/{DATASET_NAME}/data/{split}/metadata.jsonl",
+        path_in_repo=f"data/{split}/metadata.jsonl",
+        repo_id=DATASET_NAME,
+        repo_type="dataset",
+    )
 
 if gr.NO_RELOAD:
     setup()

src/label_interface.py

@@ -1,18 +1,101 @@
 """Interface for labeling concepts in images.
 """
 
+from typing import Optional
+
+import random
+
 import gradio as gr
 
+from src import global_variables
+from src.constants import CONCEPTS, ASSETS_FOLDER, DATASET_NAME
+
+def get_next_image(
+    split: str,
+    profile: gr.OAuthProfile
+):
+    username = profile.username
+    try:
+        sample_idx = random.choice(range(len(global_variables.all_metadata[split])))
+        sample = global_variables.all_metadata[split][sample_idx]
+        image_path = f"{ASSETS_FOLDER}/{DATASET_NAME}/data/{split}/{sample['file_name']}"
+        try:
+            username_votes = sample["votes"][username]
+            voted_concepts = [c for c in CONCEPTS if username_votes.get(c, False)]
+        except KeyError:
+            voted_concepts = []
+
+        return image_path, voted_concepts, f"{split}:{sample_idx}"
+    except IndexError:
+        gr.Warning("No image found for the selected filter.")
+        return None, None, None
+
+def submit_label(
+    voted_concepts: list,
+    current_image: Optional[str],
+    profile: gr.OAuthProfile
+):
+    username = profile.username
+    if current_image is None:
+        gr.Warning("No image selected.")
+        return None, None, None
+    split, idx = current_image.split(":")
+    idx = int(idx)
+    global_variables.get_metadata(split)
+    if "votes" not in global_variables.all_metadata[split][idx]:
+        global_variables.all_metadata[split][idx]["votes"] = {}
+    global_variables.all_metadata[split][idx]["votes"][username] = {c: c in voted_concepts for c in CONCEPTS}
+    vote_sum = {c: 0 for c in CONCEPTS}
+    concepts = {}
+    for c in CONCEPTS:
+        for vote in global_variables.all_metadata[split][idx]["votes"].values():
+            if c not in vote:
+                continue
+            vote_sum[c] += 2 * vote[c] - 1
+        concepts[c] = vote_sum[c] > 0 if vote_sum[c] != 0 else None
+    global_variables.all_metadata[split][idx]["concepts"] = concepts
+    global_variables.save_metadata(split)
+    return get_next_image(split, profile)
+    
 
 with gr.Blocks() as interface:
     with gr.Row():
         with gr.Column():
-            gr.Image(
+            with gr.Group():
+                gr.Markdown(
+                    "## # Image Selection",
+                )
+                split = gr.Radio(
+                    label="Split",
+                    choices=["train", "validation", "test",],
+                    value="train",
+                )
+            with gr.Group():
+                voted_concepts = gr.CheckboxGroup(
+                    label="Concepts",
+                    choices=CONCEPTS,
+                )
+            with gr.Row():
+                next_button = gr.Button(
+                    value="Next",
+                )
+                gr.LoginButton()
+                submit_button = gr.Button(
+                    value="Submit",
+                )
+
+        with gr.Column():
+            image = gr.Image(
                 label="Image",
             )
-            gr.Textbox(
-                label="Label",
-            )
-            gr.Button(
-                value="Submit",
-            )
+    current_image = gr.State(None)
+    next_button.click(
+        get_next_image, 
+        inputs=[split],
+        outputs=[image, voted_concepts, current_image],
+    )
+    submit_button.click(
+        submit_label,
+        inputs=[voted_concepts, current_image],
+        outputs=[image, voted_concepts, current_image],
+    )