Update to app.py

app.py

@@ -14,7 +14,7 @@ from lightglue.utils import load_image, rbd
 # --- CONFIG ---
 WEIGHTS_PATH = "MiewID_ArcFace_FineTun.pth"
 GALLERY_FILE = "mini_gallery.pt"
-TEST_QUERIES_DIR = "test_queries"  # Folder containing your example images
+TEST_QUERIES_DIR = "test_queries"
 IMG_SIZE = 384
 EMBEDDING_DIM = 512
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -48,6 +48,173 @@ transform = transforms.Compose([
     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
 ])
 
+def create_match_visualization(img_path1, img_path2, kpts0, kpts1, matches):
+    """
+    Creates a side-by-side visualization with lines connecting matched keypoints.
+    """
+    # FIX: Remove batch dimension if present (Shape [1, N, 2] -> [N, 2])
+    if kpts0.dim() == 3: kpts0 = kpts0[0]
+    if kpts1.dim() == 3: kpts1 = kpts1[0]
+
+    # Load images
+    img1 = Image.open(img_path1).convert("RGB")
+    img2 = Image.open(img_path2).convert("RGB")
+    
+    # Resize for display consistency (optional, but helps layout)
+    # We resize primarily to make the plot look neat side-by-side
+    target_h = 400
+    w1, h1 = img1.size
+    w2, h2 = img2.size
+    scale1 = target_h / h1
+    scale2 = target_h / h2
+    
+    img1 = img1.resize((int(w1 * scale1), target_h))
+    img2 = img2.resize((int(w2 * scale2), target_h))
+    
+    # Create plot
+    fig, ax = plt.subplots(1, 1, figsize=(10, 5))
+    ax.axis('off')
+    
+    # Concatenate images side-by-side
+    img1_np = np.array(img1)
+    img2_np = np.array(img2)
+    h1_new, w1_new, _ = img1_np.shape
+    h2_new, w2_new, _ = img2_np.shape
+    
+    # Create canvas
+    width = w1_new + w2_new
+    canvas = np.zeros((target_h, width, 3), dtype=np.uint8)
+    canvas[:, :w1_new, :] = img1_np
+    canvas[:, w1_new:, :] = img2_np
+    
+    ax.imshow(canvas)
+    
+    # Extract matched points
+    # Now that we squeezed dimensions at the top, this indexing works!
+    m_kpts0 = kpts0[matches[..., 0]].cpu().numpy()
+    m_kpts1 = kpts1[matches[..., 1]].cpu().numpy()
+    
+    # Scale points to match the resized images we are displaying
+    m_kpts0[:, 0] *= scale1
+    m_kpts0[:, 1] *= scale1
+    m_kpts1[:, 0] *= scale2
+    m_kpts1[:, 1] *= scale2
+    
+    # Plot lines
+    # Shift x-coordinates of the second image by the width of the first
+    for (x0, y0), (x1, y1) in zip(m_kpts0, m_kpts1):
+        ax.plot([x0, x1 + w1_new], [y0, y1], color="lime", linewidth=0.5, alpha=0.7)
+        ax.scatter([x0, x1 + w1_new], [y0, y1], color="lime", s=2)
+
+    plt.tight_layout()
+    
+    # Save to buffer
+    output_path = "match_viz_result.jpg"
+    plt.savefig(output_path, bbox_inches='tight', pad_inches=0)
+    plt.close()
+    
+    return output_path
+
+def predict(input_image):
+    if input_image is None: return "Upload an image!", None
+    
+    # Save input temporarily for LightGlue
+    input_path = "temp_query.jpg"
+    input_image.save(input_path)
+
+    # 1. COARSE SEARCH (ArcFace)
+    img_t = transform(input_image).unsqueeze(0).to(DEVICE)
+    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())
+        # Get Top 3
+        top_scores, top_indices = torch.topk(scores, k=min(3, len(g_paths)))
+
+    # 2. FINE SEARCH (LightGlue)
+    feats_q = extractor.extract(load_image(input_path).to(DEVICE))
+    
+    candidates = []
+    log = "🔍 **Analysis Process:**\n"
+    
+    best_score = -1
+    best_idx = -1
+    best_matches_info = None
+
+    for rank, idx in enumerate(top_indices[0]):
+        idx = idx.item()
+        path = g_paths[idx]
+        label = g_labels[idx]
+        species = g_species[idx]
+        
+        try:
+            if not os.path.exists(path):
+                log += f"- ⚠️ File missing: {path}\n"
+                continue
+
+            feats_c = extractor.extract(load_image(path).to(DEVICE))
+            with torch.no_grad():
+                matches = matcher({"image0": feats_q, "image1": feats_c})
+                matches = rbd(matches) # remove batch dim
+            
+            geo_matches = len(matches["matches"])
+            
+            log += f"- Candidate {rank+1}: **{species} / {label}** | Keypoints: {geo_matches}\n"
+            
+            if geo_matches > best_score:
+                best_score = geo_matches
+                best_idx = idx
+                # Save visualization data
+                best_matches_info = (path, feats_q['keypoints'], feats_c['keypoints'], matches['matches'])
+                
+        except Exception as e:
+            log += f"- Error processing {path}: {e}\n"
+
+    # 3. DECISION & VISUALIZATION
+    if best_idx != -1:
+        winner_path, kpts0, kpts1, matches = best_matches_info
+        winner_label = g_labels[best_idx]
+        winner_species = g_species[best_idx]
+        
+        # GENERATE VISUALIZATION
+        viz_image_path = create_match_visualization(input_path, winner_path, kpts0, kpts1, matches)
+        
+        result = f"✅ **MATCH FOUND**\nSpecies: {winner_species}\nIndividual: {winner_label}\n(Verified with {best_score} geometric matches)"
+        return result + "\n\n" + log, viz_image_path
+    else:
+        return "⚠️ No confident match found.\n\n" + log, None
+
+# --- UI SETUP ---
+examples_list = []
+if os.path.exists(TEST_QUERIES_DIR):
+    examples_list = [os.path.join(TEST_QUERIES_DIR, f) for f in os.listdir(TEST_QUERIES_DIR) if f.lower().endswith(('.jpg', '.png'))]
+
+iface = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil", label="Test Image"),
+    outputs=[
+        gr.Markdown(label="Result"),
+        gr.Image(label="Feature Matching Visualization")
+    ],
+    title="Wildlife Re-ID: Coarse-to-Fine Demo",
+    description="Click a test image below to run the identification.",
+    examples=examples_list,
+    cache_examples=False
+)
+
+iface.launch()    g_embeddings = data["embeddings"].to(DEVICE)
+    g_paths = data["paths"]
+    g_labels = data["labels"]
+    g_species = data["species"]
+else:
+    raise FileNotFoundError("Gallery file missing!")
+
+# Transform
+transform = transforms.Compose([
+    transforms.Resize((IMG_SIZE, IMG_SIZE)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+])
+
 def create_match_visualization(img_path1, img_path2, kpts0, kpts1, matches):
     """
     Creates a side-by-side visualization with lines connecting matched keypoints.