Update app.py

app.py

@@ -119,7 +119,13 @@ def create_match_visualization(img_path1, img_path2, kpts0, kpts1, matches, outp
     return output_path
 
 def predict(input_path):
-    if input_path is None: return "Upload an image!", None, None, None
+    # Default return values (empty)
+    default_header = "Please upload an image."
+    default_logs = ["", "", ""]
+    default_imgs = [None, None, None]
+    
+    if input_path is None: 
+        return default_header, default_logs[0], default_imgs[0], default_logs[1], default_imgs[1], default_logs[2], default_imgs[2]
     
     # --- 0. GROUND TRUTH ---
     true_species, true_id = "Unknown", "Unknown"
@@ -159,9 +165,11 @@ def predict(input_path):
     # --- 3. FINE SEARCH & VISUALIZATION LOOP ---
     feats_q = extractor.extract(load_image(input_path).to(DEVICE))
     
-    log = "🔍 **Analysis Process:**\n"
     best_score = -1
     best_candidate_idx = -1
+    
+    # Initialize output lists (size 3)
+    cand_logs = ["Waiting for data...", "Waiting for data...", "Waiting for data..."]
     cand_viz_paths = [None, None, None]
 
     for rank, (idx, arcface_sim) in enumerate(unique_candidates):
@@ -180,9 +188,10 @@ def predict(input_path):
             geo_matches = len(matches["matches"])
             sim_percent = arcface_sim * 100
             
-            log += f"- Candidate {rank+1}: **{species} / {label}**\n"
-            log += f"  • 🧠 Coarse Confidence: **{sim_percent:.1f}%**\n"
-            log += f"  • 📐 Geometric Matches: **{geo_matches}**\n\n"
+            # --- CREATE INDIVIDUAL LOG STRING ---
+            log_str = f"### Candidate {rank+1}: {species} / {label}\n"
+            log_str += f"**🧠 Coarse Confidence:** {sim_percent:.1f}%   |   **📐 Geometric Matches:** {geo_matches}"
+            cand_logs[rank] = log_str
             
             viz_name = f"viz_rank_{rank}"
             viz_path = create_match_visualization(
@@ -197,7 +206,7 @@ def predict(input_path):
                 best_candidate_idx = idx
                 
         except Exception as e:
-            log += f"- Error: {e}\n"
+            cand_logs[rank] = f"Error processing candidate: {e}"
 
     # --- 4. FINAL DECISION ---
     CONFIDENCE_THRESHOLD = 15
@@ -205,61 +214,63 @@ def predict(input_path):
         pred_species = g_species[best_candidate_idx]
         pred_id = g_labels[best_candidate_idx]
         is_correct = (pred_id == true_id)
-        if true_id == "Unknown": header = f"# ❓ MATCH FOUND\n"
-        elif is_correct: header = f"# ✅ CORRECT MATCH!\n"
-        else: header = f"# ❌ INCORRECT MATCH\n"
-        header += f"**Ground Truth:** {true_species} / {true_id}\n"
-        header += f"**Model Prediction:** {pred_species} / {pred_id}\n"
+        
+        if true_id == "Unknown": header = f"### ❓ MATCH FOUND (No Ground Truth)\n"
+        elif is_correct: header = f"### ✅ CORRECT MATCH!\n"
+        else: header = f"### ❌ INCORRECT MATCH\n"
+        
+        header += f"**Ground Truth:** {true_species} / {true_id}   ➡️   **Prediction:** {pred_species} / {pred_id}\n"
         header += f"*(Confirmed with {best_score} geometric keypoints)*"
     else:
-        header = "# ⚠️ UNKNOWN / NO MATCH\n"
+        header = "### ⚠️ UNKNOWN / NO MATCH\n"
         header += f"**Ground Truth:** {true_species} / {true_id}\n"
-        header += f"**Model Prediction:** None\n"
+        header += f"**Prediction:** None (Best match only had {best_score} keypoints)\n"
 
-    return header + "\n\n" + log, cand_viz_paths[0], cand_viz_paths[1], cand_viz_paths[2]
+    # Return: Header, then (Log, Img) for Cand 1, then (Log, Img) for Cand 2, etc.
+    return (header, 
+            cand_logs[0], cand_viz_paths[0],
+            cand_logs[1], cand_viz_paths[1],
+            cand_logs[2], cand_viz_paths[2])
 
-# --- NEW UI SETUP (USING BLOCKS) ---
+# --- 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'))]
 
 with gr.Blocks(title="Wildlife Re-ID: Coarse-to-Fine Demo") as demo:
     gr.Markdown("# Wildlife Re-ID: Coarse-to-Fine Demo")
-    gr.Markdown("Select a test image below. The system finds the Top 3 UNIQUE individuals using embeddings, then verifies them using geometry. Visualizations are limited to 50 lines for clarity.")
+    gr.Markdown("Select a test image. The system finds the Top 3 UNIQUE individuals using embeddings, then verifies them using geometry.")
 
     with gr.Row():
-        # --- Left Column: Input and Examples ---
+        # --- Left Column: Input ---
         with gr.Column(scale=1):
             input_img = gr.Image(type="filepath", label="Test Image", height=300)
-            
-            # EXAMPLES directly under the input image
-            gr.Examples(
-                examples=examples_list,
-                inputs=input_img,
-                label="Test Examples (Click to Select)",
-                examples_per_page=5
-            )
-            
+            gr.Examples(examples=examples_list, inputs=input_img, label="Test Examples", examples_per_page=4)
             submit_btn = gr.Button("Run Identification", variant="primary", size="lg")
 
-        # --- Right Column: Outputs ---
+        # --- Right Column: Vertical Stack of Candidates ---
         with gr.Column(scale=2):
-            output_md = gr.Markdown(label="Evaluation Result")
-            # Use tabs or a row for the 3 candidate images
-            with gr.Row():
-                 with gr.Column():
-                     output_img1 = gr.Image(label="Candidate 1 (Best ArcFace)", type="filepath")
-                 with gr.Column():
-                     output_img2 = gr.Image(label="Candidate 2 (Runner-up)", type="filepath")
-                 with gr.Column():
-                     output_img3 = gr.Image(label="Candidate 3 (Runner-up)", type="filepath")
+            header_md = gr.Markdown(label="Final Decision")
+            
+            # Candidate 1 Group
+            with gr.Group():
+                log1 = gr.Markdown()
+                img1 = gr.Image(label="Visualization", height=450, show_download_button=False)
+            
+            # Candidate 2 Group
+            with gr.Group():
+                log2 = gr.Markdown()
+                img2 = gr.Image(label="Visualization", height=450, show_download_button=False)
+                
+            # Candidate 3 Group
+            with gr.Group():
+                log3 = gr.Markdown()
+                img3 = gr.Image(label="Visualization", height=450, show_download_button=False)
 
-    # connect the button click to the function
     submit_btn.click(
         fn=predict,
         inputs=input_img,
-        outputs=[output_md, output_img1, output_img2, output_img3]
+        outputs=[header_md, log1, img1, log2, img2, log3, img3]
     )
 
-# Fix: Move CSS to launch
 demo.launch(allowed_paths=[TEST_QUERIES_DIR])