File size: 5,754 Bytes
be1de07 |
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 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 |
#!/usr/bin/env python3
"""
Web application for AI Background Removal
Serves a responsive frontend and processes images using the background_remover module
"""
import os
import uuid
from pathlib import Path
from flask import Flask, render_template, request, jsonify, send_file, url_for
from werkzeug.utils import secure_filename
from werkzeug.middleware.proxy_fix import ProxyFix
from background_remover import BackgroundRemover
app = Flask(__name__)
app.config['MAX_CONTENT_LENGTH'] = 10 * 1024 * 1024 # 10MB max file size
app.secret_key = os.environ.get("SESSION_SECRET")
# Add ProxyFix for Replit's proxy system
app.wsgi_app = ProxyFix(app.wsgi_app, x_proto=1, x_host=1)
# Create necessary directories
UPLOAD_FOLDER = Path('uploads')
OUTPUT_FOLDER = Path('outputs')
UPLOAD_FOLDER.mkdir(exist_ok=True)
OUTPUT_FOLDER.mkdir(exist_ok=True)
# Initialize background remover
bg_remover = BackgroundRemover()
# Allowed file extensions
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif', 'bmp', 'tiff', 'webp'}
def allowed_file(filename):
"""Check if the file extension is allowed."""
return '.' in filename and \
filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
@app.route('/')
def index():
"""Serve the main page."""
return render_template('index.html')
@app.route('/remove-background', methods=['POST'])
def remove_background():
"""Handle image upload and background removal."""
try:
# Check if file was uploaded
if 'image' not in request.files:
return jsonify({'error': 'No image file provided'}), 400
file = request.files['image']
# Check if file was selected
if file.filename == '':
return jsonify({'error': 'No file selected'}), 400
# Validate file type
if not file or not allowed_file(file.filename):
return jsonify({'error': 'Invalid file type. Please use JPG, PNG, BMP, TIFF, or WebP.'}), 400
# Generate unique filename
if not file.filename or '.' not in file.filename:
return jsonify({'error': 'Invalid filename'}), 400
file_extension = file.filename.rsplit('.', 1)[1].lower()
unique_id = str(uuid.uuid4())
input_filename = f"{unique_id}.{file_extension}"
output_filename = f"{unique_id}_no_bg.png"
# Save uploaded file
input_path = UPLOAD_FOLDER / input_filename
output_path = OUTPUT_FOLDER / output_filename
file.save(str(input_path))
# Process the image
try:
result_path = bg_remover.remove_background(str(input_path), str(output_path))
# Clean up input file
input_path.unlink()
# Return success response
return jsonify({
'success': True,
'output_url': url_for('get_output', filename=output_filename),
'message': 'Background removed successfully!'
})
except Exception as e:
# Clean up files on error
if input_path.exists():
input_path.unlink()
if output_path.exists():
output_path.unlink()
return jsonify({'error': f'Processing failed: {str(e)}'}), 500
except Exception as e:
return jsonify({'error': f'Upload failed: {str(e)}'}), 500
@app.route('/output/<filename>')
def get_output(filename):
"""Serve processed output images."""
try:
file_path = OUTPUT_FOLDER / secure_filename(filename)
if not file_path.exists():
return jsonify({'error': 'File not found'}), 404
return send_file(str(file_path), as_attachment=False)
except Exception as e:
return jsonify({'error': f'Failed to serve file: {str(e)}'}), 500
@app.route('/health')
def health_check():
"""Health check endpoint."""
return jsonify({'status': 'healthy', 'service': 'AI Background Remover'})
@app.errorhandler(413)
def too_large(e):
"""Handle file too large error."""
return jsonify({'error': 'File too large. Maximum size is 10MB.'}), 413
@app.errorhandler(404)
def not_found(e):
"""Handle 404 errors."""
return jsonify({'error': 'Endpoint not found'}), 404
@app.errorhandler(500)
def server_error(e):
"""Handle server errors."""
return jsonify({'error': 'Internal server error'}), 500
# Cleanup old files periodically (basic implementation)
def cleanup_old_files():
"""Remove files older than 1 hour to save space."""
import time
current_time = time.time()
for folder in [UPLOAD_FOLDER, OUTPUT_FOLDER]:
for file_path in folder.glob('*'):
if file_path.is_file():
# If file is older than 1 hour, delete it
if current_time - file_path.stat().st_mtime > 3600:
try:
file_path.unlink()
except Exception:
pass # Ignore cleanup errors
if __name__ == '__main__':
print("Starting AI Background Remover Web App...")
print("Frontend: Responsive design for mobile and desktop")
print("Backend: AI-powered background removal")
# Determine port based on environment
# Hugging Face Spaces require port 7860, Replit uses 5000
port = int(os.environ.get("PORT", 7860)) # Default to 7860 for Hugging Face
if os.environ.get("REPL_ID"): # Replit environment detected
port = 5000
print(f"Access your app at: http://localhost:{port}")
# Run cleanup on startup
cleanup_old_files()
# Start the Flask app
app.run(host='0.0.0.0', port=port, debug=False) |