import gradio as gr
import cv2
import numpy as np
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
import base64
from PIL import Image
from io import BytesIO
import torch
import clip

# Load the segmentation model
sam_checkpoint = "sam_vit_h_4b8939.pth"
model_type = "vit_h"
sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)

# Load the CLIP model
model, preprocess = clip.load("ViT-L/14@336px")
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device).eval()


def find_similarity(base64_image, text_input):
    try:
        # Decode the base64 image to bytes
        image_bytes = base64.b64decode(base64_image)

        # Convert the bytes to a PIL image
        image = Image.open(BytesIO(image_bytes))

        # Preprocess the image
        image = preprocess(image).unsqueeze(0).to(device)

        # Prepare input text
        text_tokens = clip.tokenize([text_input]).to(device)

        # Encode image and text features
        with torch.no_grad():
            image_features = model.encode_image(image)
            text_features = model.encode_text(text_tokens)

        # Normalize features and calculate similarity
        image_features /= image_features.norm(dim=-1, keepdim=True)
        text_features /= text_features.norm(dim=-1, keepdim=True)
        similarity = (text_features @ image_features.T).squeeze(0).cpu().numpy()

        return similarity
    except Exception as e:
        return str(e)


def segment_image(input_image, text_input):
    image_bytes = base64.b64decode(input_image)
    image = Image.open(BytesIO(image_bytes))

    # Convert the image to RGB color mode
    image = image.convert("RGB")

    # Convert the image to a numpy array
    image = np.array(image)

    mask_generator = SamAutomaticMaskGenerator(sam)
    masks = mask_generator.generate(image)

    segmented_regions = []  # List to store segmented regions with similarity scores

    for i, mask_dict in enumerate(masks):
        mask_data = (mask_dict['segmentation'] * 255).astype(np.uint8)

        # Create a mask with the same shape as the original image
        mask = np.zeros_like(image)
        mask[:, :] = mask_data[:, :, np.newaxis]

        # Apply the mask to the original image
        segmented_region = cv2.bitwise_and(image, mask)

        x, y, w, h = map(int, mask_dict['bbox'])
        cropped_region = segmented_region[y:y+h, x:x+w]

        if not cropped_region.size:
            # If the cropped region is empty, return the input image as is
            return input_image

        # Convert to base64 image
        _, buffer = cv2.imencode(".png", cv2.cvtColor(cropped_region, cv2.COLOR_BGR2RGB))
        segmented_image_base64 = base64.b64encode(buffer).decode()

        # Calculate similarity for the segmented image
        similarity = find_similarity(segmented_image_base64, text_input)

        # Append the segmented image and its similarity score
        segmented_regions.append({"image": segmented_image_base64, "similarity": similarity})

    # Sort the segmented images by similarity in descending order
    segmented_regions.sort(key=lambda x: x["similarity"], reverse=True)

    # Limit the output to the top 6 key-value pairs
    segmented_regions = segmented_regions[:6]

    # Return the segmented images in descending order of similarity
    return segmented_regions


# Create Gradio components
input_image = gr.Textbox(label="Base64 Image", lines=8)
text_input = gr.Textbox(label="Text Input")  # Use Textbox with a label

# Create a Gradio interface
gr.Interface(fn=segment_image, inputs=[input_image, text_input], outputs="text").launch()