Update main.py

main.py

@@ -4,17 +4,23 @@ import torch
 import torch.nn.functional as F
 import cv2
 import numpy as np
-from PIL import Image, ImageDraw
+from PIL import Image
 import io
 import base64
 from transformers import AutoImageProcessor, AutoModel, CLIPProcessor, CLIPModel
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import matplotlib.gridspec as gridspec
+from scipy.optimize import linear_sum_assignment
+import lpips
+import torchvision.transforms as transforms
 
 # ==============================================================================
 # 1. Initialize FastAPI & CORS
 # ==============================================================================
-app = FastAPI(title="Copyright Diagnostic API")
+app = FastAPI(title="Copyright Diagnostic API - 3 Pillar XAI")
 
-# Crucial for allowing your React frontend to communicate with this backend
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"], 
@@ -28,192 +34,279 @@ app.add_middleware(
 # ==============================================================================
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-print("Loading DINOv2...")
-dino_processor = AutoImageProcessor.from_pretrained("facebook/dinov2-base")
-dino_model = AutoModel.from_pretrained("facebook/dinov2-base").to(device)
+DINO_MODEL_ID = "facebook/dinov2-large"
+CLIP_MODEL_ID = "openai/clip-vit-large-patch14"
+PATCH_SIZE = 14
+
+print(f"Loading DINOv2 ({DINO_MODEL_ID})...")
+dino_processor = AutoImageProcessor.from_pretrained(DINO_MODEL_ID)
+dino_model = AutoModel.from_pretrained(DINO_MODEL_ID).to(device)
 dino_model.eval()
 
-print("Loading CLIP...")
-clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-large-patch14")
-clip_model = CLIPModel.from_pretrained("openai/clip-vit-large-patch14").to(device)
+print(f"Loading CLIP ({CLIP_MODEL_ID})...")
+clip_processor = CLIPProcessor.from_pretrained(CLIP_MODEL_ID)
+clip_model = CLIPModel.from_pretrained(CLIP_MODEL_ID).to(device)
 clip_model.eval()
 
-# Helper function to convert PIL Image to Base64 string for React
+print("Loading LPIPS (AlexNet)...")
+loss_fn_alex = lpips.LPIPS(net='alex').to(device)
+loss_fn_alex.eval()
+
+# LPIPS requires images normalized between [-1, 1]
+lpips_transform = transforms.Compose([
+    transforms.Resize((256, 256)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
+])
+
+# ==============================================================================
+# 3. Helpers & Base64 Converters
+# ==============================================================================
+
 def image_to_base64(img: Image.Image) -> str:
     buffered = io.BytesIO()
     img.save(buffered, format="JPEG")
     return base64.b64encode(buffered.getvalue()).decode("utf-8")
 
+def fig_to_base64(fig) -> str:
+    buf = io.BytesIO()
+    fig.savefig(buf, format="jpg", bbox_inches='tight', pad_inches=0.1, dpi=100)
+    plt.close(fig)
+    return base64.b64encode(buf.getvalue()).decode("utf-8")
+
+def to_edge_map(image: Image.Image) -> Image.Image:
+    """Strips style/texture, leaving structural contours for Pillar 1."""
+    img_array = np.array(image.convert("L"))
+    median = np.median(img_array)
+    low = int(max(0, 0.5 * median))
+    high = int(min(255, 1.3 * median))
+    edges = cv2.Canny(img_array, low, high)
+    kernel = np.ones((2, 2), np.uint8)
+    edges = cv2.dilate(edges, kernel, iterations=1)
+    return Image.fromarray(edges).convert("RGB")
+
+def patch_idx_to_xy(idx, grid_w, patch_size):
+    row = idx // grid_w
+    col = idx % grid_w
+    return col * patch_size + patch_size / 2, row * patch_size + patch_size / 2
+
 # ==============================================================================
-# 3. AI Pipeline Functions
+# 4. AI Pipeline Functions (Mapped to the 3 Pillars)
 # ==============================================================================
 
-def compute_semantic_similarity(img_a: Image.Image, img_b: Image.Image) -> float:
-    inputs = clip_processor(images=[img_a, img_b], return_tensors="pt").to(device)
-    
+def extract_dino_features(image: Image.Image, preprocess="color"):
+    if preprocess == "grayscale":
+        img = image.convert("L").convert("RGB")
+    elif preprocess == "edges":
+        img = to_edge_map(image)
+    else:
+        img = image
+
+    inputs = dino_processor(images=img, return_tensors="pt").to(device)
     with torch.no_grad():
-        image_features = clip_model.get_image_features(**inputs)
-        
-    # --- FIX: Handle Transformers 5.x object returns ---
-    if not isinstance(image_features, torch.Tensor):
-        if hasattr(image_features, "image_embeds"):
-            image_features = image_features.image_embeds
-        elif hasattr(image_features, "pooler_output"):
-            image_features = image_features.pooler_output
-        else:
-            image_features = image_features[1] if isinstance(image_features, tuple) and len(image_features) > 1 else image_features[0]
-            
-    image_features = F.normalize(image_features, p=2, dim=-1)
-    score = F.cosine_similarity(image_features[0].unsqueeze(0), image_features[1].unsqueeze(0))
-    return round(score.item(), 4)
-
-def compute_structural_similarity(img_a: Image.Image, img_b: Image.Image):
-    # Convert to grayscale numpy arrays
-    arr_a = np.array(img_a.convert('L'))
-    arr_b = np.array(img_b.convert('L'))
-    
-    # 1. Aggressive Gaussian Blur to eliminate stylistic texture
-    # A kernel size of (11, 11) is strong enough to blur out noise but keep main shapes
-    blur_a = cv2.GaussianBlur(arr_a, (11, 11), 0)
-    blur_b = cv2.GaussianBlur(arr_b, (11, 11), 0)
-    
-    # 2. Dynamic Auto-Canny Helper Function
-    def auto_canny(image, sigma=0.33):
-        v = np.median(image)
-        lower = int(max(0, (1.0 - sigma) * v))
-        upper = int(min(255, (1.0 + sigma) * v))
-        return cv2.Canny(image, lower, upper)
-        
-    # Extract structural edges using the blurred images and dynamic thresholds
-    edges_a = auto_canny(blur_a)
-    edges_b = auto_canny(blur_b)
-    
-    # 3. Dilate the edges to give them a "margin of error" for spatial overlap
-    kernel = np.ones((7,7), np.uint8) # Thicker kernel for better overlap
-    edges_a_thick = cv2.dilate(edges_a, kernel, iterations=1)
-    edges_b_thick = cv2.dilate(edges_b, kernel, iterations=1)
-    
-    # Resize B to match A for matrix math
-    edges_b_resized = cv2.resize(edges_b_thick, (edges_a_thick.shape[1], edges_a_thick.shape[0]))
-    
-    # Calculate Intersection over Union (IoU)
-    intersection = np.logical_and(edges_a_thick > 0, edges_b_resized > 0).sum()
-    union = np.logical_or(edges_a_thick > 0, edges_b_resized > 0).sum()
-    
-    iou_score = intersection / union if union != 0 else 0.0
-    
-    # Return the clean (non-thickened) edges for a prettier UI visualization
-    return round(iou_score, 4), Image.fromarray(edges_a), Image.fromarray(edges_b_resized)
-
-def compute_patch_similarity(img_a: Image.Image, img_b: Image.Image):
-    target_size = (518, 518)
-    img_a_resized = img_a.resize(target_size)
-    img_b_resized = img_b.resize(target_size)
-    
-    inputs_a = dino_processor(
-        images=img_a_resized, 
-        return_tensors="pt", 
-        do_resize=False, 
-        do_center_crop=False  # <--- ADD THIS
-    ).to(device)
-    
-    inputs_b = dino_processor(
-        images=img_b_resized, 
-        return_tensors="pt", 
-        do_resize=False, 
-        do_center_crop=False  # <--- ADD THIS
-    ).to(device)
-    
+        outputs = dino_model(**inputs)
+
+    cls_token = outputs.last_hidden_state[:, 0, :]
+    patch_tokens = outputs.last_hidden_state[:, 1:, :]
+    n_patches = patch_tokens.shape[1]
+    grid_size = int(n_patches ** 0.5)
+    return cls_token, patch_tokens, grid_size, grid_size
+
+def extract_clip_similarity(image_a: Image.Image, image_b: Image.Image):
+    inputs = clip_processor(images=[image_a, image_b], return_tensors="pt").to(device)
     with torch.no_grad():
-        out_a = dino_model(**inputs_a)
-        out_b = dino_model(**inputs_b)
-        
-    emb_a = F.normalize(out_a.last_hidden_state[:, 1:, :].squeeze(0), p=2, dim=-1)
-    emb_b = F.normalize(out_b.last_hidden_state[:, 1:, :].squeeze(0), p=2, dim=-1)
-    
-    sim_matrix = torch.matmul(emb_a, emb_b.T) 
-    
-    best_b_for_a = torch.argmax(sim_matrix, dim=1)
-    best_a_for_b = torch.argmax(sim_matrix, dim=0)
-    
-    matches = []
-    SIMILARITY_THRESHOLD = 0.87 
-    
-    for a_idx in range(len(best_b_for_a)):
-        b_idx = best_b_for_a[a_idx]
-        if best_a_for_b[b_idx] == a_idx: 
-            score = sim_matrix[a_idx, b_idx].item()
-            if score >= SIMILARITY_THRESHOLD:
-                matches.append((a_idx, b_idx, score))
-
-
-    patch_size = 14   
-    grid_size = target_size[0] // patch_size
-    total_patches = grid_size * grid_size
-    patch_score = len(matches) / float(total_patches)
-    
-    combined_vis = Image.new('RGB', (target_size[0] * 2, target_size[1]))
-    combined_vis.paste(img_a_resized, (0, 0))
-    combined_vis.paste(img_b_resized, (target_size[0], 0))
-    draw = ImageDraw.Draw(combined_vis)
-    
-    
-    
-    
-    for a_idx, b_idx, score in matches:
-        ay = (a_idx // grid_size) * patch_size
-        ax = (a_idx % grid_size) * patch_size
-        by = (b_idx // grid_size) * patch_size
-        bx = (b_idx % grid_size) * patch_size + target_size[0]
+        outputs = clip_model.get_image_features(**inputs)
         
-        draw.rectangle([ax, ay, ax + patch_size, ay + patch_size], outline="red", width=2)
-        draw.rectangle([bx, by, bx + patch_size, by + patch_size], outline="red", width=2)
+    if hasattr(outputs, 'image_embeds'):
+        features = outputs.image_embeds
+    elif isinstance(outputs, torch.Tensor):
+        features = outputs
+    else:
+        features = outputs.pooler_output
         
-        center_a = (ax + patch_size // 2, ay + patch_size // 2)
-        center_b = (bx + patch_size // 2, by + patch_size // 2)
-        draw.line([center_a, center_b], fill="lime", width=1)
-        
-    return round(patch_score, 4), combined_vis
+    features = F.normalize(features, dim=-1)
+    return round((features[0] @ features[1]).item(), 4)
+
+def compute_patch_matches(patches_a, patches_b):
+    pa = F.normalize(patches_a.squeeze(0), dim=-1)
+    pb = F.normalize(patches_b.squeeze(0), dim=-1)
+    sim_matrix = pa @ pb.T
+    a_to_b_scores = sim_matrix.max(dim=1).values
+    b_to_a_scores = sim_matrix.max(dim=0).values
+    return sim_matrix, a_to_b_scores, b_to_a_scores
+
+def nms_matches(matches, grid_w, patch_size, radius=2.0):
+    if not matches: return []
+    matches = sorted(matches, key=lambda m: m[2], reverse=True)
+    kept = []
+    pixel_radius = radius * patch_size
+
+    for idx_a, idx_b, score in matches:
+        xa, ya = patch_idx_to_xy(idx_a, grid_w, patch_size)
+        xb, yb = patch_idx_to_xy(idx_b, grid_w, patch_size)
+        dominated = False
+        for ka, kb, ks in kept:
+            kxa, kya = patch_idx_to_xy(ka, grid_w, patch_size)
+            kxb, kyb = patch_idx_to_xy(kb, grid_w, patch_size)
+            if (abs(xa - kxa) < pixel_radius and abs(ya - kya) < pixel_radius) or \
+               (abs(xb - kxb) < pixel_radius and abs(yb - kyb) < pixel_radius):
+                dominated = True
+                break
+        if not dominated:
+            kept.append((idx_a, idx_b, score))
+    return kept
+
+def make_correspondence_figure(image_a, image_b, patches_a, patches_b, grid_h, grid_w, max_matches=20, score_thresh=0.5):
+    pa = patches_a.squeeze(0).cpu()
+    pb = patches_b.squeeze(0).cpu()
+
+    pa_norm = F.normalize(pa, dim=-1)
+    pb_norm = F.normalize(pb, dim=-1)
+
+    sim = (pa_norm @ pb_norm.T).numpy()
+    cost = 1.0 - sim
+
+    row_ind, col_ind = linear_sum_assignment(cost)
+
+    raw_matches = []
+    for r, c in zip(row_ind, col_ind):
+        score = sim[r, c]
+        if score >= score_thresh:
+            raw_matches.append((r, c, float(score)))
+
+    matches = nms_matches(raw_matches, grid_w, PATCH_SIZE, radius=2.0)[:max_matches]
+
+    img_w, img_h = grid_w * PATCH_SIZE, grid_h * PATCH_SIZE
+    img_a_resized = image_a.resize((img_w, img_h))
+    img_b_resized = image_b.resize((img_w, img_h))
+
+    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
+
+    canvas = Image.new("RGB", (canvas_w, img_h), (15, 23, 42))
+    canvas.paste(img_a_resized, (0, 0))
+    canvas.paste(img_b_resized, (img_w + gap, 0))
+    ax.imshow(canvas)
+
+    cmap = plt.cm.get_cmap("spring", max(len(matches), 1))
+
+    for i, (idx_a, idx_b, score) in enumerate(matches):
+        xa, ya = patch_idx_to_xy(idx_a, grid_w, PATCH_SIZE)
+        xb, yb = patch_idx_to_xy(idx_b, grid_w, PATCH_SIZE)
+        xb_canvas = xb + img_w + gap
+        color = cmap(i % 20)
+
+        ax.plot([xa, xb_canvas], [ya, yb], color=color, linewidth=2, alpha=0.9)
+        ax.scatter([xa, xb_canvas], [ya, yb], color=color, s=50, zorder=5, edgecolors="white", linewidths=0.5)
+
+    ax.axis("off")
+    fig.tight_layout(pad=0)
+    return fig, len(matches)
+
+def make_combined_figure(image_a, image_b, scores_a, scores_b, grid_h, grid_w):
+    heatmap_a = scores_a.reshape(grid_h, grid_w).cpu().numpy()
+    heatmap_b = scores_b.reshape(grid_h, grid_w).cpu().numpy()
+
+    fig = plt.figure(figsize=(12, 5.5))
+    fig.patch.set_facecolor('#0f172a') # Dark mode mapping
+    gs = gridspec.GridSpec(1, 2, wspace=0.05)
+
+    ax0 = fig.add_subplot(gs[0])
+    ax0.imshow(image_a.resize((grid_w * PATCH_SIZE, grid_h * PATCH_SIZE)))
+    ax0.imshow(heatmap_a, cmap="inferno", alpha=0.6, interpolation="bilinear", extent=(0, grid_w * PATCH_SIZE, grid_h * PATCH_SIZE, 0))
+    ax0.axis("off")
+
+    ax1 = fig.add_subplot(gs[1])
+    ax1.imshow(image_b.resize((grid_w * PATCH_SIZE, grid_h * PATCH_SIZE)))
+    ax1.imshow(heatmap_b, cmap="inferno", alpha=0.6, interpolation="bilinear", extent=(0, grid_w * PATCH_SIZE, grid_h * PATCH_SIZE, 0))
+    ax1.axis("off")
+
+    fig.tight_layout(pad=0)
+    return fig
 
 # ==============================================================================
-# 4. API Endpoints
+# 5. Primary Analysis Endpoint
 # ==============================================================================
 
 @app.post("/analyze")
 async def analyze_artworks(file_a: UploadFile = File(...), file_b: UploadFile = File(...)):
     try:
-        # Read uploaded files into PIL Images
         img_a = Image.open(io.BytesIO(await file_a.read())).convert("RGB")
         img_b = Image.open(io.BytesIO(await file_b.read())).convert("RGB")
         
-        # Run pipelines 
-        semantic_score = compute_semantic_similarity(img_a, img_b)
-        struct_score, edge_a, edge_b = compute_structural_similarity(img_a, img_b)
-        patch_score, patch_vis = compute_patch_similarity(img_a, img_b)
+        # --- PILLAR 1: Idea-Expression Dichotomy ---
+        semantic_idea_score = extract_clip_similarity(img_a, img_b)
+        
+        cls_e_a, patches_e_a, gh, gw = extract_dino_features(img_a, "edges")
+        cls_e_b, patches_e_b, _, _ = extract_dino_features(img_b, "edges")
+        structural_expression_score = round(F.cosine_similarity(cls_e_a, cls_e_b).item(), 4)
+        
+        edge_a_b64 = image_to_base64(to_edge_map(img_a))
+        edge_b_b64 = image_to_base64(to_edge_map(img_b))
+
+        # --- PILLAR 2: Fragmented Literal Similarity (RESTored BEST-OF FUSION) ---
+        cls_c_a, patches_c_a, _, _ = extract_dino_features(img_a, "color")
+        cls_c_b, patches_c_b, _, _ = extract_dino_features(img_b, "color")
+        
+        cls_g_a, patches_g_a, _, _ = extract_dino_features(img_a, "grayscale")
+        cls_g_b, patches_g_b, _, _ = extract_dino_features(img_b, "grayscale")
+        
+        _, 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)
+
+        # 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)
+        
+        corr_thresh = (a2b_best.mean() + 0.5 * a2b_best.std()).item()
+        corr_thresh = min(max(corr_thresh, 0.4), 0.75)
+        
+        # We pass patches_c_a for the visual lines, but the scoring uses the fused best-of logic internally
+        corr_fig, match_count = make_correspondence_figure(img_a, img_b, patches_c_a, patches_c_b, gh, gw, score_thresh=corr_thresh)
+        correspondence_map_b64 = fig_to_base64(corr_fig)
+
+        # Restored Statistics
+        n_patches_a = a2b_best.shape[0]
+        adaptive_thresh = max((a2b_best.mean() + a2b_best.std()).item(), 0.5)
+        high_a = (a2b_best > adaptive_thresh).sum().item()
+        pct_copied = round((high_a / n_patches_a) * 100, 1)
+
+        # --- PILLAR 3: Substantial Similarity ---
+        heatmap_fig = make_combined_figure(img_a, img_b, a2b_best, b2a_best, gh, gw)
+        heatmap_b64 = fig_to_base64(heatmap_fig)
         
-        # Convert images to base64 for JSON transmission
-        patch_vis_b64 = image_to_base64(patch_vis) 
-        edge_a_b64 = image_to_base64(edge_a)    
-        edge_b_b64 = image_to_base64(edge_b)    
+        t_a = lpips_transform(img_a).unsqueeze(0).to(device)
+        t_b = lpips_transform(img_b).unsqueeze(0).to(device)
+        with torch.no_grad():
+            lpips_distance = round(loss_fn_alex(t_a, t_b).item(), 4)
 
-        # Return a clean JSON package to React
         return {
             "status": "success",
-            "scores": {
-                "semantic_idea": semantic_score,
-                "structural_layout": struct_score,
-                "fragmented_literal": patch_score
+            "pillar_1_idea_expression": {
+                "semantic_idea_score": semantic_idea_score,
+                "structural_expression_score": structural_expression_score,
+                "edge_map_a_b64": f"data:image/jpeg;base64,{edge_a_b64}",
+                "edge_map_b_b64": f"data:image/jpeg;base64,{edge_b_b64}"
+            },
+            "pillar_2_fragmented_literal": {
+                "patch_match_count": match_count,
+                "percentage_copied": pct_copied,
+                "correspondence_map_b64": f"data:image/jpeg;base64,{correspondence_map_b64}"
             },
-            "visual_evidence": {
-                "patch_mapping_image": f"data:image/jpeg;base64,{patch_vis_b64}",
-                "edge_image_a": f"data:image/jpeg;base64,{edge_a_b64}",
-                "edge_image_b": f"data:image/jpeg;base64,{edge_b_b64}"
+            "pillar_3_substantial_similarity": {
+                "perceptual_distance_lpips": lpips_distance,
+                "quantitative_heatmap_b64": f"data:image/jpeg;base64,{heatmap_b64}"
             }
         }
         
     except Exception as e:
+        import traceback
+        traceback.print_exc()
         raise HTTPException(status_code=500, detail=str(e))
 
 @app.get("/")
 def read_root():
-    return {"message": "Diagnostic API is running. Send POST requests to /analyze"}
+    return {"message": "3-Pillar Legal Diagnostic API is running."}