Update lines and patch matching

main.py

@@ -183,13 +183,14 @@ def make_correspondence_figure(image_a, image_b, patches_a, patches_b, grid_h, g
     gap = 30
     canvas_w = img_w * 2 + gap
     fig, ax = plt.subplots(1, 1, figsize=(14, 6))
-    fig.patch.set_facecolor('#0f172a') # Slate 900 for dark mode frontend
+    fig.patch.set_facecolor('#0f172a') # Slate 900 matching React frontend
 
-    canvas = Image.new("RGB", (canvas_w, img_h), (15, 23, 42))
+    canvas = Image.new("RGB", (canvas_w, img_h), (15, 23, 42)) # Slate 950 canvas
     canvas.paste(img_a_resized, (0, 0))
     canvas.paste(img_b_resized, (img_w + gap, 0))
     ax.imshow(canvas)
 
+    # Keep the clean, unified 'spring' colormap
     cmap = plt.cm.get_cmap("spring", max(len(matches), 1))
 
     for i, (idx_a, idx_b, score) in enumerate(matches):
@@ -198,9 +199,27 @@ def make_correspondence_figure(image_a, image_b, patches_a, patches_b, grid_h, g
         xb_canvas = xb + img_w + gap
         color = cmap(i % 20)
 
-        ax.plot([xa, xb_canvas], [ya, yb], color=color, linewidth=2, alpha=0.9)
+        # Draw the clean lines and dots
+        ax.plot([xa, xb_canvas], [ya, yb], color=color, linewidth=2, alpha=0.8)
         ax.scatter([xa, xb_canvas], [ya, yb], color=color, s=50, zorder=5, edgecolors="white", linewidths=0.5)
 
+        # --- THE NEW CLEAN SCORE BADGES ---
+        # Calculate the exact midpoint of the line
+        mx = (xa + xb_canvas) / 2
+        my = (ya + yb) / 2
+        
+        # Style a sleek dark pill with a glowing border matching the line color
+        bbox_props = dict(
+            boxstyle="round,pad=0.25", 
+            fc="#1e293b",     # Slate 800 background
+            ec=color,         # Border matches the line color
+            alpha=0.95,       # High opacity so it covers the line behind it
+            lw=1              # Thin, crisp border
+        )
+        
+        ax.text(mx, my, f"{score:.2f}", fontsize=8, color="white", fontweight="bold",
+                ha="center", va="center", zorder=10, bbox=bbox_props)
+
     ax.axis("off")
     fig.tight_layout(pad=0)
     return fig, len(matches)
@@ -253,10 +272,40 @@ async def analyze_artworks(file_a: UploadFile = File(...), file_b: UploadFile =
         cls_g_a, patches_g_a, _, _ = extract_dino_features(img_a, "grayscale")
         cls_g_b, patches_g_b, _, _ = extract_dino_features(img_b, "grayscale")
         
+        cls_e_a, patches_e_a, _, _ = extract_dino_features(img_a, "edges")
+        cls_e_b, patches_e_b, _, _ = extract_dino_features(img_b, "edges")
+        
         _, a2b_c, b2a_c = compute_patch_matches(patches_c_a, patches_c_b)
         _, a2b_g, b2a_g = compute_patch_matches(patches_g_a, patches_g_b)
         _, a2b_e, b2a_e = compute_patch_matches(patches_e_a, patches_e_b)
 
+        # 1. Combine domains to defeat color-shifting (used for Heatmap stats)
+        a2b_best = torch.max(torch.max(a2b_c, a2b_g), a2b_e)
+        b2a_best = torch.max(torch.max(b2a_c, b2a_g), b2a_e)
+        
+        corr_thresh = (a2b_best.mean() + 0.5 * a2b_best.std()).item()
+        corr_thresh = min(max(corr_thresh, 0.4), 0.75)
+        
+        # 2. THE FIX: Dynamically select the best patch domain for the visual lines
+        mode_scores = {
+            "color": a2b_c.mean().item(),
+            "grayscale": a2b_g.mean().item(),
+            "edges": a2b_e.mean().item()
+        }
+        best_mode = max(mode_scores, key=mode_scores.get)
+        mode_patches = {
+            "color": (patches_c_a, patches_c_b),
+            "grayscale": (patches_g_a, patches_g_b),
+            "edges": (patches_e_a, patches_e_b),
+        }
+        corr_pa, corr_pb = mode_patches[best_mode]
+
+        # 3. Pass the dynamically selected patches to the figure generator
+        corr_fig, match_count = make_correspondence_figure(
+            img_a, img_b, corr_pa, corr_pb, gh, gw, score_thresh=corr_thresh
+        )
+        correspondence_map_b64 = fig_to_base64(corr_fig)
+
         # The missing magic: Combine domains to defeat color-shifting
         a2b_best = torch.max(torch.max(a2b_c, a2b_g), a2b_e)
         b2a_best = torch.max(torch.max(b2a_c, b2a_g), b2a_e)