handling users

app.py

@@ -26,27 +26,28 @@ except Exception as e:
     dataset = Dataset.from_dict({'image_id': [], 'caption': []}, features=features)
     dataset.push_to_hub(dataset_name)  # Push the empty dataset to Hugging Face
 
-image_folder = "images"
+image_folder = "test"
 image_files = [f for f in os.listdir(image_folder) if f.endswith(('.png', '.jpg', '.jpeg'))]
 lock = threading.Lock()
-current_image = None
 
 # Function to get a random image that hasn’t been annotated or skipped
-def get_next_image():
+def get_next_image(session_data):
     with lock:
-        available_images = [img for img in image_files if img not in dataset["image_id"]]
-        if available_images:
-            return os.path.join(image_folder, random.choice(available_images))
-    return None
+        annotated_images = set(dataset["image_id"])  # Set of annotated images
+        available_images = [img for img in image_files if img not in annotated_images]
+        # Check if the user already has an image
+        if session_data["current_image"] is None and available_images:
+            # Assign a new random image to the user
+            session_data["current_image"] = random.choice(available_images)
+    return os.path.join(image_folder, session_data["current_image"]) if session_data["current_image"] else None
 
 # Function to save the annotation to Hugging Face dataset and fetch the next image
-def save_annotation(caption):
-    global current_image
-    if current_image is None:
+def save_annotation(caption, session_data):
+    if session_data["current_image"] is None:
         return gr.update(visible=False), gr.update(value="All images have been annotated!")
 
     with lock:
-        image_id = os.path.basename(current_image)
+        image_id = session_data["current_image"]
 
         # Save caption or "skipped" based on user input
         if caption.strip().lower() == "skip":
@@ -61,23 +62,25 @@ def save_annotation(caption):
         dataset.push_to_hub(dataset_name)
         print("Pushed updated dataset")
 
+        # Clear user's current image so they get a new one next time
+        session_data["current_image"] = None
+
     # Fetch the next image
-    current_image = get_next_image()
-    if current_image:
-        return gr.update(value=current_image), gr.update(value="")
+    next_image = get_next_image(session_data)
+    if next_image:
+        return gr.update(value=next_image), gr.update(value="")
     else:
         return gr.update(visible=False), gr.update(value="All images have been annotated!")
 
 # Function to skip the current image
-def skip_image():
-    return save_annotation("skip")
+def skip_image(session_data):
+    return save_annotation("skip", session_data)
 
 # Function to initialize the interface
-def initialize_interface():
-    global current_image
-    current_image = get_next_image()
-    if current_image:
-        return gr.update(value=current_image), gr.update(value="")
+def initialize_interface(session_data):
+    next_image = get_next_image(session_data)
+    if next_image:
+        return gr.update(value=next_image), gr.update(value="")
     else:
         return gr.update(visible=False), gr.update(value="All images have been annotated!")
 
@@ -86,6 +89,8 @@ with gr.Blocks() as demo:
     gr.Markdown("# Image Captioning Tool")
     gr.Markdown("Please provide a caption for each image displayed. Click 'Submit' after writing your caption, or type 'skip' if you don’t want to annotate this image.")
 
+    session_data = gr.State({"current_image": None})  # Session-specific state
+
     with gr.Row():
         image = gr.Image()
         caption = gr.Textbox(placeholder="Enter caption here...")
@@ -93,10 +98,10 @@ with gr.Blocks() as demo:
         skip = gr.Button("Skip")  # Skip button
 
     # Define actions for buttons
-    submit.click(fn=save_annotation, inputs=caption, outputs=[image, caption])
-    skip.click(fn=skip_image, inputs=None, outputs=[image, caption])
+    submit.click(fn=save_annotation, inputs=[caption, session_data], outputs=[image, caption])
+    skip.click(fn=skip_image, inputs=session_data, outputs=[image, caption])
 
     # Load initial image
-    demo.load(fn=initialize_interface, inputs=None, outputs=[image, caption])
+    demo.load(fn=initialize_interface, inputs=session_data, outputs=[image, caption])
 
 demo.launch(share=True)