vote not in metadata

src/global_variables.py

@@ -8,7 +8,7 @@ import jsonlines
 
 import gradio as gr
 
-from src.constants import DATASET_NAME, HF_TOKEN, ASSETS_FOLDER
+from src.constants import DATASET_NAME, HF_TOKEN, ASSETS_FOLDER, CONCEPTS
 
 hf_api: HfApi
 all_metadata: dict
@@ -64,6 +64,76 @@ def save_metadata(split):
         repo_type="dataset",
     )
 
+def update_votes(
+    username: str,
+    current_image: str,
+    voted_concepts: list,
+):
+    global all_metadata
+    global all_votes
+    current_split, idx = current_image.split(":")
+    idx = int(idx)
+    s_id = all_metadata[current_split][idx]["id"]
+    if s_id not in all_votes:
+        all_votes[s_id] = {}
+    all_votes[s_id][username] = {c: c in voted_concepts for c in CONCEPTS}
+    new_concepts = compute_concepts(all_votes[s_id])
+    all_metadata[current_split][idx]["concepts"] = new_concepts
+    
+def compute_concepts(votes):
+    vote_sum = {c: 0 for c in CONCEPTS}
+    for vote in votes.values():
+        for c in CONCEPTS:
+            if c not in vote:
+                continue
+            vote_sum[c] += 2 * vote[c] - 1
+    return {c: vote_sum[c] > 0 if vote_sum[c] != 0 else None for c in CONCEPTS}
+
+def save_current_work(
+    username: str,
+):
+    global all_metadata
+    global all_votes
+    global hf_api
+    hf_api.snapshot_download(
+        local_dir=f"{ASSETS_FOLDER}/{DATASET_NAME}",
+        repo_id=DATASET_NAME,
+        repo_type="dataset",
+        allow_patterns=["*/metadata.jsonl", "votes/*"],
+    )
+    new_votes = {}
+    for file in os.listdir(f"{ASSETS_FOLDER}/{DATASET_NAME}/votes"):
+        with open(f"{ASSETS_FOLDER}/{DATASET_NAME}/votes/{file}") as f:
+            key = file.split(".")[0]
+            new_votes[key] = json.load(f)
+    for key in all_votes:
+        if username in all_votes[key]:
+            if key not in new_votes:
+                new_votes[key] = {}
+            new_votes[key][username] = all_votes[key][username]
+    for key in new_votes:
+        with open(f"{ASSETS_FOLDER}/{DATASET_NAME}/votes/{key}.json", "w") as f:
+            json.dump(new_votes[key], f)
+    all_votes = new_votes
+    new_metadata = {}
+    for split in ["train", "validation", "test"]:
+        new_metadata[split] = []
+        with jsonlines.open(f"{ASSETS_FOLDER}/{DATASET_NAME}/data/{split}/metadata.jsonl") as reader:
+            for row in reader:
+                s_id = row["id"]
+                if s_id in all_votes:
+                    row["concepts"] = compute_concepts(all_votes[s_id])
+                new_metadata[split].append(row)
+        with jsonlines.open(f"{ASSETS_FOLDER}/{DATASET_NAME}/data/{split}/metadata.jsonl", mode='w') as writer:
+            writer.write_all(new_metadata[split])
+    all_metadata = new_metadata
+    hf_api.upload_folder(
+        folder_path=f"{ASSETS_FOLDER}/{DATASET_NAME}",
+        repo_id=DATASET_NAME,
+        repo_type="dataset",
+        allow_patterns=["*/metadata.jsonl", "votes/*"],
+    )
+
 def get_votes(key):
     global all_votes
     global hf_api

src/label_interface.py

