kendrickfff commited on
Commit
cf6d9c5
·
verified ·
1 Parent(s): b616460

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +25 -37
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 for each class
88
- bin_colors = {
89
- 'battery': 'Merah (Red)', # Limbah berbahaya (B3)
90
- 'biological': 'Hijau (Green)', # Limbah organik
91
- 'brown-glass': 'Kuning (Yellow or trash banks / recycling centers)', # Gelas berwarna coklat (anorganik/daur ulang)
92
- 'cardboard': 'Biru (Blue)', # Kertas (daur ulang)
93
- 'clothes': 'Kuning atau Bank Sampah (Yellow or trash banks / recycling centers)', # Pakaian (dimasukkan sebagai daur ulang)
94
- 'green-glass': 'Kuning (Yellow)', # Gelas berwarna hijau (anorganik/daur ulang)
95
- 'metal': 'Kuning (Yellow)', # Logam (anorganik/daur ulang)
96
- 'paper': 'Biru (Blue)', # Kertas (daur ulang)
97
- 'plastic': 'Kuning (Yellow)', # Plastik (anorganik/daur ulang)
98
- 'shoes': 'Kuning atau Bank Sampah (Yellow or trash banks / recycling centers)', # Sepatu (dimasukkan sebagai daur ulang)
99
- 'trash': 'Abu-abu (Gray)', # Limbah umum
100
- 'white-glass': 'Kuning (Yellow or trash banks / recycling centers)' # Gelas berwarna putih (anorganik/daur ulang)
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 = bin_colors[class_name] # Get the corresponding bin color
115
- return class_name, bin_color # Return both class name and bin color
116
 
117
- # Make Gradio interface
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") # 2 output with label
 
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>Note: Untuk masker dan pampers dikategorikan sebagai trash</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