app.py

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','jpg'])
    
    if uploaded_file!=None:

        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()