from transformers import pipeline
from PIL import Image
import requests
from io import BytesIO

classifier = pipeline("zero-shot-image-classification", model="google/siglip-base-patch16-224")

st.title("Image classifier model demo")

file_name = st.file_uploader("Upload an image")

def scan_image(image, label, tolerance = 0.01):
    predictions = classifier(image, candidate_labels = [label, "other"])
    dict = {}
    for prediction in predictions:
        dict[prediction['label']] = prediction['score']
    # print(json.dumps(dict, indent = 3))
    return (dict[label] > (dict['other'] + tolerance), dict)

if file_name is not None:
    col1, col2 = st.columns(2)

    image = Image.open(file_name)
    col1.image(image, use_column_width=True)

    label = st.text_input("What to look for in the image?", "Cats")

    predictions = scan_image(image, label)
    
    col2.header("Probabilities")
    for key, value in predictions[1]:
        col2.subheader(f"{ key }: { round(value * 100, 1)}%")

    if predictions[1]:
        st.header("The object is present in the given image")
    else: st.header("The object is not found in the given image")