Update app.py

app.py

@@ -55,7 +55,7 @@ with gr.Blocks() as demo:
             # workaround we just delete the repo
             delete_repo(repo_id="active-learning/to_label_samples", repo_type="dataset")
             
-            # Save to dataset
+            # Push to training dataset
             labeled_dataset = load_dataset("active-learning/labeled_samples")["train"]
             feature = datasets.Image(decode=False)
             for img, label in labeled_data:
@@ -64,16 +64,20 @@ with gr.Blocks() as demo:
                   "image": feature.encode_example(Image.fromarray(img)), 
                   "label": label
               })
-        labeled_dataset.push_to_hub("active-learning/labeled_samples")
-        labeled_data = []
-        idx = 0
-        return {
-            img_to_label: gr.update(visible=False),
-            label_dropdown: gr.update(visible=False),
-            save_btn: gr.update(visible=False),
-            output_box: gr.update(visible=True),
-            reload_btn: gr.update(visible=True)
-        }
+            labeled_dataset.push_to_hub("active-learning/labeled_samples")
+
+            # Clean up data
+            labeled_data = []
+            idx = 0
+
+            # Update UI
+            return {
+                img_to_label: gr.update(visible=False),
+                label_dropdown: gr.update(visible=False),
+                save_btn: gr.update(visible=False),
+                output_box: gr.update(visible=True),
+                reload_btn: gr.update(visible=True)
+            }
         else:
             return {
                 img_to_label: gr.update(value=get_image())
@@ -82,6 +86,8 @@ with gr.Blocks() as demo:
     def reload_data():
         global data_to_label 
         global imgs
+
+        # See if there is new data to be labeled
         data_to_label = load_dataset("active-learning/to_label_samples")
         imgs = data_to_label["train"]["image"]
         if len(imgs) == 0: