Create main.py

main.py

@@ -0,0 +1,181 @@
+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, ImageDraw
+import io
+import base64
+from transformers import AutoImageProcessor, AutoModel, CLIPProcessor, CLIPModel
+
+# ==============================================================================
+# 1. Initialize FastAPI & CORS
+# ==============================================================================
+app = FastAPI(title="Copyright Diagnostic API")
+
+# Crucial for allowing your React frontend to communicate with this backend
+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"
+
+print("Loading DINOv2...")
+dino_processor = AutoImageProcessor.from_pretrained("facebook/dinov2-base")
+dino_model = AutoModel.from_pretrained("facebook/dinov2-base").to(device)
+dino_model.eval()
+
+print("Loading CLIP...")
+clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(device)
+clip_model.eval()
+
+# Helper function to convert PIL Image to Base64 string for React
+def image_to_base64(img: Image.Image) -> str:
+    buffered = io.BytesIO()
+    img.save(buffered, format="JPEG")
+    return base64.b64encode(buffered.getvalue()).decode("utf-8")
+
+# ==============================================================================
+# 3. AI Pipeline Functions
+# ==============================================================================
+
+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)
+    
+    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):
+    arr_a = np.array(img_a.convert('L'))
+    arr_b = np.array(img_b.convert('L'))
+    
+    edges_a = cv2.Canny(arr_a, 100, 200)
+    edges_b = cv2.Canny(arr_b, 100, 200)
+    
+    edges_b_resized = cv2.resize(edges_b, (edges_a.shape[1], edges_a.shape[0]))
+    
+    intersection = np.logical_and(edges_a > 0, edges_b_resized > 0).sum()
+    union = np.logical_or(edges_a > 0, edges_b_resized > 0).sum()
+    
+    iou_score = intersection / union if union != 0 else 0.0
+    return round(iou_score, 4), Image.fromarray(edges_a), Image.fromarray(edges_b)
+
+def compute_patch_similarity(img_a: Image.Image, img_b: Image.Image):
+    target_size = (224, 224)
+    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").to(device)
+    inputs_b = dino_processor(images=img_b_resized, 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.85 
+    
+    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_score = len(matches) / 256.0
+    
+    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)
+    
+    grid_size = 16
+    patch_size = 14
+    
+    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]
+        
+        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)
+        
+        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
+
+# ==============================================================================
+# 4. API Endpoints
+# ==============================================================================
+
+@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)
+        
+        # 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)    
+
+        # Return a clean JSON package to React
+        return {
+            "status": "success",
+            "scores": {
+                "semantic_idea": semantic_score,
+                "structural_layout": struct_score,
+                "fragmented_literal": patch_score
+            },
+            "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}"
+            }
+        }
+        
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/")
+def read_root():
+    return {"message": "Diagnostic API is running. Send POST requests to /analyze"}