updates

app.py

@@ -36,6 +36,7 @@ lock = threading.Lock()
 with open('nearest_neighbors_with_captions.json', 'r') as f:
     results = json.load(f)
 
+
 def get_caption_for_image_id(image_path):
     """
     Retrieve the caption for a given image_id from the JSON data.
@@ -62,17 +63,27 @@ def get_caption_for_image_id(image_path):
     print("Caption not found for image_id:", image_id)  # Debugging line
     return None
 
-# Function to get a random image that hasn’t been annotated or skipped
+
+# Function to get a random image that hasn’t been fully annotated
 def get_next_image(session_data):
     with lock:
-        annotated_images = set(dataset["image_id"])  # Set of annotated images
-        available_images = [img for img in image_files if img not in annotated_images]
+        annotated_images = {item["image_id"]: item["annotation_count"] for item in dataset}
+
+        # Available images filter
+        available_images = [
+            img for img in image_files
+            if img not in annotated_images or
+               ("val" in img and annotated_images[img] < 5) or
+               ("val" not in img and annotated_images[img] == 0)
+        ]
+
         # 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, session_data):
     global dataset  # Declare global dataset at the start of the function
@@ -107,8 +118,9 @@ def save_annotation(caption, session_data):
         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
+        # Clear user's current image if the validation image has been annotated five times
+        if ("val" not in image_id) or (annotation_count > 2):
+            session_data["current_image"] = None
 
     # Fetch the next image
     next_image = get_next_image(session_data)
@@ -118,6 +130,7 @@ def save_annotation(caption, session_data):
     else:
         return gr.update(visible=False), gr.update(value="All images have been annotated!"), gr.update(value="")
 
+
 def initialize_interface(session_data):
     next_image = get_next_image(session_data)
     if next_image:
@@ -127,6 +140,7 @@ def initialize_interface(session_data):
     else:
         return gr.update(visible=False), gr.update(value="All images have been annotated!")
 
+
 # Build the Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown("# Image Captioning Tool")