app.py

import gradio as gr
import numpy as np
import clip
import torch
from PIL import Image
import base64
from io import BytesIO
from decimal import Decimal

# Load the CLIP model
model, preprocess = clip.load("ViT-B/32")
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device).eval()

# Define a function to find similarity
def find_similarity(base64_image, text_input):
    # 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)

    # Tokenize the text input
    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)

    # Calculate cosine similarity
    similarity = (image_features @ text_features.T).squeeze(0).cpu().numpy()

    # Convert each element in the similarity array to Decimal
    similarity_decimal = [Decimal(float(score)) for score in similarity]

    # Format Decimal values as floats with specific precision (e.g., 4 decimal places)
    formatted_similarity = [f'{float(score):.5f}' for score in similarity_decimal]

    return formatted_similarity

# Create a Gradio interface
iface = gr.Interface(
    fn=find_similarity,
    inputs=[
        gr.inputs.Textbox(label="Base64 Image", lines=8),
        "text"
    ],
    outputs="text",
    live=True,
    interpretation="default",
    title="CLIP Model Image-Text Cosine Similarity",
    description="Upload a base64 image and enter text to find their cosine similarity.",
)

iface.launch()