Update app.py

app.py

@@ -1,58 +1,57 @@
 import time
 import streamlit as st
 import numpy as np
-import PIL
+from PIL import Image
 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()
 
-    labels = gen_labels()
-    predicted_class = np.argmax(prediction, axis=1)[0]
-    classification_result = labels[predicted_class]
+html_temp = '''
+    <div style="padding-bottom: 20px; padding-top: 20px; padding-left: 5px; padding-right: 5px">
+    <center><h1>Garbage Segregation</h1></center>
+    </div>
+    '''
+st.markdown(html_temp, unsafe_allow_html=True)
 
-    # Get the confidence (probability) of the predicted class
-    confidence = prediction[0][predicted_class] * 100  # Convert probability to percentage
+html_temp = '''
+    <div>
+    <h2></h2>
+    <center><h3>Please upload Waste Image to find its Category</h3></center>
+    </div>
+    '''
+st.markdown(html_temp, unsafe_allow_html=True)
 
-    return classification_result, confidence
+opt = st.selectbox("How do you want to upload the image for classification?\n", ('Please Select', 'Upload image via link', 'Upload image from device'))
 
-# 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'))
+image = None  # Initialize image variable
 
 if opt == 'Upload image from device':
     file = st.file_uploader('Select', type=['jpg', 'png', 'jpeg'])
     if file is not None:
-        pil_image = Image.open(file)
-        image = preprocess(pil_image)
+        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:
-        pil_image = Image.open(urllib.request.urlopen(img_url))
-        image = preprocess(pil_image)
+        img = st.text_input('Enter the Image Address')
+        image = Image.open(urllib.request.urlopen(img)).resize((256, 256), Image.LANCZOS)
     except ValueError:
-        st.error("Please Enter a valid Image Address!")
+        if st.button('Submit'):
+            show = st.error("Please Enter a valid Image Address!")
+            time.sleep(4)
+            show.empty()
+
+try:
+  if image is not None:
+    st.image(image, width = 300, caption = 'Uploaded Image')
+    if st.button('Predict'):
+        img = preprocess(image)
 
-if 'image' in locals():  # Check if image variable exists
-    st.image(image, width=300, caption='Uploaded Image')
+        model = model_arc()
+        #model.load_weights("classify_model.h5")
 
-    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}")
+        prediction = model.predict(img[np.newaxis, ...])
+        st.info('Hey! The uploaded image has been classified as " {} waste " '.format(labels[np.argmax(prediction[0], axis=-1)]))
+except Exception as e:
+  st.info(e)
+  pass