Create app.py

app.py

@@ -0,0 +1,106 @@
+import gradio as gr
+import torch
+import cv2
+import numpy as np
+from PIL import Image
+from transformers import CLIPProcessor, CLIPModel
+import pytesseract
+from ultralytics import YOLO
+from sklearn.metrics.pairwise import cosine_similarity
+
+# Load CLIP model and processor
+model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
+processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+
+# Load YOLOv8 model
+yolo = YOLO("yolov8x.pt")  # change to yolov8n.pt for faster inference if needed
+
+# Set Tesseract OCR path if needed
+pytesseract.pytesseract.tesseract_cmd = "/usr/bin/tesseract"
+
+def extract_clip_features(image):
+    inputs = processor(images=image, return_tensors="pt")
+    with torch.no_grad():
+        features = model.get_image_features(**inputs)
+    return features / features.norm(p=2, dim=-1, keepdim=True)
+
+def detect_text(image):
+    return pytesseract.image_to_string(image)
+
+def detect_objects(image):
+    results = yolo(image)
+    boxes = results[0].boxes.xywh.cpu().numpy()  # (x, y, w, h)
+    return boxes
+
+def compare_images(img1, img2):
+    img1 = img1.convert("RGB")
+    img2 = img2.convert("RGB")
+
+    # Extract features
+    feat1 = extract_clip_features(img1)
+    feat2 = extract_clip_features(img2)
+
+    # Cosine similarity
+    sim_score = cosine_similarity(feat1, feat2)[0][0]
+
+    # OCR text similarity
+    text1 = detect_text(img1)
+    text2 = detect_text(img2)
+
+    vec1 = np.array([ord(c) for c in text1[:100]] + [0]*100)[:100]
+    vec2 = np.array([ord(c) for c in text2[:100]] + [0]*100)[:100]
+    text_sim = cosine_similarity([vec1], [vec2])[0][0]
+
+    # Object detection
+    obj1 = detect_objects(np.array(img1))
+    obj2 = detect_objects(np.array(img2))
+    shape_diff = abs(len(obj1) - len(obj2)) / max(len(obj1), 1)
+
+    # Geometric center distance (avg heuristic)
+    dist_penalty = 0
+    for i in range(min(len(obj1), len(obj2))):
+        dist_penalty += np.linalg.norm(obj1[i][:2] - obj2[i][:2])
+    dist_penalty /= max(len(obj1), 1)
+
+    # Final similarity score
+    final_score = (0.5 * sim_score) + (0.3 * text_sim) + (0.2 * (1 - shape_diff))
+    rating = round(final_score * 5, 2)
+    rating_clamped = min(5.0, max(0.0, rating))
+
+    return {
+        "Similarity Score (%)": f"{final_score*100:.2f}%",
+        "Rating (0–5)": f"{rating_clamped:.1f} ⭐",
+        "Text in Image 1": text1.strip()[:200],
+        "Text in Image 2": text2.strip()[:200],
+        "Detected Objects (img1, img2)": f"{len(obj1)} vs {len(obj2)}"
+    }
+
+def gradio_ui(img1, img2):
+    result = compare_images(img1, img2)
+    return (
+        result["Similarity Score (%)"],
+        result["Rating (0–5)"],
+        result["Text in Image 1"],
+        result["Text in Image 2"],
+        result["Detected Objects (img1, img2)"]
+    )
+
+# Gradio Interface
+demo = gr.Interface(
+    fn=gradio_ui,
+    inputs=[
+        gr.Image(type="pil", label="Input Image"),
+        gr.Image(type="pil", label="Output Image")
+    ],
+    outputs=[
+        gr.Text(label="Similarity Score (%)"),
+        gr.Text(label="Predicted Rating (0-5 Stars)"),
+        gr.Textbox(label="Extracted Text from Input"),
+        gr.Textbox(label="Extracted Text from Output"),
+        gr.Text(label="Part/Object Count Comparison")
+    ],
+    title="🛠️ CAD Image Comparison AI",
+    description="Upload two CAD/Technical images and get a full feature-based similarity score including shapes, geometry, dimensions, and text."
+)
+
+demo.launch()