Update app.py
Browse files
app.py
CHANGED
@@ -54,21 +54,8 @@ else:
|
|
54 |
|
55 |
# Model training and testing in separate directory at ipynb file (Copy of ai-portfolio Kendrick.ipynb)
|
56 |
|
57 |
-
from PIL import Image
|
58 |
-
import gradio as gr
|
59 |
-
|
60 |
-
# Load model
|
61 |
-
def load_model():
|
62 |
-
model = models.resnet50(weights='DEFAULT') # Using default weights for initialization
|
63 |
-
num_ftrs = model.fc.in_features
|
64 |
-
model.fc = nn.Linear(num_ftrs, 12) # Adjust to the number of classes you have
|
65 |
-
|
66 |
-
# Load the state dict
|
67 |
-
model.load_state_dict(torch.load('resnet50_garbage_classificationv1.2.pth', map_location=torch.device('cpu')))
|
68 |
-
|
69 |
-
model.eval() # Set to evaluation mode
|
70 |
-
return model
|
71 |
|
|
|
72 |
model = load_model()
|
73 |
|
74 |
# Define image transformations
|
@@ -79,25 +66,25 @@ transform = transforms.Compose([
|
|
79 |
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
|
80 |
])
|
81 |
|
82 |
-
# Class names
|
83 |
class_names = ['battery', 'biological', 'brown-glass', 'cardboard',
|
84 |
'clothes', 'green-glass', 'metal', 'paper',
|
85 |
'plastic', 'shoes', 'trash', 'white-glass']
|
86 |
|
87 |
-
# Define bin colors
|
88 |
-
|
89 |
-
'battery': 'Merah (Red)',
|
90 |
-
'biological': 'Hijau (Green)',
|
91 |
-
'brown-glass': 'Kuning (Yellow
|
92 |
-
'cardboard': 'Biru (Blue)',
|
93 |
-
'clothes': 'Kuning
|
94 |
-
'green-glass': 'Kuning (Yellow)',
|
95 |
-
'metal': 'Kuning (Yellow)',
|
96 |
-
'paper': 'Biru (Blue)',
|
97 |
-
'plastic': 'Kuning (Yellow)',
|
98 |
-
'shoes': 'Kuning
|
99 |
-
'trash': 'Abu-abu (Gray)',
|
100 |
-
'white-glass': 'Kuning (Yellow
|
101 |
}
|
102 |
|
103 |
# Define the prediction function
|
@@ -111,23 +98,24 @@ def predict(image):
|
|
111 |
_, predicted = torch.max(outputs, 1)
|
112 |
|
113 |
class_name = class_names[predicted.item()] # Return predicted class name
|
114 |
-
bin_color =
|
115 |
-
return class_name, bin_color # Return
|
116 |
|
117 |
-
#
|
118 |
iface = gr.Interface(
|
119 |
fn=predict,
|
120 |
inputs=gr.Image(type="numpy", label="Unggah Gambar"),
|
121 |
outputs=[
|
122 |
-
gr.Textbox(label="Jenis Sampah"),
|
123 |
-
gr.Textbox(label="Tong Sampah yang Sesuai")
|
|
|
124 |
],
|
125 |
title="Klasifikasi Sampah dengan ResNet50 v1.2",
|
126 |
description="Unggah gambar sampah, dan model kami akan mengklasifikasikannya ke dalam salah satu dari 12 kategori bersama dengan warna tempat sampah yang sesuai. "
|
127 |
"<strong>Model ini bisa memprediksi jenis sampah dari ke-12 jenis berikut:</strong> Baterai, Sampah organik, Gelas Kaca Coklat, "
|
128 |
-
"Kardus, Pakaian, Gelas Kaca Hijau, Metal, Kertas, Plastik, Sepatu/sandal, Popok/pampers, Gelas Kaca bening
|
129 |
-
"<strong>
|
130 |
)
|
131 |
|
132 |
-
iface.launch(share=True)
|
133 |
|
|
|
54 |
|
55 |
# Model training and testing in separate directory at ipynb file (Copy of ai-portfolio Kendrick.ipynb)
|
56 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
57 |
|
58 |
+
# Load the model
|
59 |
model = load_model()
|
60 |
|
61 |
# Define image transformations
|
|
|
66 |
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
|
67 |
])
|
68 |
|
69 |
+
# Class names and corresponding bin images
|
70 |
class_names = ['battery', 'biological', 'brown-glass', 'cardboard',
|
71 |
'clothes', 'green-glass', 'metal', 'paper',
|
72 |
'plastic', 'shoes', 'trash', 'white-glass']
|
73 |
|
74 |
+
# Define bin colors and image paths
|
75 |
+
bin_info = {
|
76 |
+
'battery': ('Merah (Red)', '/main/red_bin.png'),
|
77 |
+
'biological': ('Hijau (Green)', '/main/green_bin.png'),
|
78 |
+
'brown-glass': ('Kuning (Yellow)', '/main/yellow_bin.png'),
|
79 |
+
'cardboard': ('Biru (Blue)', '/main/blue_bin.png'),
|
80 |
+
'clothes': ('Kuning (Yellow)', '/main/yellow_bin.png'),
|
81 |
+
'green-glass': ('Kuning (Yellow)', '/main/yellow_bin.png'),
|
82 |
+
'metal': ('Kuning (Yellow)', '/main/yellow_bin.png'),
|
83 |
+
'paper': ('Biru (Blue)', '/main/blue_bin.png'),
|
84 |
+
'plastic': ('Kuning (Yellow)', '/main/yellow_bin.png'),
|
85 |
+
'shoes': ('Kuning (Yellow)', '/main/yellow_bin.png'),
|
86 |
+
'trash': ('Abu-abu (Gray)', '/main/gray_bin.png'),
|
87 |
+
'white-glass': ('Kuning (Yellow)', '/main/yellow_bin.png')
|
88 |
}
|
89 |
|
90 |
# Define the prediction function
|
|
|
98 |
_, predicted = torch.max(outputs, 1)
|
99 |
|
100 |
class_name = class_names[predicted.item()] # Return predicted class name
|
101 |
+
bin_color, bin_image = bin_info[class_name] # Get bin color and image
|
102 |
+
return class_name, bin_color, bin_image # Return class name, bin color, and bin image
|
103 |
|
104 |
+
# Gradio interface with 3 outputs
|
105 |
iface = gr.Interface(
|
106 |
fn=predict,
|
107 |
inputs=gr.Image(type="numpy", label="Unggah Gambar"),
|
108 |
outputs=[
|
109 |
+
gr.Textbox(label="Jenis Sampah"),
|
110 |
+
gr.Textbox(label="Tong Sampah yang Sesuai"),
|
111 |
+
gr.Image(label="Gambar Tong Sampah") # Display bin image
|
112 |
],
|
113 |
title="Klasifikasi Sampah dengan ResNet50 v1.2",
|
114 |
description="Unggah gambar sampah, dan model kami akan mengklasifikasikannya ke dalam salah satu dari 12 kategori bersama dengan warna tempat sampah yang sesuai. "
|
115 |
"<strong>Model ini bisa memprediksi jenis sampah dari ke-12 jenis berikut:</strong> Baterai, Sampah organik, Gelas Kaca Coklat, "
|
116 |
+
"Kardus, Pakaian, Gelas Kaca Hijau, Metal, Kertas, Plastik, Sepatu/sandal, Popok/pampers, Gelas Kaca bening,"
|
117 |
+
"<strong> NB: untuk masker, pampers disebut trash, tapi tong sampah masih sesuai </strong>"
|
118 |
)
|
119 |
|
120 |
+
iface.launch(share=True)
|
121 |
|