Spaces:
Running
Running
| import time | |
| import streamlit as st | |
| import numpy as np | |
| import PIL | |
| import urllib.request | |
| from utils import * | |
| #from fastai.data.external import * | |
| # Function to classify the garbage | |
| def classify_garbage(img_path, model): | |
| processed_img = preprocess(img_path) | |
| prediction = model.predict(processed_img) | |
| labels = gen_labels() | |
| predicted_class = np.argmax(prediction, axis=1)[0] | |
| classification_result = labels[predicted_class] | |
| # Get the confidence (probability) of the predicted class | |
| confidence = prediction[0][predicted_class] * 100 # Convert probability to percentage | |
| return classification_result, confidence | |
| # Streamlit app layout | |
| st.markdown('<center><h1>Garbage Segregation</h1></center>', unsafe_allow_html=True) | |
| st.markdown('<center><h3>Please upload Waste Image to find its Category</h3></center>', unsafe_allow_html=True) | |
| opt = st.selectbox("How do you want to upload the image for classification?", ('Please Select', 'Upload image via link', 'Upload image from device')) | |
| if opt == 'Upload image from device': | |
| file = st.file_uploader('Select', type=['jpg', 'png', 'jpeg']) | |
| if file is not None: | |
| image = Image.open(file).resize((256, 256))#, Image.LANCZOS) | |
| elif opt == 'Upload image via link': | |
| img_url = st.text_input('Enter the Image Address') | |
| try: | |
| image = Image.open(urllib.request.urlopen(img_url)).resize((256, 256))#, Image.LANCZOS) | |
| except ValueError: | |
| st.error("Please Enter a valid Image Address!") | |
| if 'image' in locals(): # Check if image variable exists | |
| st.image(image, width=300, caption='Uploaded Image') | |
| if st.button('Predict'): | |
| try: | |
| model = model_arc() # Initialize your model | |
| # Ensure image shape is correct and add batch dimension | |
| img_array = preprocess(image) # This should return an array of shape (1, 256, 256, 3) | |
| predicted_class, confidence = classify_garbage(img_array, model) | |
| st.info('The uploaded image has been classified as "{}" with {:.2f}% confidence.'.format(predicted_class, confidence)) | |
| except Exception as e: | |
| st.error(f"An error occurred: {e}") |