ysfad commited on
Commit
e15cf70
Β·
verified Β·
1 Parent(s): f599ff2

Upload app.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. app.py +203 -0
app.py ADDED
@@ -0,0 +1,203 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """Gradio app for waste classification using finetuned MAE ViT-Base model."""
3
+
4
+ import os
5
+ import gradio as gr
6
+ from PIL import Image
7
+ from mae_waste_classifier import MAEWasteClassifier
8
+
9
+ print("πŸš€ Initializing MAE waste classifier...")
10
+ try:
11
+ # Load the finetuned MAE model from Hugging Face Hub
12
+ classifier = MAEWasteClassifier(hf_model_id="ysfad/mae-waste-classifier")
13
+ print("βœ… MAE Classifier ready!")
14
+ except Exception as e:
15
+ print(f"❌ Error loading MAE classifier: {e}")
16
+ raise
17
+
18
+ def classify_waste(image):
19
+ """Classify waste item and provide disposal instructions."""
20
+ if image is None:
21
+ return "Please upload an image.", "", "", ""
22
+
23
+ try:
24
+ # Classify the image
25
+ result = classifier.classify_image(image, top_k=5)
26
+
27
+ if not result['success']:
28
+ return f"Error: {result['error']}", "", "", ""
29
+
30
+ # Get model info
31
+ model_info = classifier.get_model_info()
32
+
33
+ # Format main prediction
34
+ main_prediction = f"""
35
+ **🎯 Predicted Class:** {result['predicted_class']}
36
+ **🎲 Confidence:** {result['confidence']:.3f}
37
+ **πŸ€– Model:** {model_info['model_name']}
38
+ **πŸ† Validation Accuracy:** 93.27%
39
+ """
40
+
41
+ # Get disposal instructions
42
+ disposal_text = classifier.get_disposal_instructions(result['predicted_class'])
43
+
44
+ # Format detailed results table
45
+ if result['top_predictions']:
46
+ table_rows = []
47
+ for i, pred in enumerate(result['top_predictions'], 1):
48
+ table_rows.append([
49
+ str(i),
50
+ pred['class'],
51
+ f"{pred['confidence']:.3f}"
52
+ ])
53
+
54
+ # Create HTML table
55
+ table_html = f"""
56
+ <div style="margin-top: 15px;">
57
+ <h4>πŸ” Top {len(result['top_predictions'])} Predictions</h4>
58
+ <table style="width: 100%; border-collapse: collapse;">
59
+ <thead>
60
+ <tr style="background-color: #f0f0f0;">
61
+ <th style="border: 1px solid #ddd; padding: 8px; text-align: left;">#</th>
62
+ <th style="border: 1px solid #ddd; padding: 8px; text-align: left;">Class</th>
63
+ <th style="border: 1px solid #ddd; padding: 8px; text-align: left;">Confidence</th>
64
+ </tr>
65
+ </thead>
66
+ <tbody>
67
+ """
68
+
69
+ for row in table_rows:
70
+ # Color coding based on confidence
71
+ confidence_val = float(row[2])
72
+ if confidence_val > 0.7:
73
+ row_color = "#e8f5e8" # Light green
74
+ elif confidence_val > 0.4:
75
+ row_color = "#fff3cd" # Light yellow
76
+ else:
77
+ row_color = "#f8d7da" # Light red
78
+
79
+ table_html += f"""
80
+ <tr style="background-color: {row_color};">
81
+ <td style="border: 1px solid #ddd; padding: 8px;">{row[0]}</td>
82
+ <td style="border: 1px solid #ddd; padding: 8px;"><strong>{row[1]}</strong></td>
83
+ <td style="border: 1px solid #ddd; padding: 8px;">{row[2]}</td>
84
+ </tr>
85
+ """
86
+
87
+ table_html += """
88
+ </tbody>
89
+ </table>
90
+ </div>
91
+ """
92
+ else:
93
+ table_html = "<p>No predictions available.</p>"
94
+
95
+ # Format model info
96
+ model_info_text = f"""
97
+ **Architecture:** {model_info['architecture']}
98
+ **Pretrained:** {model_info['pretrained']}
99
+ **Classes:** {model_info['num_classes']} waste categories
100
+ **Device:** {model_info['device'].upper()}
101
+ **Training:** Finetuned on RealWaste dataset (4,752 images)
102
+ **Performance:** 93.27% validation accuracy
103
+ **Model Hub:** [ysfad/mae-waste-classifier](https://huggingface.co/ysfad/mae-waste-classifier)
104
+ """
105
+
106
+ return main_prediction, disposal_text, table_html, model_info_text
107
+
108
+ except Exception as e:
109
+ return f"Error during classification: {str(e)}", "", "", ""
110
+
111
+ # Create Gradio interface
112
+ with gr.Blocks(title="πŸ—‚οΈ MAE Waste Classifier", theme=gr.themes.Soft()) as demo:
113
+ gr.Markdown("""
114
+ # πŸ—‚οΈ MAE Waste Classification System
115
+
116
+ Upload an image of waste item to get **classification** and **disposal instructions**.
117
+
118
+ Uses a **finetuned MAE ViT-Base model** achieving **93.27% validation accuracy** on 9 waste categories!
119
+
120
+ **Model:** [ysfad/mae-waste-classifier](https://huggingface.co/ysfad/mae-waste-classifier)
121
+ """)
122
+
123
+ with gr.Row():
124
+ with gr.Column(scale=1):
125
+ # Input section
126
+ gr.Markdown("### πŸ“Έ Upload Image")
127
+ image_input = gr.Image(
128
+ type="pil",
129
+ label="Upload waste item image",
130
+ height=300
131
+ )
132
+
133
+ classify_btn = gr.Button(
134
+ "πŸ” Classify Waste",
135
+ variant="primary",
136
+ size="lg"
137
+ )
138
+
139
+ # Model info section
140
+ gr.Markdown("### πŸ€– Model Information")
141
+ model_info_output = gr.Markdown("")
142
+
143
+ with gr.Column(scale=1):
144
+ # Results section
145
+ gr.Markdown("### 🎯 Classification Results")
146
+ prediction_output = gr.Markdown("")
147
+
148
+ gr.Markdown("### ♻️ Disposal Instructions")
149
+ disposal_output = gr.Textbox(
150
+ label="How to dispose of this item",
151
+ lines=4,
152
+ interactive=False
153
+ )
154
+
155
+ # Detailed results
156
+ gr.Markdown("### πŸ“Š Detailed Results")
157
+ detailed_output = gr.HTML("")
158
+
159
+ # Example images section (if available)
160
+ if os.path.exists("examples"):
161
+ gr.Markdown("### πŸ’‘ Try these examples:")
162
+ gr.Examples(
163
+ examples=[
164
+ ["examples/plastic_bottle.jpg"],
165
+ ["examples/cardboard_box.jpg"],
166
+ ["examples/aluminum_can.jpg"],
167
+ ["examples/glass_bottle.jpg"],
168
+ ["examples/battery.jpg"]
169
+ ],
170
+ inputs=image_input,
171
+ outputs=[prediction_output, disposal_output, detailed_output, model_info_output],
172
+ fn=classify_waste,
173
+ cache_examples=False
174
+ )
175
+
176
+ # Event handlers
177
+ classify_btn.click(
178
+ fn=classify_waste,
179
+ inputs=image_input,
180
+ outputs=[prediction_output, disposal_output, detailed_output, model_info_output]
181
+ )
182
+
183
+ image_input.change(
184
+ fn=classify_waste,
185
+ inputs=image_input,
186
+ outputs=[prediction_output, disposal_output, detailed_output, model_info_output]
187
+ )
188
+
189
+ # Footer
190
+ gr.Markdown("""
191
+ ---
192
+ **πŸ”¬ About:** This system uses a **MAE (Masked Autoencoder) ViT-Base** model finetuned on the RealWaste dataset.
193
+ The model was pretrained with MAE self-supervised learning and then finetuned for waste classification.
194
+
195
+ **⚑ Performance:** Achieved **93.27% validation accuracy** on 9 waste categories with 4,752 training images.
196
+
197
+ **πŸ“Š Categories:** Cardboard, Food Organics, Glass, Metal, Miscellaneous Trash, Paper, Plastic, Textile Trash, Vegetation
198
+
199
+ **πŸ€— Model:** [ysfad/mae-waste-classifier](https://huggingface.co/ysfad/mae-waste-classifier)
200
+ """)
201
+
202
+ if __name__ == "__main__":
203
+ demo.launch()