app.py

import os
import numpy as np
import tensorflow as tf
from flask import Flask, request, render_template, jsonify
from tensorflow.keras.utils import load_img, img_to_array
from werkzeug.utils import secure_filename
from datetime import datetime
from huggingface_hub import hf_hub_download # Crucial for downloading model from HF Hub
import time # Used for potential retry logic, though not explicitly in hf_hub_download here

app = Flask(__name__)

# --- Model Loading Configuration ---
MODEL_FILE_NAME = "model.keras"
# IMPORTANT: Replace "YOUR_USERNAME/garbage-detection-model" with the actual REPO ID of YOUR MODEL on Hugging Face Hub.
# This is the repository where your 'model.keras' file is stored.
# Example: "your_huggingface_username/your_model_repo_name"
MODEL_REPO_ID = "nonamelife/garbage-detection-model" # <--- MAKE SURE THIS IS YOUR CORRECT MODEL REPO ID!

model = None # Initialize model as None

# --- Model Loading Logic ---
try:
    print(f"Attempting to download '{MODEL_FILE_NAME}' from Hugging Face Hub ({MODEL_REPO_ID})...")
    
    # hf_hub_download returns the FULL PATH to the downloaded file in the cache.
    # We specify local_dir within /app to ensure write permissions and consistency.
    # local_dir_use_symlinks=False is important for Docker environments to avoid symlink issues.
    downloaded_model_path = hf_hub_download(
        repo_id=MODEL_REPO_ID, 
        filename=MODEL_FILE_NAME,
        local_dir="/app/.cache/huggingface/models", # This directory will be created if it doesn't exist
        local_dir_use_symlinks=False 
    )
    
    print(f"'{MODEL_FILE_NAME}' downloaded successfully to: {downloaded_model_path}")

    # Now, load the model directly from the downloaded path.
    # This check is a safeguard, as hf_hub_download should ensure existence.
    if os.path.exists(downloaded_model_path): 
        model = tf.keras.models.load_model(downloaded_model_path)
        print("Model loaded successfully!")
    else:
        # This message indicates a very unusual state if download was reported successful.
        print(f"ERROR: Download reported success, but file not found at expected path: {downloaded_model_path}")

except Exception as e:
    # Catch any exceptions during download or loading and log them.
    print(f"FATAL: Could not download or load model from Hugging Face Hub: {e}")
    model = None # Ensure model remains None if there's an error

# --- End Model Loading Logic ---


# Configurations for Flask app
UPLOAD_FOLDER = os.path.join('static', 'uploads')
ALLOWED_EXTENSIONS = {'jpg', 'jpeg', 'png'}
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
# Ensure the uploads directory exists within the container
os.makedirs(UPLOAD_FOLDER, exist_ok=True) 

def allowed_file(filename):
    """Checks if the uploaded file has an allowed extension."""
    return '.' in filename and \
           filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

def preprocess_image(image_path):
    """Loads and preprocesses an image for model prediction."""
    img = load_img(image_path, target_size=(224, 224)) # Load image, resize to 224x224
    img_array = img_to_array(img) / 255.0 # Convert to array and normalize pixel values
    return np.expand_dims(img_array, axis=0) # Add batch dimension for model input

# --- Flask Routes ---
@app.route('/')
def index():
    """Renders the home page."""
    return render_template('home.html')

@app.route('/tool')
def tool():
    """Renders the image upload tool page."""
    return render_template('tool.html')

@app.route('/about')
def about():
    """Renders the about page."""
    return render_template('about.html')

@app.route('/contact')
def contact():
    """Renders the contact page."""
    return render_template('contact.html')

@app.route('/predict', methods=['POST'])
def predict():
    """Handles image uploads and returns predictions."""
    # Check if the model was loaded successfully at startup
    if model is None:
        return jsonify({'error': 'Model not loaded. Please check server logs.'}), 500

    # Check if a file was part of the request
    if 'file' not in request.files:
        return jsonify({'error': 'No files uploaded'}), 400

    files = request.files.getlist('file')
    if not files or all(f.filename == '' for f in files):
        return jsonify({'error': 'No files selected'}), 400

    results = []
    for file in files:
        file_path = None
        if file and allowed_file(file.filename):
            # Secure filename and create a unique name to prevent collisions
            filename = secure_filename(file.filename)
            timestamp = datetime.now().strftime("%Y%m%d%H%M%S%f")
            unique_filename = f"{timestamp}_{filename}"
            file_path = os.path.join(app.config['UPLOAD_FOLDER'], unique_filename)
            file.save(file_path) # Save the uploaded file temporarily

            try:
                img_array = preprocess_image(file_path) # Preprocess the image
                prediction = model.predict(img_array)[0][0] # Get prediction from the model
                
                # Determine label and confidence based on sigmoid output
                label = "Dirty" if prediction > 0.5 else "Clean"
                confidence = prediction if label == "Dirty" else 1 - prediction

                results.append({
                    'label': label,
                    'confidence': f"{confidence:.2%}", # Format confidence as percentage
                    'image_url': f"/static/uploads/{unique_filename}" # URL for displaying the image
                })
            except Exception as e:
                # Catch any errors during prediction or processing
                results.append({
                    'label': 'Error',
                    'confidence': 'N/A',
                    'image_url': None,
                    'error': str(e)
                })
        else:
            # Handle invalid file types
            results.append({
                'label': 'Error',
                'confidence': 'N/A',
                'image_url': None,
                'error': f"Invalid file type: {file.filename}"
            })

    # Render the results page with predictions
    return render_template('results.html', results=results)

# --- Main execution block ---
if __name__ == '__main__':
    # Hugging Face Spaces sets the PORT environment variable for the app to listen on.
    # We default to 7860 as it's common for HF Spaces apps.
    # Debug mode should be OFF for production deployments (like Hugging Face Spaces) for security.
    port = int(os.environ.get('PORT', 7860))
    app.run(host='0.0.0.0', port=port, debug=False)