Update app.py

app.py

@@ -1,14 +1,15 @@
 import numpy as np
 import streamlit as st
 from PIL import Image
-import tensorflow as tf
+import torch
+import torch.nn as nn
+import torchvision.transforms as transforms
 from utils import preprocess_image
 
-
-
 # Initialize labels and model
 labels = ['cardboard', 'glass', 'metal', 'paper', 'plastic', 'trash']
-model = tf.keras.models.load_model('classify_model.h5')
+model = torch.load('classify_model.pth')
+model.eval()
 
 # Customized Streamlit layout
 st.set_page_config(
@@ -86,7 +87,7 @@ st.title("EcoIdentify by EcoClim Solutions")
 st.header("Upload a waste image to find its category")
 
 # Note
-st.markdown("* Please note that our dataset is trained primarily with images that contain a white background.  Therefore, images with white background would produce maximum accuracy *")
+st.markdown("* Please note that our dataset is trained primarily with images that contain a white background. Therefore, images with white background would produce maximum accuracy *")
 
 # Image upload section
 opt = st.selectbox("How do you want to upload the image for classification?", ("Please Select", "Upload image from device"))
@@ -102,8 +103,14 @@ try:
     if image is not None:
         st.image(image, width=256, caption='Uploaded Image')
         if st.button('Predict'):
-            prediction = model.predict(image[np.newaxis, ...])
-            st.success(f'Prediction: {labels[np.argmax(prediction[0], axis=-1)]}')
+            transform = transforms.Compose([
+                transforms.Resize((256, 256)),
+                transforms.ToTensor(),
+                transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+            ])
+            image = transform(image).unsqueeze(0)
+            with torch.no_grad():
+                prediction = model(image)
+            st.success(f'Prediction: {labels[torch.argmax(prediction, dim=1).item()]}')
 except Exception as e:
-    st.error(f"An error occurred: {e}.  Please contact us EcoClim Solutions at EcoClimSolutions.wordpress.com.")
-    
+    st.error(f"An error occurred: {e}. Please contact us EcoClim Solutions at EcoClimSolutions.wordpress.com.")