updates

app.py

@@ -18,6 +18,8 @@ dataset_name = "GeorgeIbrahim/EGYCOCO"  # Replace with your dataset name
 with open('nearest_neighbors_with_captions.json', 'r') as f:
     results = json.load(f)
 
+print(results)
+
 # Load or create the dataset
 try:
     dataset = load_dataset(dataset_name, split="train")
@@ -125,8 +127,14 @@ def save_annotation(caption, session_data):
         return gr.update(visible=False), gr.update(value="All images have been annotated!"), gr.update(value="")
 
     with lock:
+        match = re.search(r'_(\d+)\.', image_path)
         image_id = session_data["current_image"]
-        split = "dev" if image_id in results else "train"
+        print("Image ID before: ", 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)
+        
         # Save caption or "skipped" based on user input
         if caption.strip().lower() == "skip":
             caption = "skipped"