Update app.py

app.py

@@ -1,65 +1,52 @@
-import os, tempfile
+import matplotlib
+matplotlib.use('Agg')  # Non-GUI backend for plotting
+
+import os
 import cv2
 import numpy as np
-import pandas as pd
-from datetime import datetime
-from flask import Flask, render_template, request, redirect, send_file
+from flask import Flask, render_template, request, redirect, send_file, Response
 from tensorflow.keras.applications.mobilenet_v2 import preprocess_input
 from tensorflow.keras.models import load_model
 from fpdf import FPDF
 import matplotlib.pyplot as plt
+# from ultralytics import YOLO  # Commented out for now
 
 # -----------------------------
-# Environment fixes for Hugging Face
-# -----------------------------
-# Fix matplotlib cache dir
-os.environ["MPLCONFIGDIR"] = "/tmp/matplotlib"
-os.makedirs("/tmp/matplotlib", exist_ok=True)
-
-# Uploads go into /tmp
-UPLOAD_DIR = os.path.join(tempfile.gettempdir(), "uploads")
-os.makedirs(UPLOAD_DIR, exist_ok=True)
-
-# Logs and reports also in /tmp
-LOG_PATH = os.path.join(tempfile.gettempdir(), "waste_log.csv")
-REPORT_PDF = os.path.join(tempfile.gettempdir(), "waste_report.pdf")
-CHART_PATH = os.path.join(tempfile.gettempdir(), "stats_chart.png")
-
-# -----------------------------
-# Flask App Config
+# Flask Config
 # -----------------------------
 app = Flask(__name__)
-app.config["UPLOAD_FOLDER"] = UPLOAD_DIR
+app.config["UPLOAD_FOLDER"] = "static/uploads"
+os.makedirs(app.config["UPLOAD_FOLDER"], exist_ok=True)
 
 # -----------------------------
-# Model path
-# -----------------------------
-MODEL_DIR = "model"
-os.makedirs(MODEL_DIR, exist_ok=True)
-MODEL_PATH = os.path.join(MODEL_DIR, "/efficientnet_b0_best.keras")
-
-if not os.path.exists(MODEL_PATH):
-    raise FileNotFoundError(f"Model file not found at {MODEL_PATH}. Place your .keras model in /model.")
-
-best_model = load_model(MODEL_PATH)
+# Load Keras classification model
+try:
+    best_model = load_model("efficientnet_b0_best.keras")
+    print("✅ Model loaded successfully!")
+except Exception as e:
+    print(f"❌ Error loading model: {e}")
+    best_model = None
 
 IMG_SIZE = 128
 
-# Class labels
 CLASS_LABELS = ['biological', 'brown-glass', 'cardboard', 'green-glass',
                 'metal', 'paper', 'plastic', 'shoes', 'trash', 'white-glass']
 
-# Waste categories
 RECYCLABLE = ["brown-glass", "green-glass", "white-glass", "metal", "plastic", "paper", "cardboard"]
 NON_RECYCLABLE = ["trash", "biological", "shoes"]
 
-# Initialize statistics
 stats = {}
 
 # -----------------------------
-# Preprocess image
+# Load YOLOv8 model (disabled for Hugging Face Spaces)
+# yolo_model = YOLO("best.pt")  # Uncomment and add your YOLO model file if needed
+
+# -----------------------------
+# Preprocess image for classification
 def preprocess_image(file_path):
     img = cv2.imread(file_path)
+    if img is None:
+        raise ValueError("Could not load image")
     img_rgb = cv2.cvtColor(cv2.resize(img, (IMG_SIZE, IMG_SIZE)), cv2.COLOR_BGR2RGB)
     img_input = preprocess_input(img_rgb.astype("float32"))
     img_input = np.expand_dims(img_input, axis=0)
@@ -68,103 +55,150 @@ def preprocess_image(file_path):
 # -----------------------------
 # Log predictions
 def log_prediction(class_label):
-    log_df = pd.DataFrame([[datetime.now(), class_label]], columns=["Timestamp", "Class"])
+    stats[class_label] = stats.get(class_label, 0) + 1
+    total_items = sum(stats.values())
     try:
-        old_df = pd.read_csv(LOG_PATH)
-        new_df = pd.concat([old_df, log_df], ignore_index=True)
-    except FileNotFoundError:
-        new_df = log_df
-    new_df.to_csv(LOG_PATH, index=False)
+        with open("waste_log.csv", "w") as f:
+            f.write("Waste Classification Report\n")
+            f.write(f"Total Items Processed: {total_items}\n")
+            for category, count in stats.items():
+                f.write(f"{category}: {count}\n")
+    except Exception as e:
+        print(f"Error writing log: {e}")
 
 # -----------------------------
-# PDF Report
+# Generate PDF report
 def generate_pdf_report():
-    pdf = FPDF()
-    pdf.add_page()
-    pdf.set_font("Arial", size=14)
-    pdf.cell(200, 10, txt="Waste Classification Report", ln=True, align="C")
-    pdf.ln(10)
-
-    total_items = sum(stats.values()) if stats else 0
-    pdf.set_font("Arial", size=12)
-    pdf.cell(0, 10, txt=f"Total Items Processed: {total_items}", ln=True)
-
-    for category, count in stats.items():
-        pdf.cell(0, 10, txt=f"{category}: {count}", ln=True)
-
-    pdf.output(REPORT_PDF)
-    return REPORT_PDF
+    try:
+        pdf = FPDF()
+        pdf.add_page()
+        pdf.set_font("Arial", size=14)
+        pdf.cell(200, 10, txt="Waste Classification Report", ln=True, align="C")
+        pdf.ln(10)
+
+        total_items = sum(stats.values())
+        pdf.set_font("Arial", size=12)
+        pdf.cell(0, 10, txt=f"Total Items Processed: {total_items}", ln=True)
+
+        for category, count in stats.items():
+            pdf.cell(0, 10, txt=f"{category}: {count}", ln=True)
+
+        pdf_file = "waste_report.pdf"
+        pdf.output(pdf_file)
+        return pdf_file
+    except Exception as e:
+        print(f"Error generating PDF: {e}")
+        return None
 
 # -----------------------------
-# Routes
+# Image classification route
 @app.route("/", methods=["GET", "POST"])
 def index():
     if request.method == "POST":
         file = request.files.get("file")
-        if file is None or file.filename == "":
+        if not file or file.filename == "":
             return redirect(request.url)
 
-        file_path = os.path.join(app.config["UPLOAD_FOLDER"], file.filename)
-        file.save(file_path)
-
-        # Preprocess & predict
-        img_rgb, img_input = preprocess_image(file_path)
-        preds = best_model.predict(img_input)
-        class_idx = np.argmax(preds, axis=1)[0]
-        class_label = CLASS_LABELS[class_idx]
-        confidence = preds[0][class_idx]
-
-        # Update stats & log
-        stats[class_label] = stats.get(class_label, 0) + 1
-        log_prediction(class_label)
-
-        # Determine bin type
-        if class_label in RECYCLABLE:
-            bin_type = "Recyclable ♻️"
-        elif class_label in NON_RECYCLABLE:
-            bin_type = "Non-Recyclable 🗑️"
-        else:
-            bin_type = "Unknown ⚠️"
-
-        return render_template(
-            "result.html",
-            image=file.filename,
-            label=class_label,
-            confidence=f"{confidence*100:.2f}%",
-            bin_type=bin_type
-        )
+        # Validate file type
+        allowed_extensions = {'png', 'jpg', 'jpeg', 'gif', 'bmp'}
+        file_extension = file.filename.rsplit('.', 1)[1].lower() if '.' in file.filename else ''
+        if file_extension not in allowed_extensions:
+            return render_template("index.html", error="Please upload a valid image file (PNG, JPG, JPEG, GIF, BMP)")
+
+        try:
+            file_path = os.path.join(app.config["UPLOAD_FOLDER"], file.filename)
+            file.save(file_path)
+
+            if best_model is None:
+                return render_template("index.html", error="Model not loaded. Please check if the model file exists.")
+
+            img_rgb, img_input = preprocess_image(file_path)
+            preds = best_model.predict(img_input)
+            class_idx = np.argmax(preds, axis=1)[0]
+            class_label = CLASS_LABELS[class_idx]
+            confidence = preds[0][class_idx]
+
+            log_prediction(class_label)
+
+            if class_label in RECYCLABLE:
+                bin_type = "Recyclable ♻️"
+            elif class_label in NON_RECYCLABLE:
+                bin_type = "Non-Recyclable 🗑️"
+            else:
+                bin_type = "Unknown ⚠️"
+
+            return render_template(
+                "result.html",
+                image=file.filename,
+                label=class_label,
+                confidence=f"{confidence*100:.2f}%",
+                bin_type=bin_type
+            )
+        except Exception as e:
+            return render_template("index.html", error=f"Error processing image: {str(e)}")
 
     return render_template("index.html")
 
