a

app.py

@@ -5,6 +5,7 @@ import random
 from datasets import load_dataset, Dataset, Features, Value, concatenate_datasets
 from huggingface_hub import login
 import json
+import re
 
 # Authenticate with Hugging Face
 token = os.getenv("HUGGINGFACE_TOKEN")
@@ -33,23 +34,32 @@ lock = threading.Lock()
 
 with open('nearest_neighbors_with_captions.json', 'r') as f:
     results = json.load(f)
-    print("Loaded JSON data:", results)  # Add this line to verify the data
+    # print("Loaded JSON data:", results)  # Add this line to verify the data
 
-def get_caption_for_image_id(image_id):
+def get_caption_for_image_id(image_path):
     """
     Retrieve the caption for a given image_id from the JSON data.
     """
-    # Check if image_id is a test image
-    if image_id in results:
-        return results[image_id]["caption"]
-
-    # If image_id is not a test image, search in nearest neighbors
-    for test_image_data in results.values():
-        for neighbor in test_image_data["nearest_neighbors"]:
-            if neighbor["image_id"] == image_id:
-                return neighbor["caption"]
+    # Extract the numeric part of the image ID
+    match = re.search(r'(\d+)', image_path)
+    if match:
+        image_id = match.group(1).lstrip('0')  # Remove leading zeros
+        print("Searching for image_id:", image_id)  # Debugging line
+
+        # Check if image_id is a test image
+        if image_id in results:
+            print("Found caption in results:", results[image_id]["caption"])  # Debugging line
+            return results[image_id]["caption"]
+
+        # If image_id is not a test image, search in nearest neighbors
+        for test_image_data in results.values():
+            for neighbor in test_image_data["nearest_neighbors"]:
+                if neighbor["image_id"] == image_id:
+                    print("Found caption in nearest neighbors:", neighbor["caption"])  # Debugging line
+                    return neighbor["caption"]
 
     # Return None if the image_id is not found
+    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