Update app.py

app.py

@@ -39,27 +39,18 @@ if os.path.exists(GALLERY_FILE):
 else:
     raise FileNotFoundError("Gallery file missing!")
 
-# --- BUILD GROUND TRUTH LOOKUP (The Fix) ---
-# We map file_size -> (Species, ID) so we can identify images even if Gradio renames them.
+# --- BUILD GROUND TRUTH LOOKUP ---
 GT_LOOKUP = {}
 if os.path.exists(TEST_QUERIES_DIR):
-    print("Building Ground Truth Lookup Table...")
     for f in os.listdir(TEST_QUERIES_DIR):
         if f.lower().endswith(('.jpg', '.png', '.jpeg')):
             full_path = os.path.join(TEST_QUERIES_DIR, f)
             try:
-                # key = file size in bytes
                 f_size = os.path.getsize(full_path)
-                
-                # Parse filename: "Species_ID_QUERY.jpg"
                 parts = f.split("_")
                 if len(parts) >= 2:
-                    species = parts[0]
-                    ind_id = parts[1]
-                    GT_LOOKUP[f_size] = (species, ind_id)
-            except Exception as e:
-                print(f"Skipping {f}: {e}")
-    print(f"Indexed {len(GT_LOOKUP)} test images.")
+                    GT_LOOKUP[f_size] = (parts[0], parts[1])
+            except: pass
 
 # Transform
 transform = transforms.Compose([
@@ -120,24 +111,18 @@ def create_match_visualization(img_path1, img_path2, kpts0, kpts1, matches):
 def predict(input_path):
     if input_path is None: return "Upload an image!", None
     
-    # --- 0. IDENTIFY GROUND TRUTH (ROBUST) ---
-    true_species = "Unknown"
-    true_id = "Unknown"
-    
-    # Method A: Check File Size (Robust against renaming)
+    # --- 0. GROUND TRUTH ---
+    true_species, true_id = "Unknown", "Unknown"
     try:
         input_size = os.path.getsize(input_path)
         if input_size in GT_LOOKUP:
             true_species, true_id = GT_LOOKUP[input_size]
         else:
-            # Method B: Fallback to filename parsing
             filename = os.path.basename(input_path)
             if "_QUERY" in filename:
                 parts = filename.split("_")
-                true_species = parts[0]
-                true_id = parts[1]
-    except:
-        pass
+                true_species, true_id = parts[0], parts[1]
+    except: pass
 
     # Load Image
     input_image = Image.open(input_path).convert("RGB")
@@ -147,18 +132,38 @@ def predict(input_path):
     with torch.no_grad():
         q_emb = torch.nn.functional.normalize(model(img_t), p=2, dim=1)
         scores = torch.mm(q_emb, g_embeddings.t())
-        top_scores, top_indices = torch.topk(scores, k=min(3, len(g_paths)))
+        
+        # NEW: Fetch top 50 matches so we can filter duplicates
+        # We need enough candidates to find 3 unique individuals
+        top_scores, top_indices = torch.topk(scores, k=min(50, len(g_paths)))
 
-    # --- 2. FINE SEARCH (LightGlue) ---
+    # --- 2. UNIQUE CANDIDATE FILTERING ---
+    unique_candidates = []
+    seen_individuals = set()
+    
+    # Loop through results until we have 3 unique people
+    for i in range(len(top_indices[0])):
+        if len(unique_candidates) >= 3: break
+        
+        idx = top_indices[0][i].item()
+        score = top_scores[0][i].item() # This is the ArcFace Similarity!
+        label = g_labels[idx]
+        
+        # If we haven't seen this individual yet, add them!
+        if label not in seen_individuals:
+            seen_individuals.add(label)
+            unique_candidates.append((idx, score))
+
+    # --- 3. FINE SEARCH (LightGlue) ---
     feats_q = extractor.extract(load_image(input_path).to(DEVICE))
     
     log = "🔍 **Analysis Process:**\n"
     best_score = -1
-    best_idx = -1
+    best_candidate_idx = -1
     best_matches_info = None
 
-    for rank, idx in enumerate(top_indices[0]):
-        idx = idx.item()
+    # Process only our 3 UNIQUE candidates
+    for rank, (idx, arcface_sim) in enumerate(unique_candidates):
         path = g_paths[idx]
         label = g_labels[idx]
         species = g_species[idx]
@@ -173,29 +178,33 @@ def predict(input_path):
             
             geo_matches = len(matches["matches"])
             
-            log += f"- Candidate {rank+1}: **{species} / {label}** | Keypoints: {geo_matches}\n"
+            # Format ArcFace score as percentage
+            sim_percent = arcface_sim * 100
+            
+            log += f"- Candidate {rank+1}: **{species} / {label}**\n"
+            log += f"  • 🧠 Coarse Confidence (ArcFace): **{sim_percent:.1f}%**\n"
+            log += f"  • 📐 Geometric Matches (LightGlue): **{geo_matches}**\n\n"
             
+            # We still pick the winner based on LightGlue (Fine-grained)
             if geo_matches > best_score:
                 best_score = geo_matches
-                best_idx = idx
+                best_candidate_idx = idx
                 best_matches_info = (path, feats_q['keypoints'], feats_c['keypoints'], matches['matches'])
                 
         except Exception as e:
             log += f"- Error: {e}\n"
 
-    # --- 3. FINAL DECISION ---
+    # --- 4. FINAL DECISION ---
     CONFIDENCE_THRESHOLD = 15
+    viz_path = None
     
-    if best_idx != -1 and best_score > CONFIDENCE_THRESHOLD:
-        pred_species = g_species[best_idx]
-        pred_id = g_labels[best_idx]
+    if best_candidate_idx != -1 and best_score > CONFIDENCE_THRESHOLD:
+        pred_species = g_species[best_candidate_idx]
+        pred_id = g_labels[best_candidate_idx]
         
-        # Check correctness
         is_correct = (pred_id == true_id)
-        
-        # Logic: If Truth is Unknown, we can't say it's Incorrect.
         if true_id == "Unknown":
-            header = f"# ❓ MATCH FOUND (No Ground Truth)\n"
+            header = f"# ❓ MATCH FOUND\n"
         elif is_correct:
             header = f"# ✅ CORRECT MATCH!\n"
         else:
@@ -203,7 +212,7 @@ def predict(input_path):
             
         header += f"**Ground Truth:** {true_species} / {true_id}\n"
         header += f"**Model Prediction:** {pred_species} / {pred_id}\n"
-        header += f"*(Confidence: {best_score} geometric keypoints)*"
+        header += f"*(Confirmed with {best_score} geometric keypoints)*"
         
         winner_path, kpts0, kpts1, matches = best_matches_info
         viz_path = create_match_visualization(input_path, winner_path, kpts0, kpts1, matches)
@@ -211,12 +220,10 @@ def predict(input_path):
     else:
         header = "# ⚠️ UNKNOWN / NO MATCH\n"
         header += f"**Ground Truth:** {true_species} / {true_id}\n"
-        header += f"**Model Prediction:** None (Best candidate only had {best_score} matches)\n"
-        
+        header += f"**Model Prediction:** None\n"
         if true_id == "Unknown" or true_species != "Unknown":
-             header += "\n*Model correctly rejected a non-matching image!*"
-        viz_path = None
-        
+             header += "\n*Model correctly rejected non-matching candidates.*"
+
     return header + "\n\n" + log, viz_path
 
 # --- UI SETUP ---
@@ -232,7 +239,7 @@ iface = gr.Interface(
         gr.Image(label="Visualization")
     ],
     title="Wildlife Re-ID: Coarse-to-Fine Demo",
-    description="Select a test image. The system will reveal the Ground Truth (hidden in filename) and compare it with the Model's Prediction.",
+    description="Select a test image. The system finds the Top 3 UNIQUE individuals using embeddings, then verifies the best match using geometry.",
     examples=examples_list,
     cache_examples=False
 )