darkmortal's picture
Update app.py
cc2125a verified
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")