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',
    ]
    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'))