Update app.py

app.py

@@ -1,34 +1,35 @@
 import os
+import base64
+import tempfile
 import tensorflow as tf
-from flask import Flask, request, render_template, redirect, url_for
-from werkzeug.utils import secure_filename
+from flask import Flask, request, render_template, redirect
+from io import BytesIO
 
 # Initialize the Flask application
 app = Flask(__name__)
 
-# --- Load the Clean, Compatible .h5 Model ---
-# This model was saved with save_format='h5' for maximum compatibility.
-MODEL_PATH = 'waste_classifier_final_5.h5'
+# --- Load the Model ---
+# This now points to the directory created by model.export()
+MODEL_PATH = 'waste_classifier_final_5.h5' # IMPORTANT: Ensure this matches your uploaded model's filename
 try:
     model = tf.keras.models.load_model(MODEL_PATH)
-    print("Image classification model loaded successfully!")
+    print("✅ Image classification model loaded successfully!")
 except Exception as e:
-    print(f"Error loading image model: {e}")
+    print(f"❌ Error loading image model: {e}")
     exit()
 
-# Define the class names in the correct order for the model's output
+# --- CRITICAL: Ensure this list EXACTLY matches the output from your training script ---
+# Example: ['cardboard', 'glass', 'metal', 'paper', 'plastic', 'trash']
 CLASS_NAMES = ['cardboard', 'glass', 'metal', 'paper', 'plastic', 'trash']
 
+
 def preprocess_image(image_path):
-    """
-    Loads an image from a file path and preprocesses it for the model.
-    This function ensures the input image matches the format used during training.
-    """
+    """Loads and preprocesses an image for the model."""
     img = tf.keras.preprocessing.image.load_img(image_path, target_size=(224, 224))
     img_array = tf.keras.preprocessing.image.img_to_array(img)
-    img_array = tf.expand_dims(img_array, 0) # Create a batch of one
-    # Apply the MobileNetV2-specific preprocessing
-    return tf.keras.applications.mobilenet_v2.preprocess_input(img_array)
+    img_array = tf.expand_dims(img_array, 0)
+    # --- UPDATED: Switched to the correct preprocessing for EfficientNet ---
+    return tf.keras.applications.efficientnet.preprocess_input(img_array)
 
 @app.route('/', methods=['GET'])
 def index():
@@ -37,8 +38,7 @@ def index():
 
 @app.route('/predict', methods=['POST'])
 def predict():
-    """Handles the image upload, prediction, and renders the result."""
-    # Check if a file was uploaded
+    """Handles image upload, prediction, and renders the result."""
     if 'file' not in request.files:
         return redirect(request.url)
     file = request.files['file']
@@ -46,30 +46,31 @@ def predict():
         return redirect(request.url)
 
     if file:
-        # Save the file securely
-        filename = secure_filename(file.filename)
-        filepath = os.path.join('static/uploads', filename)
-        file.save(filepath)
+        with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(file.filename)[1]) as tmp_file:
+            filepath = tmp_file.name
+            file.save(filepath)
+
+        with open(filepath, "rb") as f:
+            image_data = f.read()
+        
+        encoded_image = base64.b64encode(image_data).decode('utf-8')
+        image_to_display = f"data:image/jpeg;base64,{encoded_image}"
 
-        # Preprocess the image and get a prediction
         preprocessed_image = preprocess_image(filepath)
         prediction = model.predict(preprocessed_image)
         
-        # Decode the prediction
         predicted_class_index = tf.argmax(prediction[0]).numpy()
         predicted_class = CLASS_NAMES[predicted_class_index]
         confidence = tf.reduce_max(prediction[0]).numpy() * 100
 
-        # Pass the results to the HTML template
+        os.remove(filepath)
+
         return render_template('index.html',
                                prediction=f'Prediction: {predicted_class}',
                                confidence=f'Confidence: {confidence:.2f}%',
-                               uploaded_image=filepath)
+                               uploaded_image=image_to_display)
+    
     return redirect(request.url)
 
 if __name__ == '__main__':
-    # Ensure the upload folder exists
-    os.makedirs('static/uploads', exist_ok=True)
-    # This host and port configuration is important for deployment services like Hugging Face
-    app.run(host='0.0.0.0', port=int(os.environ.get('PORT', 7860)))
-
+    app.run(host='0.0.0.0', port=int(os.environ.get('PORT', 7860)))