Update app.py

app.py

@@ -5,80 +5,100 @@ 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 # New import!
-import time
+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"
-# REPLACE THIS WITH YOUR HUGGING FACE MODEL REPO ID
-# Format: "your-username/your-model-repo-name"
-HF_MODEL_REPO_ID = "nonamelife/garbage-detection-model" # Example!
-
-# Check if model exists, if not, try to download it from Hugging Face Hub
-if not os.path.exists(MODEL_FILE_NAME):
-    print(f"'{MODEL_FILE_NAME}' not found locally. Attempting to download from Hugging Face Hub...")
-    try:
-        # Download the model from Hugging Face Hub
-        # The downloaded file will be in a cache directory by default,
-        # so we'll move it to the current directory for easier loading.
-        model_path = hf_hub_download(repo_id=HF_MODEL_REPO_ID, filename=MODEL_FILE_NAME)
-        # Move the downloaded file to the root directory for app.py to find it easily
-        os.rename(model_path, MODEL_FILE_NAME)
-        print(f"'{MODEL_FILE_NAME}' downloaded successfully from Hugging Face Hub.")
-    except Exception as e:
-        print(f"FATAL: Could not download model from Hugging Face Hub: {e}")
-        # If download fails, the model will remain None, and prediction attempts will fail.
-        model = None
-
-# Load the trained model
-model = None # Initialize model to None
+# 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:
-    if os.path.exists(MODEL_FILE_NAME):
-        model = tf.keras.models.load_model(MODEL_FILE_NAME)
-        print(f"Model loaded successfully from {MODEL_FILE_NAME}")
+    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:
-        print(f"Model file '{MODEL_FILE_NAME}' still not found after download attempt.")
+        # 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:
-    print(f"Error loading model from {MODEL_FILE_NAME}: {e}")
-    model = None # Ensure model is None if loading fails
+    # 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
+# Configurations for Flask app
 UPLOAD_FOLDER = os.path.join('static', 'uploads')
 ALLOWED_EXTENSIONS = {'jpg', 'jpeg', 'png'}
 app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
-os.makedirs(UPLOAD_FOLDER, exist_ok=True) # Ensure uploads directory exists
+# Ensure the uploads directory exists within the container
+os.makedirs(UPLOAD_FOLDER, exist_ok=True) 
 
 def allowed_file(filename):
-    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
+    """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):
-    img = load_img(image_path, target_size=(224, 224))
-    img_array = img_to_array(img) / 255.0
-    return np.expand_dims(img_array, axis=0)
+    """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
 
@@ -90,24 +110,28 @@ def predict():
     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)
+            file.save(file_path) # Save the uploaded file temporarily
 
             try:
-                img_array = preprocess_image(file_path)
-                prediction = model.predict(img_array)[0][0]
+                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%}",
-                    'image_url': f"/static/uploads/{unique_filename}"
+                    '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',
@@ -115,7 +139,7 @@ def predict():
                     'error': str(e)
                 })
             finally:
-                # Clean up the uploaded file after processing
+                # Clean up: delete the temporary uploaded file
                 if file_path and os.path.exists(file_path):
                     try:
                         os.remove(file_path)
@@ -123,6 +147,7 @@ def predict():
                     except Exception as e:
                         print(f"Error deleting file {file_path}: {e}")
         else:
+            # Handle invalid file types
             results.append({
                 'label': 'Error',
                 'confidence': 'N/A',
@@ -130,10 +155,13 @@ def predict():
                 '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
-    # Default to 7860 as it's common for HF Spaces apps
+    # 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=True)
+    app.run(host='0.0.0.0', port=port, debug=False)