import gradio as gr
import open_clip
import torch
import requests
import numpy as np
from PIL import Image

shapes = ["leggings", "jogger", 
          "palazzo", "cargo",
          "dresspants", "chinos"]

model, preprocess_train, preprocess_val = open_clip.create_model_and_transforms('hf-hub:Marqo/marqo-fashionCLIP')
tokenizer = open_clip.get_tokenizer('hf-hub:Marqo/marqo-fashionCLIP')


shapes_desc = list(map(lambda x: "a " + x + " pants shape", shapes))
text = tokenizer(shapes_desc)

with torch.no_grad(), torch.cuda.amp.autocast():
  text_features = model.encode_text(text)
  text_features /= text_features.norm(dim=-1, keepdim=True)

def predict(inp):
  image = preprocess_val(inp).unsqueeze(0)

  with torch.no_grad(), torch.cuda.amp.autocast():
      image_features = model.encode_image(image)
      image_features /= image_features.norm(dim=-1, keepdim=True)

      text_probs = (100 * image_features @ text_features.T).softmax(dim=-1)
      confidences = {shapes[i]: float(text_probs[0, i]) for i in range(6)}
  return confidences


gr.Interface(fn=predict,
             inputs=gr.Image(type="pil"),
             outputs=gr.Label(num_top_classes=6),
             examples=["imgs/cargo.jpg", "imgs/palazzo.jpg",
                      "imgs/leggings.jpg", "imgs/jogger.jpg",
                      "imgs/chinos.jpg", "imgs/dresspants.jpg"]).launch(share=True)