from PIL import Image
from difflib import SequenceMatcher
from skimage.metrics import structural_similarity as ssim
import gradio as gr
import pytesseract
import io


# Function to calculate SSIM between two images
def calculate_ssim(img1, img2):
    img1_gray = Image.open(io.BytesIO(img1)).convert("L")
    img2_gray = Image.open(io.BytesIO(img2)).convert("L")
    return ssim(img1_gray, img2_gray)


# Function to extract text from an image using OCR
def extract_text(image):
    image = Image.open(image)
    text = pytesseract.image_to_string(image).lower()
    return text

# Function to compare trademarks based on text similarity
def compare_text(trademark1, trademark2):
    text1 = extract_text(trademark1)
    text2 = extract_text(trademark2)
    similarity_ratio = SequenceMatcher(None, text1, text2).ratio()
    return similarity_ratio

# Function to compare trademarks based on color similarity
def compare_colors(trademark1, trademark2):
    trademark1 = trademark1.convert("RGB")
    trademark2 = trademark2.convert("RGB")
    colors1 = trademark1.getcolors(trademark1.size[0] * trademark1.size[1])
    colors2 = trademark2.getcolors(trademark2.size[0] * trademark2.size[1])
    color_vector1 = sum([(count * np.array(color)) for count, color in colors1]) / (trademark1.size[0] * trademark1.size[1])
    color_vector2 = sum([(count * np.array(color)) for count, color in colors2]) / (trademark2.size[0] * trademark2.size[1])
    color_similarity = 1 - np.linalg.norm(color_vector1 - color_vector2)
    return color_similarity

# Function to compare trademarks based on multiple aspects
def compare_trademarks(trademark1, trademark2):
    ssim_score = calculate_ssim(trademark1, trademark2)
    text_similarity = compare_text(trademark1, trademark2)
    color_similarity = compare_colors(trademark1, trademark2)

    # Adjust weights based on the importance of each aspect
    ssim_weight = 0.6
    text_weight = 0.2
    color_weight = 0.2

    overall_similarity = (ssim_weight * ssim_score) + (text_weight * text_similarity) + (color_weight * color_similarity)
    return overall_similarity

# Function to perform trademark conflict prevention
def prevent_trademark_conflict(trademark1, trademark2):
    similarity_score = compare_trademarks(trademark1, trademark2)
    result = f"Trademark Similarity Score: {similarity_score:.4f}"
    return result

iface = gr.Interface(
    fn=prevent_trademark_conflict,
    inputs=[
        gr.inputs.File(label="Trademark Image 1"),
        gr.inputs.File(label="Trademark Image 2")
    ],
    outputs="text",
    title="Trademark Conflict Prevention",
    description="Upload two trademark images. Get the trademark similarity score."
)

iface.launch()