import streamlit as st
from PIL import Image

from transformers import pipeline
import numpy as np
from transformers import AutoFeatureExtractor, AutoModelForImageClassification


st.set_page_config(layout='wide',
                   page_title='Garbage image classification'
                   )




def main():
    


    st.header("Try it out!")

    uploaded_file = st.file_uploader("Upload Files",type=['png','jpeg'])
    img=Image.open(uploaded_file)

    extractor = AutoFeatureExtractor.from_pretrained("yangy50/garbage-classification")
    model = AutoModelForImageClassification.from_pretrained("yangy50/garbage-classification")

    inputs = extractor(img,return_tensors="pt")
    outputs = model(**inputs)
    label_num=outputs.logits.softmax(1).argmax(1)
    label_num=label_num.item()

    if label_num==0:
        st.write("cardboard")
    elif label_num==1:
        st.write("glass")
    elif label_num==2:
        st.write("metal")
    elif label_num==3:
        st.write("paper")
    elif label_num==4:
        st.write("plastic")
    else:
        st.write("trash")



if __name__ == '__main__':
    main()