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())