Update app.py

app.py

@@ -54,21 +54,8 @@ else:
 
 # Model training and testing in separate directory at ipynb file (Copy of ai-portfolio Kendrick.ipynb)
 
-from PIL import Image
-import gradio as gr
-
-# Load model
-def load_model():
-    model = models.resnet50(weights='DEFAULT')  # Using default weights for initialization
-    num_ftrs = model.fc.in_features
-    model.fc = nn.Linear(num_ftrs, 12)  # Adjust to the number of classes you have
-    
-    # Load the state dict
-    model.load_state_dict(torch.load('resnet50_garbage_classificationv1.2.pth', map_location=torch.device('cpu')))
-    
-    model.eval()  # Set to evaluation mode
-    return model
 
+# Load the model
 model = load_model()
 
 # Define image transformations
@@ -79,25 +66,25 @@ transform = transforms.Compose([
     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
 ])
 
-# Class names
+# Class names and corresponding bin images
 class_names = ['battery', 'biological', 'brown-glass', 'cardboard', 
                'clothes', 'green-glass', 'metal', 'paper', 
                'plastic', 'shoes', 'trash', 'white-glass']
 
-# Define bin colors for each class
-bin_colors = {
-    'battery': 'Merah (Red)',                # Limbah berbahaya (B3)
-    'biological': 'Hijau (Green)',           # Limbah organik
-    'brown-glass': 'Kuning (Yellow or trash banks / recycling centers)',  # Gelas berwarna coklat (anorganik/daur ulang)
-    'cardboard': 'Biru (Blue)',              # Kertas (daur ulang)
-    'clothes': 'Kuning atau Bank Sampah (Yellow or trash banks / recycling centers)',  # Pakaian (dimasukkan sebagai daur ulang)
-    'green-glass': 'Kuning (Yellow)',        # Gelas berwarna hijau (anorganik/daur ulang)
-    'metal': 'Kuning (Yellow)',              # Logam (anorganik/daur ulang)
-    'paper': 'Biru (Blue)',                  # Kertas (daur ulang)
-    'plastic': 'Kuning (Yellow)',            # Plastik (anorganik/daur ulang)
-    'shoes': 'Kuning atau Bank Sampah (Yellow or trash banks / recycling centers)',  # Sepatu (dimasukkan sebagai daur ulang)
-    'trash': 'Abu-abu (Gray)',               # Limbah umum
-    'white-glass': 'Kuning (Yellow or trash banks / recycling centers)'         # Gelas berwarna putih (anorganik/daur ulang)
+# Define bin colors and image paths
+bin_info = {
+    'battery': ('Merah (Red)', '/main/red_bin.png'),
+    'biological': ('Hijau (Green)', '/main/green_bin.png'),
+    'brown-glass': ('Kuning (Yellow)', '/main/yellow_bin.png'),
+    'cardboard': ('Biru (Blue)', '/main/blue_bin.png'),
+    'clothes': ('Kuning (Yellow)', '/main/yellow_bin.png'),
+    'green-glass': ('Kuning (Yellow)', '/main/yellow_bin.png'),
+    'metal': ('Kuning (Yellow)', '/main/yellow_bin.png'),
+    'paper': ('Biru (Blue)', '/main/blue_bin.png'),
+    'plastic': ('Kuning (Yellow)', '/main/yellow_bin.png'),
+    'shoes': ('Kuning (Yellow)', '/main/yellow_bin.png'),
+    'trash': ('Abu-abu (Gray)', '/main/gray_bin.png'),
+    'white-glass': ('Kuning (Yellow)', '/main/yellow_bin.png')
 }
 
 # Define the prediction function
@@ -111,23 +98,24 @@ def predict(image):
         _, predicted = torch.max(outputs, 1)
     
     class_name = class_names[predicted.item()]  # Return predicted class name
-    bin_color = bin_colors[class_name]  # Get the corresponding bin color
-    return class_name, bin_color  # Return both class name and bin color
+    bin_color, bin_image = bin_info[class_name]  # Get bin color and image
+    return class_name, bin_color, bin_image  # Return class name, bin color, and bin image
 
-# Make Gradio interface
+# Gradio interface with 3 outputs
 iface = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="numpy", label="Unggah Gambar"),
     outputs=[
-        gr.Textbox(label="Jenis Sampah"), 
-        gr.Textbox(label="Tong Sampah yang Sesuai")  # 2 output with label
+        gr.Textbox(label="Jenis Sampah"),
+        gr.Textbox(label="Tong Sampah yang Sesuai"),
+        gr.Image(label="Gambar Tong Sampah")  # Display bin image
     ],
     title="Klasifikasi Sampah dengan ResNet50 v1.2",
     description="Unggah gambar sampah, dan model kami akan mengklasifikasikannya ke dalam salah satu dari 12 kategori bersama dengan warna tempat sampah yang sesuai. "
                 "<strong>Model ini bisa memprediksi jenis sampah dari ke-12 jenis berikut:</strong> Baterai, Sampah organik, Gelas Kaca Coklat, "
-                "Kardus, Pakaian, Gelas Kaca Hijau, Metal, Kertas, Plastik, Sepatu/sandal, Popok/pampers, Gelas Kaca bening. "
-                "<strong>Note: Untuk masker dan pampers dikategorikan sebagai trash</strong>"
+                "Kardus, Pakaian, Gelas Kaca Hijau, Metal, Kertas, Plastik, Sepatu/sandal, Popok/pampers, Gelas Kaca bening,"
+                "<strong> NB: untuk masker, pampers disebut trash, tapi tong sampah masih sesuai </strong>"
 )
 
-iface.launch(share=True)  
+iface.launch(share=True)