Update app.py

app.py

@@ -2,65 +2,78 @@ import gradio as gr
 import cv2
 import numpy as np
 from skimage.metrics import structural_similarity as ssim
-from PIL import Image
-import pytesseract
-import matplotlib.pyplot as plt
+from transformers import BertForSequenceClassification, BertTokenizer
+import torch
+
+
+# Load pre-trained model and tokenizer
+model = BertForSequenceClassification.from_pretrained("bert-base-uncased")
+tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
+
 
-# Function to calculate SSIM between two images
 def calculate_ssim(img1, img2):
-    img1_gray = img1.convert("L")
-    img2_gray = img2.convert("L")
-    img1_resized = img1_gray.resize((256, 256))
-    img2_resized = img2_gray.resize((256, 256))
-    return ssim(np.array(img1_resized), np.array(img2_resized))
-
-# Function to extract text from an image
-def extract_text(image):
-    text = pytesseract.image_to_string(image)
-    return text
-
-# Function to calculate color difference between two images
-def calculate_color_difference(img1, img2):
-    hist1 = cv2.calcHist([np.array(img1)], [0, 1, 2], None, [8, 8, 8], [0, 256, 0, 256, 0, 256])
-    hist2 = cv2.calcHist([np.array(img2)], [0, 1, 2], None, [8, 8, 8], [0, 256, 0, 256, 0, 256])
-    return cv2.compareHist(hist1, hist2, cv2.HISTCMP_CHISQR)
-
-# Function to compare two trademarks
+    img1_gray = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
+    img2_gray = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)
+    return ssim(img1_gray, img2_gray)
+
+
+def calculate_text_similarity(text1, text2):
+    encoded_text1 = tokenizer(text1, truncation=True, padding=True, return_tensors="pt")
+    encoded_text2 = tokenizer(text2, truncation=True, padding=True, return_tensors="pt")
+
+    with torch.no_grad():
+        outputs_text1 = model(**encoded_text1)
+        outputs_text2 = model(**encoded_text2)
+
+    embeddings_text1 = outputs_text1.last_hidden_state.squeeze(0)
+    embeddings_text2 = outputs_text2.last_hidden_state.squeeze(0)
+
+    text_similarity = ssim(embeddings_text1.numpy(), embeddings_text2.numpy())
+
+    return text_similarity
+
+
+def calculate_color_similarity(img1, img2):
+    img1_lab = cv2.cvtColor(img1, cv2.COLOR_BGR2Lab)
+    img2_lab = cv2.cvtColor(img2, cv2.COLOR_BGR2Lab)
+
+    color_similarity = ssim(img1_lab, img2_lab, multichannel=True)
+
+    return color_similarity
+
+
 def compare_trademarks(trademark1, trademark2):
-    ssim_score = calculate_ssim(trademark1, trademark2)
-    text1 = extract_text(trademark1)
-    text2 = extract_text(trademark2)
+    img1 = cv2.imread(trademark1)
+    img2 = cv2.imread(trademark2)
+
+    ssim_score = calculate_ssim(img1, img2)
+
+    text1 = "Trademark text 1"
+    text2 = "Trademark text 2"
+
     text_similarity = calculate_text_similarity(text1, text2)
-    color_difference = calculate_color_difference(trademark1, trademark2)
-    return ssim_score, text_similarity, color_difference
 
-# Function to calculate text similarity ratio
-def calculate_text_similarity(text1, text2):
-    if text1 == "" or text2 == "":
-        return 0.0
-    text1 = text1.lower()
-    text2 = text2.lower()
-    intersection = len(set(text1.split()) & set(text2.split()))
-    union = len(set(text1.split()) | set(text2.split()))
-    return intersection / union
-
-# Function to handle the Gradio interface
+    color_similarity = calculate_color_similarity(img1, img2)
+
+    return ssim_score, text_similarity, color_similarity
+
+
 def prevent_trademark_conflict(trademark1, trademark2):
-    ssim_score, text_similarity, color_difference = compare_trademarks(trademark1, trademark2)
-    result = f"SSIM Score: {ssim_score:.2f}\nText Similarity: {text_similarity:.2f}\nColor Difference: {color_difference:.2f}"
-    return result
+    similarity_scores = compare_trademarks(trademark1, trademark2)
+    return similarity_scores
+
 
-# Prepare Gradio interface
-iface = gr.Interface(
+# Interface
+trademark_comparison_interface = gr.Interface(
     fn=prevent_trademark_conflict,
     inputs=[
         gr.inputs.Image(type="pil", label="Trademark Image 1"),
-        gr.inputs.Image(type="pil", label="Trademark Image 2")
+        gr.inputs.Image(type="pil", label="Trademark Image 2"),
     ],
     outputs="text",
-    title="Trademark Similarity",
-    description="Compare two trademarks to check for similarity."
+    title="Trademark Comparison",
+    description="Compare two trademarks based on SSIM, text similarity, and color similarity.",
 )
 
 # Launch the interface
-iface.launch()
+trademark_comparison_interface.launch()