updates

app.py

@@ -96,27 +96,28 @@ def get_caption_for_image_id(image_path):
     print("Caption not found for image_id:", image_id)  # Debugging line
     return None
 
+
+# Initialize a dictionary to keep track of how many times each 'dev' image has been shown
 shown_counts = {}
 
 # Function to get a random image that hasn’t been fully annotated
 def get_next_image(session_data):
     with lock:
-        # Filter available images based on annotation counts and split
+        # Filter available images based on annotation counts and the 'split' column
         available_images = [
             img for img in image_files
             if img not in annotation_counts or
-               ("val" in img and annotation_counts.get(img, 0) < 2) or
-               ("val" not in img and annotation_counts.get(img, 0) == 0)
+               (dataset[img]['split'] == "dev" and annotation_counts.get(img, 0) < 2) or
+               (dataset[img]['split'] == "train" and annotation_counts.get(img, 0) == 0)
         ]
 
         print("Available images:", available_images)  # Debugging line
 
-        # Select an image to show based on split type
+        # Select an image to show based on the 'split' column
         if session_data["current_image"] is None and available_images:
             random.shuffle(available_images)  # Shuffle for randomness
             for img in available_images:
-                image_id = re.search(r'_(\d+)\.', img).group(1).lstrip('0')  # Extract image ID
-                split = "dev" if image_id in results else "train"
+                split = dataset[img]['split']  # Get the split value from the dataset
 
                 # Show 'dev' images twice
                 if split == "dev":
@@ -144,10 +145,9 @@ def save_annotation(caption, session_data):
         image_id = session_data["current_image"]
         print("Image ID before: ", image_id)
                 
-        match = re.search(r'_(\d+)\.', image_id)
-        image_id_2 = match.group(1).lstrip('0') 
-        split = "dev" if image_id_2 in results else "train"
-        print("Image ID after: ", image_id_2)
+        # Determine the split directly from the dataset
+        split = dataset[image_id]['split']
+        print("Split for image:", split)
 
         # Save caption or "skipped" based on user input
         if caption.strip().lower() == "skip":
@@ -193,7 +193,8 @@ 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: