app.py

#!/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)