from fastapi import FastAPI, File, UploadFile, HTTPException from fastapi.middleware.cors import CORSMiddleware import torch import torch.nn.functional as F import cv2 import numpy as np 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 - 3 Pillar XAI") app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # ============================================================================== # 2. Global Initialization & Memory Management # ============================================================================== device = "cuda" if torch.cuda.is_available() else "cpu" 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(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() 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 # ============================================================================== # 4. AI Pipeline Functions (Mapped to the 3 Pillars) # ============================================================================== 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(): 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(): outputs = clip_model.get_image_features(**inputs) if hasattr(outputs, 'image_embeds'): features = outputs.image_embeds elif isinstance(outputs, torch.Tensor): features = outputs else: features = outputs.pooler_output 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))) # 1. Calculate ALL valid matches that pass the threshold all_valid_matches = nms_matches(raw_matches, grid_w, PATCH_SIZE, radius=2.0) total_match_count = len(all_valid_matches) # 2. Slice the list to only draw the top N matches to prevent visual clutter vis_matches = all_valid_matches[: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') 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(vis_matches), 1)) # 3. Only iterate over the sliced vis_matches for drawing for i, (idx_a, idx_b, score) in enumerate(vis_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) tl_xa = xa - PATCH_SIZE / 2 tl_ya = ya - PATCH_SIZE / 2 tl_xb = xb_canvas - PATCH_SIZE / 2 tl_yb = yb - PATCH_SIZE / 2 rect_a = plt.Rectangle((tl_xa, tl_ya), PATCH_SIZE, PATCH_SIZE, linewidth=1.5, edgecolor='#ef4444', facecolor='none', alpha=0.9, zorder=4) ax.add_patch(rect_a) rect_b = plt.Rectangle((tl_xb, tl_yb), PATCH_SIZE, PATCH_SIZE, linewidth=1.5, edgecolor='#ef4444', facecolor='none', alpha=0.9, zorder=4) ax.add_patch(rect_b) 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) mx = (xa + xb_canvas) / 2 my = (ya + yb) / 2 bbox_props = dict(boxstyle="round,pad=0.25", fc="#1e293b", ec=color, alpha=0.95, lw=1) 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) # 4. Return the TOTAL count, not the sliced count return fig, total_match_count 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 # ============================================================================== # 5. Primary Analysis Endpoint # ============================================================================== @app.post("/analyze") async def analyze_artworks(file_a: UploadFile = File(...), file_b: UploadFile = File(...)): try: img_a = Image.open(io.BytesIO(await file_a.read())).convert("RGB") img_b = Image.open(io.BytesIO(await file_b.read())).convert("RGB") # --- 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") 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) 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) 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] # Call the figure generator exactly ONCE. It now returns the total count. corr_fig, total_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) n_patches_a = a2b_best.shape[0] # 1. Calculate the raw statistical threshold raw_adaptive_thresh = (a2b_best.mean() + a2b_best.std()).item() # 2. Clamp it: never lower than 0.5 (noise), never higher than 0.85 (identical images) smart_thresh = min(max(raw_adaptive_thresh, 0.5), 0.85) # 3. Count patches above this dynamic, safe threshold high_a = (a2b_best > smart_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) 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 { "status": "success", "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": total_match_count, "percentage_copied": pct_copied, "correspondence_map_b64": f"data:image/jpeg;base64,{correspondence_map_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": "3-Pillar Legal Diagnostic API is running."}