import gradio as gr
from transformers import pipeline
import numpy as np
from PIL import Image

test1_labels = "me lill;;general store;;goods;;nike;;bags;;shoe;;apparel store;;48 hours sale;;shoes;;8.8;;roadmate;;multi brand outlet;;shoe store;;find your way;;bag shop;;mall;;ralph lauren;;adidas originals;;single luxury;;flat 50%;;kailasa;;yy sports;;sale upto 50%;;laurel;;converse;;pony;;store;;kappa;;glamshops;;grocery store;;fruits;;orefici;;sale;;vans;;convienence store;;met;;yy;;grocery;;milano;;Fresh eggs;;multi brand luxury;;reebok;;bombay electric;;cloth;;manequin;;delhi electric;;lee;;caps;;zara;;crocs;;pepe jenes;;adidas;;sandals;;napa;;project x;;asics;;niche fashion;;puma;;yonex;;t-shirt;;macys;;adidas shoes;;cloth store;;le mill;;sweet monster;;shirt store;;shirt"

test2_labels = "Bata;;Tiger;;Mevious;;ministreet;;Bagong;;pepsi;;Persian Cat;;Nike;;birds;;Chesterfield;;Parcel;;Adidas;;American Express;;lions;;Ciggerete;;Puma;;xyz;;Truck;;Good food;;square one;;Courier;;Cat;;circle one;;Cignal;;Classic;;No Face masks no entry;;open;;Ciggerate;;Pall Mall;;Fortune;;Winston;;Fed Express;;Chefs Bakery;;Chefs Kitchen;;RC;;Leapard;;Dog;;JTI;;Bakery;;sophie store;;Cheetah;;mini street 27;;Hygenic;;coca cola;;Marvel;;abc;;dogs;;Malboro;;chef jhez kitchen marife patiag;;Sphynix Cat;;cats;;British American Tobacco;;BAT;;Fed Ex;;mini street convenience store albert josephparren;;Absolute;;Dunhill"

pipe = pipeline("zero-shot-image-classification", model="openai/clip-vit-base-patch32")

def shot(image, labels_text):
    PIL_image = Image.fromarray(np.uint8(image)).convert('RGB')
    labels = labels_text.split(";;")
    res = pipe(images=PIL_image, 
           candidate_labels=labels,
           hypothesis_template="This is a photo of {}")
    return {dic["label"]: dic["score"] for dic in res}
    
iface = gr.Interface(shot, 
                    ["image", "text"], 
                    "label", 
                    examples=[
							  ["examples/test1/1.jpg", "ralph lauren;;apparel store;;ralph lauren store;;shirts;;wardrobe;;white flower"], 
                              ["examples/test1/2.JPG", "adidas;;apparel store;;adidas store;;shirts;;wardrobe;;women training;;shoes"], 
                              ["examples/test1/3.jpg", "project x;;sweet monster;;bags store;;glass windows;;hanging lights"], 
                              ["examples/test1/4.JPG", "multi brand store;;multi brand shoe store;;shoe store;;mannequins;;adidas store;;reebok store;;puma store"],
							  
                              #["examples/test2/1.jpg", test2_labels],
                              #["examples/test2/2.jpg", test2_labels],
							 ],
                    description="Add a picture and a list of labels separated by ;;",
                    title="Zero-shot Image Classification")

iface.launch()