DevBM's picture
Update app.py
65a4186 verified
raw
history blame
1.59 kB
import streamlit as st
from transformers import pipeline
import requests
from PIL import Image
from io import BytesIO
import pandas as pd
st.subheader("Image Classification", divider='orange')
if st.toggle(label='Show Pipe4'):
models = [
'google/vit-base-patch16-224',
'WinKawaks/vit-tiny-patch16-224',
'microsoft/resnet-50',
'facebook/deit-base-distilled-patch16-224',
'facebook/convnext-large-224',
'apple/mobilevit-small',
'timm/resnet50.a1_in1k'
]
model_name = st.selectbox(
label='Select Model',
options=models,
placeholder='google/vit-base-patch16-224',
)
pipe = pipeline("image-classification", model=model_name)
url = 'https://media.istockphoto.com/id/182756302/photo/hot-dog-with-grilled-peppers.jpg?s=1024x1024&w=is&k=20&c=NCHo2P94a-PfRDKzWSe4h6oACQZ-_ubZqUBj5CMSEWY='
response = requests.get(url=url)
image_bytes = BytesIO(response.content)
image = Image.open(image_bytes)
# image = Image.open(BytesIO(requests.get(url).content))
# use_default = st.checkbox(label='Use default image')
file = st.file_uploader(label='Upload image')
if file is not None:
image = Image.open(file)
res = pipe(image)
if st.toggle(label='Show row data'):
st.write(res)
p = pd.DataFrame(res)
p = p.sort_values(by='score',ascending=False)
col1, col2 = st.columns(2)
col1.write(image)
col2.write(p['label'])
st.bar_chart(p.set_index('label'))
st.area_chart(p.set_index('label'))
# col2.bar_chart(p.set_index('label'))