@@ -86,25 +86,9 @@ def submit_label(
     if current_image is None:
         gr.Warning("No image selected.")
         return None, None, None, None, None, None, None, filtered_indices, selected_concepts, selected_sample_type
-    current_split, idx = current_image.split(":")
-    idx = int(idx)
-    global_variables.get_metadata(current_split)
-    s_id = global_variables.all_metadata[current_split][idx]["id"]
-    global_variables.get_votes(s_id)
-    if s_id not in global_variables.all_votes:
-        global_variables.all_votes[s_id] = {}
-    global_variables.all_votes[s_id][username] = {c: c in voted_concepts for c in CONCEPTS}
-    vote_sum = {c: 0 for c in CONCEPTS}
-    new_concepts = {}
-    for c in CONCEPTS:
-        for vote in global_variables.all_votes[s_id].values():
-            if c not in vote:
-                continue
-            vote_sum[c] += 2 * vote[c] - 1
-        new_concepts[c] = vote_sum[c] > 0 if vote_sum[c] != 0 else None
-    global_variables.all_metadata[current_split][idx]["concepts"] = new_concepts
-    global_variables.save_metadata(current_split)
-    global_variables.save_votes(s_id)
+    
+    global_variables.update_votes(username, current_image, voted_concepts)
+    
     gr.Info("Submit success")
     return get_next_image(
         split, 
@@ -116,6 +100,12 @@ def submit_label(
         profile
     )
     
+def save_current_work(
+    profile: gr.OAuthProfile,
+):
+    username = profile.username
+    global_variables.save_current_work(username)
+    gr.Info("Save success")
 
 with gr.Blocks() as interface:
     with gr.Row():
@@ -160,7 +150,11 @@ with gr.Blocks() as interface:
                 )
                 gr.LoginButton()
                 submit_button = gr.Button(
-                    value="Submit",
+                    value="Local Submit",
+                )
+            with gr.Row():
+                save_button = gr.Button(
+                    value="Save",
                 )
             with gr.Group():
                 gr.Markdown(
@@ -205,3 +199,7 @@ with gr.Blocks() as interface:
         inputs=[voted_concepts, current_image, split, concepts, sample_type, filtered_indices, selected_concepts, selected_sample_type],
         outputs=common_output
     )
+    save_button.click(
+        save_current_work,
+        outputs=[image]
+    )

src/sample_interface.py

@@ -77,25 +77,9 @@ def submit_label(
     if current_image is None:
         gr.Warning("No image selected.")
         return None, None, None, None, None, None, None, index, filtered_indices
-    current_split, idx = current_image.split(":")
-    idx = int(idx)
-    global_variables.get_metadata(current_split)
-    s_id = global_variables.all_metadata[current_split][idx]["id"]
-    global_variables.get_votes(s_id)
-    if s_id not in global_variables.all_votes:
-        global_variables.all_votes[s_id] = {}
-    global_variables.all_votes[s_id][username] = {c: c in voted_concepts for c in CONCEPTS}
-    vote_sum = {c: 0 for c in CONCEPTS}
-    new_concepts = {}
-    for c in CONCEPTS:
-        for vote in global_variables.all_votes[s_id].values():
-            if c not in vote:
-                continue
-            vote_sum[c] += 2 * vote[c] - 1
-        new_concepts[c] = vote_sum[c] > 0 if vote_sum[c] != 0 else None
-    global_variables.all_metadata[current_split][idx]["concepts"] = new_concepts
-    global_variables.save_metadata(current_split)
-    global_variables.save_votes(s_id)
+    
+    global_variables.update_votes(username, current_image, voted_concepts)
+
     gr.Info("Submit success")
     return get_next_image(
         split, 
@@ -104,6 +88,12 @@ def submit_label(
         profile
     )
     
+def save_current_work(
+    profile: gr.OAuthProfile,
+):
+    username = profile.username
+    global_variables.save_current_work(username)
+    gr.Info("Save success")
 
 with gr.Blocks() as interface:
     with gr.Row():
@@ -145,7 +135,11 @@ with gr.Blocks() as interface:
                 )
                 gr.LoginButton()
                 submit_button = gr.Button(
-                    value="Submit",
+                    value="Local Submit",
+                )
+            with gr.Row():
+                save_button = gr.Button(
+                    value="Save",
                 )
             with gr.Group():
                 gr.Markdown(
@@ -198,3 +192,8 @@ with gr.Blocks() as interface:
         inputs=common_input,
         outputs=common_output,
     )
+
+    save_button.click(
+        save_current_work,
+        outputs=[image]
+    )