Update app.py

app.py

@@ -1,3 +1,4 @@
+# app.py
 import gradio as gr
 from PIL import Image
 import matplotlib.pyplot as plt
@@ -7,61 +8,58 @@ import os
 import numpy as np
 from pycocotools.coco import COCO
 import io
+import hashlib
 
-# ==== Update these paths to match your environment ====
+# ==== Paths (update if needed) ====
 annotation_path = "test/test_annotations.coco.json"
 image_folder = "test/images"
-# ======================================================
+# ==================================
 
-# Load COCO annotations once at start
+# Load COCO annotations once
 coco = COCO(annotation_path)
 categories = coco.loadCats(coco.getCatIds())
 cat_id_to_name = {cat['id']: cat['name'] for cat in categories}
 image_ids = coco.getImgIds()
-image_file_name_map = {coco.loadImgs(img_id)[0]['file_name']: coco.loadImgs(img_id)[0] for img_id in image_ids}
+
+# Helper to compute image content hash (used for matching uploads)
+def compute_md5(image_array):
+    return hashlib.md5(image_array.tobytes()).hexdigest()
+
+# Precompute hashes for all dataset images
+image_hash_map = {}
+for img_id in image_ids:
+    img_info = coco.loadImgs(img_id)[0]
+    img_path = os.path.join(image_folder, img_info['file_name'])
+    if os.path.exists(img_path):
+        img_np = cv2.imread(img_path)
+        img_np = cv2.cvtColor(img_np, cv2.COLOR_BGR2RGB)
+        img_hash = compute_md5(img_np)
+        image_hash_map[img_hash] = img_info
 
 def show_annotated_teeth_image(image_file):
-    # === Handle various image input types ===
-    if isinstance(image_file, str):  # example image from filepath
-        filename = os.path.basename(image_file)
+    # Convert uploaded image to numpy
+    if isinstance(image_file, str):  # example file
         image_np = cv2.imread(image_file)
         image_np = cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB)
-    elif isinstance(image_file, Image.Image):  # uploaded PIL image
-        filename = getattr(image_file, 'filename', None)
+    elif isinstance(image_file, Image.Image):
         image_np = np.array(image_file.convert("RGB"))
-    else:  # raw file upload
-        filename = getattr(image_file, 'name', None)
+    else:
         image_np = cv2.imdecode(np.frombuffer(image_file.read(), np.uint8), cv2.IMREAD_COLOR)
         image_np = cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB)
 
-    filename = os.path.basename(filename) if filename else None
-    print(f"Input filename: {filename}")
-
-    # === Try to match by filename ===
-    matched_info = image_file_name_map.get(filename)
-
-    # === Fallback: Try to match by content ===
-    if not matched_info:
-        print("⚠️ No direct filename match. Trying to match image by content...")
-        for fname, info in image_file_name_map.items():
-            test_img_path = os.path.join(image_folder, fname)
-            test_img = cv2.imread(test_img_path)
-            test_img = cv2.cvtColor(test_img, cv2.COLOR_BGR2RGB)
-
-            # Match by shape and pixel similarity
-            if test_img.shape == image_np.shape and np.allclose(test_img, image_np, atol=5):
-                matched_info = info
-                print(f"✅ Matched by content with: {fname}")
-                break
+    # Compute hash of input image
+    image_hash = compute_md5(image_np)
 
+    matched_info = image_hash_map.get(image_hash)
     if not matched_info:
-        print("❌ No annotation match found. Returning plain image.")
+        print("No matching annotation found for this image.")
         return Image.fromarray(image_np)
 
-    # === Load annotations and draw ===
+    # Load annotations
     ann_ids = coco.getAnnIds(imgIds=matched_info['id'])
     annotations = coco.loadAnns(ann_ids)
 
+    # Plot the image and annotations
     fig, ax = plt.subplots(1, figsize=(8, 8))
     ax.imshow(image_np)
 
@@ -69,8 +67,7 @@ def show_annotated_teeth_image(image_file):
         x, y, w, h = ann['bbox']
         class_id = ann['category_id']
         class_name = cat_id_to_name.get(class_id, "Unknown")
-        rect = patches.Rectangle((x, y), w, h, linewidth=2, edgecolor='blue', facecolor='none')
-        ax.add_patch(rect)
+        ax.add_patch(patches.Rectangle((x, y), w, h, linewidth=2, edgecolor='blue', facecolor='none'))
         ax.text(x, y, class_name, fontsize=10, color='blue', backgroundcolor='black')
 
     ax.axis('off')
@@ -84,16 +81,17 @@ def show_annotated_teeth_image(image_file):
 
     return annotated_image
 
-# === Prepare examples ===
-example_files = list(image_file_name_map.keys())[:3]
-example_paths = [os.path.join(image_folder, fname) for fname in example_files]
+# Example paths for Gradio buttons
+example_files = list(image_hash_map.values())[:3]
+example_paths = [os.path.join(image_folder, img['file_name']) for img in example_files]
 
+# Gradio Interface
 interface = gr.Interface(
     fn=show_annotated_teeth_image,
-    inputs=gr.Image(type="filepath", label="Upload a Dental X-ray or choose an example"),
+    inputs=gr.Image(type="filepath", label="Upload a Dental X-ray or choose a default example"),
     outputs=gr.Image(type="pil", label="Annotated Image"),
     title="🦷 Dental X-ray Annotator",
-    description="Upload a dental X-ray or select a sample to visualize tooth annotations (COCO format).",
+    description="Upload a dental X-ray image or select a default example to visualize teeth annotations using COCO format.",
     examples=example_paths,
 )