app.py

import gradio as gr
import torch
import torch.nn.functional as F
import cv2
import numpy as np
from PIL import Image, ImageDraw
from transformers import AutoImageProcessor, AutoModel, CLIPProcessor, CLIPModel

# ==============================================================================
# 1. Global Initialization & Memory Management
# ==============================================================================
# We load models globally so they cache in memory on startup, not on every click.
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() # Prevent gradient tracking

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() # Prevent gradient tracking

# ==============================================================================
# 2. Pipeline Functions
# ==============================================================================

def compute_semantic_similarity(img_a: Image.Image, img_b: Image.Image) -> float:
    """
    Model: CLIP
    Legal Concept: The "Idea" / Semantic Referent
    """
    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:
            # Fallback for tuple-like object behavior
            image_features = image_features[1] if isinstance(image_features, tuple) and len(image_features) > 1 else image_features[0]
    # ---------------------------------------------------
            
    # Normalize and compute cosine similarity
    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):
    """
    Model: OpenCV Canny
    Legal Concept: "Substantial Similarity" (Layout / Composition)
    """
    # Convert to grayscale numpy arrays
    arr_a = np.array(img_a.convert('L'))
    arr_b = np.array(img_b.convert('L'))
    
    # Extract structural edges
    edges_a = cv2.Canny(arr_a, 100, 200)
    edges_b = cv2.Canny(arr_b, 100, 200)
    
    # Calculate structural overlap using Intersection over Union (IoU) of edges
    # We resize edges_b to match edges_a to ensure matrix math works
    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):
    """
    Model: DINOv2
    Legal Concept: "Fragmented Literal Similarity" (Scattered Literal Copying)
    """
    # Resize to a fixed multiple of patch size (14) so we have a known grid
    # 224x224 gives us a 16x16 grid of patches (256 total patches)
    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)
        
    # Isolate patches (skip CLS token) and normalize. Shape: (256, 768)
    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)
    
    # Compute N x M similarity matrix using dot product
    sim_matrix = torch.matmul(emb_a, emb_b.T) # Shape: (256, 256)
    
    # Mutual Nearest Neighbors logic to filter out noise
    best_b_for_a = torch.argmax(sim_matrix, dim=1)
    best_a_for_b = torch.argmax(sim_matrix, dim=0)
    
    matches = []
    # Threshold for what we consider "copied" (adjust based on testing)
    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: # It's a mutual match
            score = sim_matrix[a_idx, b_idx].item()
            if score >= SIMILARITY_THRESHOLD:
                matches.append((a_idx, b_idx, score))
                
    # Calculate overall patch score based on percentage of matching patches
    patch_score = len(matches) / 256.0
    
    # --- Visual Evidence Generation ---
    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:
        # Image A coordinates
        ay = (a_idx // grid_size) * patch_size
        ax = (a_idx % grid_size) * patch_size
        
        # Image B coordinates (shifted X by the width of Image A)
        by = (b_idx // grid_size) * patch_size
        bx = (b_idx % grid_size) * patch_size + target_size[0]
        
        # Draw bounding boxes
        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)
        
        # Draw connecting line
        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

# ==============================================================================
# 3. Main Orchestration Function
# ==============================================================================

def analyze_images(image_a, image_b):
    if image_a is None or image_b is None:
        raise gr.Error("Please upload both images.")

    # 1. Semantic Match
    semantic_score = compute_semantic_similarity(image_a, image_b)
    
    # 2. Structural Match
    struct_score, edge_a, edge_b = compute_structural_similarity(image_a, image_b)
    
    # 3. Patch Match
    patch_score, patch_vis = compute_patch_similarity(image_a, image_b)
    
    return (
        semantic_score, 
        struct_score, 
        patch_score, 
        patch_vis,
        edge_a, 
        edge_b
    )

# ==============================================================================
# 4. Gradio UI / UX
# ==============================================================================

with gr.Blocks() as demo:
    gr.Markdown("# Assistive Diagnostic Framework for Copyright Infringement")
    gr.Markdown("Upload two images to compare them across semantic, structural, and literal fragment dimensions.")
    
    with gr.Row():
        with gr.Column():
            img_in_a = gr.Image(type="pil", label="Image A (Original)")
        with gr.Column():
            img_in_b = gr.Image(type="pil", label="Image B (Suspected Copy)")
            
    btn_analyze = gr.Button("Analyze Similarity", variant="primary")
    
    gr.Markdown("### Assessment Metrics & Legal Context")
    
    # Context pulled directly from the paper
    gr.Markdown("""
    * **Semantic Match (Idea-Expression Dichotomy):** Evaluates if works share an underlying conceptual basis. Models like CLIP link visual data to semantic concepts to bridge the "semantic gap". This acts as an initial threshold, as copyright protects concrete expression rather than mere abstract ideas.
    * **Structural Match (Substantial Similarity):** Detects structural relationships through spatial mapping. This addresses how the specific combination and arrangement of individual elements can form part of the protected expression.
    * **Patch Match (Fragmented Literal Similarity):** Identifies "scattered literal copying" where specific fragments of a protected work are reproduced directly or in a highly proximate manner, without the new work as a whole having to resemble the original. Localized patch-based matching algorithms divide the image to detect these copied fragments regardless of where they are placed in the new works.
    """)

    with gr.Row():
        score_semantic = gr.Number(label="Semantic Match (CLIP) - Idea", show_label=True)
        score_struct = gr.Number(label="Structural Match (Edge IoU) - Layout", show_label=True)
        score_patch = gr.Number(label="Patch Match (DINOv2) - Fragmented Literal", show_label=True)
        
    gr.Markdown("### Visual Evidence")
    
    # Replaced Tabs with a Row containing two Columns for side-by-side display
    with gr.Row():
        with gr.Column():
            gr.Markdown("#### Fragmented Literal Similarity (DINOv2)")
            gr.Markdown("**Red boxes and green lines indicate mutually correlating local patches (Similarity > 0.85)**")
            vis_patch = gr.Image(label="Patch Mapping Visualization", type="pil")
            
        with gr.Column():
            gr.Markdown("#### Substantial Similarity (Edge Detection)")
            gr.Markdown("**Comparison of spatial arrangements and structural boundaries.**")
            with gr.Row():
                vis_edge_a = gr.Image(label="Image A Edges", type="pil")
                vis_edge_b = gr.Image(label="Image B Edges", type="pil")

    btn_analyze.click(
        fn=analyze_images,
        inputs=[img_in_a, img_in_b],
        outputs=[
            score_semantic, 
            score_struct, 
            score_patch, 
            vis_patch,
            vis_edge_a, 
            vis_edge_b 
        ]
    )

if __name__ == "__main__":
    # Remember to keep your theme in the launch method if you are using Gradio 6.0!
    demo.launch(theme=gr.themes.Soft())