-
+# -----------------------------
+# Show statistics
 @app.route("/stats")
 def show_stats():
     if stats:
-        categories = list(stats.keys())
-        counts = list(stats.values())
-        plt.figure(figsize=(6, 4))
-        plt.bar(categories, counts)
-        plt.xlabel("Category")
-        plt.ylabel("Count")
-        plt.title("Waste Classification Statistics")
-        plt.tight_layout()
-        plt.savefig(CHART_PATH)
-        plt.close()
-    return render_template("report.html", stats=stats, chart_path=CHART_PATH)
-
+        try:
+            categories = list(stats.keys())
+            counts = list(stats.values())
+            plt.figure(figsize=(10, 6))
+            plt.bar(categories, counts, color="green", alpha=0.7)
+            plt.xlabel("Category")
+            plt.ylabel("Count")
+            plt.title("Waste Classification Statistics")
+            plt.xticks(rotation=45)
+            plt.tight_layout()
+            
+            # Ensure static directory exists
+            os.makedirs("static", exist_ok=True)
+            plt.savefig("static/stats_chart.png", dpi=150, bbox_inches='tight')
+            plt.close()
+        except Exception as e:
+            print(f"Error generating chart: {e}")
+    
+    return render_template("report.html", stats=stats)
 
+# -----------------------------
+# Download reports
 @app.route("/download_pdf")
 def download_pdf():
-    pdf_path = generate_pdf_report()
-    return send_file(pdf_path, as_attachment=True)
-
+    try:
+        pdf_path = generate_pdf_report()
+        if pdf_path and os.path.exists(pdf_path):
+            return send_file(pdf_path, as_attachment=True)
+        else:
+            return "Error generating PDF report", 500
+    except Exception as e:
+        return f"Error: {str(e)}", 500
 
 @app.route("/download_csv")
 def download_csv():
-    return send_file(LOG_PATH, as_attachment=True)
+    try:
+        if os.path.exists("waste_log.csv"):
+            return send_file("waste_log.csv", as_attachment=True)
+        else:
+            return "No data to download", 404
+    except Exception as e:
+        return f"Error: {str(e)}", 500
 
 # -----------------------------
-# Run Flask (useful locally)
+# Real-time camera detection (disabled for Hugging Face Spaces)
+@app.route("/camera")
+def camera():
+    return render_template("camera_disabled.html")
+
+# Health check endpoint
+@app.route("/health")
+def health():
+    return {"status": "healthy", "model_loaded": best_model is not None}
+
+# Run Flask
 if __name__ == "__main__":
-    port = int(os.environ.get("PORT", 8080))
-    app.run(host="0.0.0.0", port=port, debug=True)
+    port = int(os.environ.get("PORT", 7860))  # Hugging Face Spaces uses port 7860
+    app.run(host="0.0.0.0", port=port, debug=False)