Spaces:
Runtime error
Runtime error
| 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() | |