import streamlit as st
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?")
    if label == '':
        st.warning('Please enter a object label', icon="⚠️")
    else:
        if st.button("Scan Image"):
            predictions = scan_image(image, label)
            
            col2.header("Probabilities") 
            for key in predictions[1].keys():
                col2.subheader(f"{ key }: { round(predictions[1][key] * 100, 1)}%")
    
            if predictions[0]:
                st.header("The object is present in the given image")
            else: st.header("The object is not found in the given image")