Update app.py

app.py

@@ -3,28 +3,16 @@ import streamlit as st
 import numpy as np
 from PIL import Image
 import urllib.request
-from utils import *
+from utils import *  # Assuming the gen_labels() and preprocess() functions are in this module
 
+# Load labels
 labels = gen_labels()
 
-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)
+# 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)
 
-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)
-
-opt = st.selectbox("How do you want to upload the image for classification?\n", ('Please Select', 'Upload image via link', 'Upload image from device'))
-
-image = None  # Initialize image variable
+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'])
@@ -32,26 +20,26 @@ if opt == 'Upload image from device':
         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:
-        img = st.text_input('Enter the Image Address')
-        image = Image.open(urllib.request.urlopen(img)).resize((256, 256), Image.LANCZOS)
+        image = Image.open(urllib.request.urlopen(img_url)).resize((256, 256), Image.LANCZOS)
     except ValueError:
-        if st.button('Submit'):
-            show = st.error("Please Enter a valid Image Address!")
-            time.sleep(4)
-            show.empty()
+        st.error("Please Enter a valid Image Address!")
 
-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")
-
-        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
+    if st.button('Predict'):
+        try:
+            img_array = preprocess(image)
+            model = model_arc()  # Assuming this function initializes and returns your model
+            prediction = model.predict(img_array[np.newaxis, ...])
+            
+            # Get the predicted class name
+            predicted_class_index = np.argmax(prediction[0], axis=-1)
+            predicted_class_name = labels[predicted_class_index]
+            
+            st.info('The uploaded image has been classified as "{}" waste.'.format(predicted_class_name))
+            
+        except Exception as e:
+            st.error(f"An error occurred: {e}")