app.py

# app.py

import os
from flask import Flask, request, send_file, jsonify
from rembg import remove
from PIL import Image
import io

# --- Create the Flask App ---
app = Flask(__name__)

# --- Configuration ---
# Get the API Key from an environment variable for security.
API_KEY = os.environ.get("BG_REMOVER_API_KEY") 

# --- API Endpoints ---

# A simple root endpoint to check if the server is running.
@app.route('/')
def index():
    return "Background Remover API is running!"

# The main endpoint for removing the background.
@app.route('/remove-bg', methods=['POST'])
def remove_background_api():
    # 1. --- API Key Authentication ---
    api_key_header = request.headers.get('x-api-key')
    if not api_key_header or api_key_header != API_KEY:
        return jsonify({"error": "Unauthorized. Invalid or missing API Key."}), 401

    # 2. --- Image Validation ---
    if 'file' not in request.files:
        return jsonify({"error": "No file part in the request"}), 400

    file = request.files['file']

    if file.filename == '':
        return jsonify({"error": "No selected file"}), 400

    # 3. --- Image Processing ---
    if file:
        try:
            input_image_bytes = file.read()
            output_image_bytes = remove(input_image_bytes)
            
            # 4. --- Send the Response ---
            return send_file(
                io.BytesIO(output_image_bytes),
                mimetype='image/png',
                as_attachment=True,
                download_name='background_removed.png'
            )
        except Exception as e:
            return jsonify({"error": "Failed to process image", "details": str(e)}), 500

    return jsonify({"error": "An unknown error occurred"}), 500

# --- Run the App ---
if __name__ == '__main__':
    # For local testing, ensure the environment variable is set.
    if not API_KEY:
        # If you're using the temporary hard-coded key method, you can ignore this error.
        # Otherwise, this reminds you to set the variable.
        print("WARNING: BG_REMOVER_API_KEY is not set. The API will not be secured.")
        print("For local testing, run 'set BG_REMOVER_API_KEY=your-key' (Windows) or 'export BG_REMOVER_API_KEY=your-key' (Mac/Linux) first.")
    
    app.run(debug=True, port=5000)