import clip
import torch
import gradio as gr
import torchvision.transforms as T
from PIL import Image
try:
    from torchvision.transforms import InterpolationMode
    BICUBIC = InterpolationMode.BICUBIC
except ImportError:
    BICUBIC = Image.BICUBIC
import warnings
warnings.filterwarnings("ignore")

#MODEL_PATH = '/media/delta/S/clipmodel.pth'                        #CHANGE THIS IF YOU WANT TO CHANGE THE MODEL PATH
MODEL_PATH ='/media/delta/S/clipmodel_large.pth'                    #CHANGE THIS IF YOU WANT TO CHANGE THE MODEL PATH


device = "cuda" if torch.cuda.is_available() else "cpu"
model = clip.model.build_model(torch.load(MODEL_PATH)).to(device)
preprocess = clip.clip._transform(model.visual.input_resolution)

def zeroshot_detection(Press_Clear_Dont_Stack_Image):
    inp = Press_Clear_Dont_Stack_Image

    captions = "photo of a guardrail, no guardrail in the photo"   #CHANGE THIS IF YOU WANT TO CHANGE THE PREDICTION: separate by commas

    captions = captions.split(',')
    caption = clip.tokenize(captions).to(device)
    image = preprocess(inp).unsqueeze(0).to(device)
    with torch.no_grad():
        image_features = model.encode_image(image)
        text_features = model.encode_text(caption)
    image_features /= image_features.norm(dim=-1, keepdim=True)
    text_features /= text_features.norm(dim=-1, keepdim=True)
    similarity = (100.0 * image_features @ text_features.T).softmax(dim=-1)
    values, indices = similarity[0].topk(len(captions))
    return {captions[indices[i].item()]: float(values[i].item()) for i in range(len(values))}   

gr.Interface(fn=zeroshot_detection, 
             inputs=[gr.Image(type="pil")],
             outputs=gr.Label(num_top_classes=1)).launch()