Upload app.py with huggingface_hub

app.py

@@ -0,0 +1,203 @@
+#!/usr/bin/env python3
+"""Gradio app for waste classification using finetuned MAE ViT-Base model."""
+
+import os
+import gradio as gr
+from PIL import Image
+from mae_waste_classifier import MAEWasteClassifier
+
+print("🚀 Initializing MAE waste classifier...")
+try:
+    # Load the finetuned MAE model from Hugging Face Hub
+    classifier = MAEWasteClassifier(hf_model_id="ysfad/mae-waste-classifier")
+    print("✅ MAE Classifier ready!")
+except Exception as e:
+    print(f"❌ Error loading MAE classifier: {e}")
+    raise
+
+def classify_waste(image):
+    """Classify waste item and provide disposal instructions."""
+    if image is None:
+        return "Please upload an image.", "", "", ""
+    
+    try:
+        # Classify the image
+        result = classifier.classify_image(image, top_k=5)
+        
+        if not result['success']:
+            return f"Error: {result['error']}", "", "", ""
+        
+        # Get model info
+        model_info = classifier.get_model_info()
+        
+        # Format main prediction
+        main_prediction = f"""
+        **🎯 Predicted Class:** {result['predicted_class']}
+        **🎲 Confidence:** {result['confidence']:.3f}
+        **🤖 Model:** {model_info['model_name']}
+        **🏆 Validation Accuracy:** 93.27%
+        """
+        
+        # Get disposal instructions
+        disposal_text = classifier.get_disposal_instructions(result['predicted_class'])
+        
+        # Format detailed results table
+        if result['top_predictions']:
+            table_rows = []
+            for i, pred in enumerate(result['top_predictions'], 1):
+                table_rows.append([
+                    str(i),
+                    pred['class'],
+                    f"{pred['confidence']:.3f}"
+                ])
+            
+            # Create HTML table
+            table_html = f"""
+            <div style="margin-top: 15px;">
+                <h4>🔍 Top {len(result['top_predictions'])} Predictions</h4>
+                <table style="width: 100%; border-collapse: collapse;">
+                    <thead>
+                        <tr style="background-color: #f0f0f0;">
+                            <th style="border: 1px solid #ddd; padding: 8px; text-align: left;">#</th>
+                            <th style="border: 1px solid #ddd; padding: 8px; text-align: left;">Class</th>
+                            <th style="border: 1px solid #ddd; padding: 8px; text-align: left;">Confidence</th>
+                        </tr>
+                    </thead>
+                    <tbody>
+            """
+            
+            for row in table_rows:
+                # Color coding based on confidence
+                confidence_val = float(row[2])
+                if confidence_val > 0.7:
+                    row_color = "#e8f5e8"  # Light green
+                elif confidence_val > 0.4:
+                    row_color = "#fff3cd"  # Light yellow
+                else:
+                    row_color = "#f8d7da"  # Light red
+                
+                table_html += f"""
+                        <tr style="background-color: {row_color};">
+                            <td style="border: 1px solid #ddd; padding: 8px;">{row[0]}</td>
+                            <td style="border: 1px solid #ddd; padding: 8px;"><strong>{row[1]}</strong></td>
+                            <td style="border: 1px solid #ddd; padding: 8px;">{row[2]}</td>
+                        </tr>
+                """
+            
+            table_html += """
+                    </tbody>
+                </table>
+            </div>
+            """
+        else:
+            table_html = "<p>No predictions available.</p>"
+        
+        # Format model info
+        model_info_text = f"""
+        **Architecture:** {model_info['architecture']} 
+        **Pretrained:** {model_info['pretrained']}
+        **Classes:** {model_info['num_classes']} waste categories
+        **Device:** {model_info['device'].upper()}
+        **Training:** Finetuned on RealWaste dataset (4,752 images)
+        **Performance:** 93.27% validation accuracy
+        **Model Hub:** [ysfad/mae-waste-classifier](https://huggingface.co/ysfad/mae-waste-classifier)
+        """
+        
+        return main_prediction, disposal_text, table_html, model_info_text
+        
+    except Exception as e:
+        return f"Error during classification: {str(e)}", "", "", ""
+
+# Create Gradio interface
+with gr.Blocks(title="🗂️ MAE Waste Classifier", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("""
+    # 🗂️ MAE Waste Classification System
+    
+    Upload an image of waste item to get **classification** and **disposal instructions**.
+    
+    Uses a **finetuned MAE ViT-Base model** achieving **93.27% validation accuracy** on 9 waste categories!
+    
+    **Model:** [ysfad/mae-waste-classifier](https://huggingface.co/ysfad/mae-waste-classifier)
+    """)
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            # Input section
+            gr.Markdown("### 📸 Upload Image")
+            image_input = gr.Image(
+                type="pil",
+                label="Upload waste item image",
+                height=300
+            )
+            
+            classify_btn = gr.Button(
+                "🔍 Classify Waste",
+                variant="primary",
+                size="lg"
+            )
+            
+            # Model info section
+            gr.Markdown("### 🤖 Model Information")
+            model_info_output = gr.Markdown("")
+        
+        with gr.Column(scale=1):
+            # Results section
+            gr.Markdown("### 🎯 Classification Results")
+            prediction_output = gr.Markdown("")
+            
+            gr.Markdown("### ♻️ Disposal Instructions")
+            disposal_output = gr.Textbox(
+                label="How to dispose of this item",
+                lines=4,
+                interactive=False
+            )
+            
+            # Detailed results
+            gr.Markdown("### 📊 Detailed Results")
+            detailed_output = gr.HTML("")
+    
+    # Example images section (if available)
+    if os.path.exists("examples"):
+        gr.Markdown("### 💡 Try these examples:")
+        gr.Examples(
+            examples=[
+                ["examples/plastic_bottle.jpg"],
+                ["examples/cardboard_box.jpg"],
+                ["examples/aluminum_can.jpg"],
+                ["examples/glass_bottle.jpg"],
+                ["examples/battery.jpg"]
+            ],
+            inputs=image_input,
+            outputs=[prediction_output, disposal_output, detailed_output, model_info_output],
+            fn=classify_waste,
+            cache_examples=False
+        )
+    
+    # Event handlers
+    classify_btn.click(
+        fn=classify_waste,
+        inputs=image_input,
+        outputs=[prediction_output, disposal_output, detailed_output, model_info_output]
+    )
+    
+    image_input.change(
+        fn=classify_waste,
+        inputs=image_input,
+        outputs=[prediction_output, disposal_output, detailed_output, model_info_output]
+    )
+    
+    # Footer
+    gr.Markdown("""
+    ---
+    **🔬 About:** This system uses a **MAE (Masked Autoencoder) ViT-Base** model finetuned on the RealWaste dataset.
+    The model was pretrained with MAE self-supervised learning and then finetuned for waste classification.
+    
+    **⚡ Performance:** Achieved **93.27% validation accuracy** on 9 waste categories with 4,752 training images.
+    
+    **📊 Categories:** Cardboard, Food Organics, Glass, Metal, Miscellaneous Trash, Paper, Plastic, Textile Trash, Vegetation
+    
+    **🤗 Model:** [ysfad/mae-waste-classifier](https://huggingface.co/ysfad/mae-waste-classifier)
+    """)
+
+if __name__ == "__main__":
+    demo.launch()