Upload enhanced-waste-classification-webapp.py

enhanced-waste-classification-webapp.py

@@ -0,0 +1,229 @@
+import gradio as gr
+from PIL import Image
+import torch
+from transformers import CLIPProcessor, CLIPModel
+import numpy as np
+import json
+import os
+from datetime import datetime
+
+# Define classification categories
+waste_categories = ["General Waste", "Recyclable Waste"]
+
+# More specific waste item types for tracking
+waste_items = [
+    "Water Bottle", "Plastic Bag", "Food Container", "Paper", "Cardboard",
+    "Glass Bottle", "Aluminum Can", "Food Waste", "Coffee Cup", "Other"
+]
+
+# Load CLIP model and processor (will be loaded from Hugging Face Hub)
+model_name = "openai/clip-vit-base-patch16"
+model = CLIPModel.from_pretrained(model_name)
+processor = CLIPProcessor.from_pretrained(model_name)
+
+# File to store historical data
+HISTORY_FILE = "waste_history.json"
+
+def load_history():
+    """Load historical waste classification data from file."""
+    if os.path.exists(HISTORY_FILE):
+        try:
+            with open(HISTORY_FILE, "r") as f:
+                return json.load(f)
+        except:
+            return {"classifications": []}
+    return {"classifications": []}
+
+def save_history(history):
+    """Save historical waste classification data to file."""
+    with open(HISTORY_FILE, "w") as f:
+        json.dump(history, f)
+
+def classify_waste_category(image):
+    """Classify waste as General or Recyclable using CLIP model."""
+    if image is None:
+        return None, None
+    
+    # Convert to PIL Image if needed
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
+    
+    # Process the image with CLIP for waste category
+    inputs = processor(text=waste_categories, images=image, return_tensors="pt", padding=True)
+    
+    with torch.no_grad():
+        outputs = model(**inputs)
+
+    logits_per_image = outputs.logits_per_image
+    probs = logits_per_image.softmax(dim=1).numpy()[0]
+    
+    # Get the prediction and confidence
+    predicted_class_idx = np.argmax(probs)
+    predicted_category = waste_categories[predicted_class_idx]
+    confidence = probs[predicted_class_idx]
+    
+    return predicted_category, confidence
+
+def classify_waste_item(image):
+    """Identify the specific waste item type using CLIP model."""
+    if image is None:
+        return None, None
+    
+    # Convert to PIL Image if needed
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
+    
+    # Process the image with CLIP for specific waste item
+    inputs = processor(text=waste_items, images=image, return_tensors="pt", padding=True)
+    
+    with torch.no_grad():
+        outputs = model(**inputs)
+
+    logits_per_image = outputs.logits_per_image
+    probs = logits_per_image.softmax(dim=1).numpy()[0]
+    
+    # Get the prediction and confidence
+    predicted_item_idx = np.argmax(probs)
+    predicted_item = waste_items[predicted_item_idx]
+    confidence = probs[predicted_item_idx]
+    
+    return predicted_item, confidence
+
+def process_image(image, waste_item_selected=None):
+    """Main function to process images from webcam or upload."""
+    if image is None:
+        return "No image provided", None, None, None
+    
+    # Classify waste category (General or Recyclable)
+    category, category_confidence = classify_waste_category(image)
+    
+    # Get specific waste item type
+    if waste_item_selected and waste_item_selected != "Auto-detect":
+        item = waste_item_selected
+        item_confidence = 1.0  # User manually selected, so confidence is 100%
+    else:
+        item, item_confidence = classify_waste_item(image)
+    
+    # Format the result
+    result = f"This appears to be: {item}\nClassified as: {category}\n"
+    result += f"Category confidence: {category_confidence:.2f}"
+    
+    # Create dictionaries for visualizations
+    category_probs = {
+        waste_categories[0]: 1-category_confidence if category == waste_categories[1] else category_confidence,
+        waste_categories[1]: category_confidence if category == waste_categories[1] else 1-category_confidence
+    }
+    
+    # Record this classification in history
+    history = load_history()
+    history["classifications"].append({
+        "timestamp": datetime.now().isoformat(),
+        "waste_item": item,
+        "waste_category": category,
+        "confidence": float(category_confidence)
+    })
+    save_history(history)
+    
+    # Generate statistics for the history tab
+    item_counts = {}
+    category_counts = {"General Waste": 0, "Recyclable Waste": 0}
+    
+    for entry in history["classifications"]:
+        item = entry["waste_item"]
+        category = entry["waste_category"]
+        
+        if item not in item_counts:
+            item_counts[item] = {"General Waste": 0, "Recyclable Waste": 0}
+        
+        item_counts[item][category] += 1
+        category_counts[category] += 1
+    
+    # Format history data for visualization
+    history_data = []
+    for item, counts in item_counts.items():
+        for category, count in counts.items():
+            if count > 0:
+                history_data.append({"item": item, "category": category, "count": count})
+    
+    return result, category_probs, history_data, category_counts
+
+# Create Gradio interface
+with gr.Blocks(title="EcoCan - Waste Classification") as demo:
+    gr.Markdown("""
+    # EcoCan - AI Waste Classification
+    
+    Use your webcam or phone camera to take a photo of waste and classify it as General Waste or Recyclable Waste.
+    """)
+    
+    with gr.Tabs():
+        with gr.Tab("Classify Waste"):
+            with gr.Row():
+                with gr.Column():
+                    # Input options: webcam, upload, or file upload
+                    input_image = gr.Image(
+                        sources=["webcam", "upload"], 
+                        type="numpy", 
+                        label="Take a photo or upload an image"
+                    )
+                    
+                    # Optional manual selection
+                    waste_item_dropdown = gr.Dropdown(
+                        choices=["Auto-detect"] + waste_items,
+                        value="Auto-detect",
+                        label="Or manually select waste item type (optional)"
+                    )
+                    
+                    # Classify button
+                    classify_button = gr.Button("Classify Waste", variant="primary")
+                
+                with gr.Column():
+                    # Output components
+                    result_text = gr.Textbox(label="Classification Result")
+                    confidence_chart = gr.BarPlot(
+                        x="category", 
+                        y="confidence", 
+                        title="Classification Confidence",
+                        y_lim=[0, 1],
+                        tooltip=["category", "confidence"],
+                        color="category"
+                    )
+        
+        with gr.Tab("History"):
+            gr.Markdown("### Waste Classification History")
+            
+            with gr.Row():
+                with gr.Column():
+                    history_chart = gr.BarPlot(
+                        x="item",
+                        y="count",
+                        color="category",
+                        group="category",
+                        title="Items Classified by Category",
+                        tooltip=["item", "category", "count"],
+                        height=400
+                    )
+                
+                with gr.Column():
+                    category_pie = gr.PieChart(
+                        label="Waste Categories Distribution",
+                        type="value",
+                        height=400
+                    )
+    
+    # Set up event handlers
+    classify_button.click(
+        process_image, 
+        inputs=[input_image, waste_item_dropdown], 
+        outputs=[result_text, confidence_chart, history_chart, category_pie]
+    )
+    
+    # Also process when image is captured without clicking button
+    input_image.change(
+        process_image,
+        inputs=[input_image, waste_item_dropdown],
+        outputs=[result_text, confidence_chart, history_chart, category_pie]
+    )
+
+# Launch locally for testing
+if __name__ == "__main__":
+    demo.launch()