Spaces:
Sleeping
Sleeping
File size: 5,310 Bytes
21c551e 7463015 21c551e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 |
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 # New import!
import time
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
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}")
else:
print(f"Model file '{MODEL_FILE_NAME}' still not found after download attempt.")
except Exception as e:
print(f"Error loading model from {MODEL_FILE_NAME}: {e}")
model = None # Ensure model is None if loading fails
# Configurations
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
def allowed_file(filename):
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)
@app.route('/')
def index():
return render_template('home.html')
@app.route('/tool')
def tool():
return render_template('tool.html')
@app.route('/about')
def about():
return render_template('about.html')
@app.route('/contact')
def contact():
return render_template('contact.html')
@app.route('/predict', methods=['POST'])
def predict():
if model is None:
return jsonify({'error': 'Model not loaded. Please check server logs.'}), 500
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):
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)
try:
img_array = preprocess_image(file_path)
prediction = model.predict(img_array)[0][0]
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}"
})
except Exception as e:
results.append({
'label': 'Error',
'confidence': 'N/A',
'image_url': None,
'error': str(e)
})
finally:
# Clean up the uploaded file after processing
if file_path and os.path.exists(file_path):
try:
os.remove(file_path)
print(f"Deleted uploaded file: {file_path}")
except Exception as e:
print(f"Error deleting file {file_path}: {e}")
else:
results.append({
'label': 'Error',
'confidence': 'N/A',
'image_url': None,
'error': f"Invalid file type: {file.filename}"
})
return render_template('results.html', results=results)
if __name__ == '__main__':
# Hugging Face Spaces sets the PORT environment variable
# Default to 7860 as it's common for HF Spaces apps
port = int(os.environ.get('PORT', 7860))
app.run(host='0.0.0.0', port=port, debug=True) |