from flask import Flask, request, jsonify, send_file
from flask_cors import CORS
from transformers import pipeline
from PIL import Image
import io

app = Flask(__name__)
CORS(app, resources={r"/remove-background": {"origins": "*"}})

# Initialize the image segmentation pipeline
pipe = pipeline("image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True)

# Function to remove background from an image
def remove_background(image):
    pillow_image = pipe(image)
    return pillow_image

@app.route('/remove-background', methods=['POST'])
def remove_background_endpoint():
    try:
        # Handle raw image data in the request body
        if request.data:
            image = Image.open(io.BytesIO(request.data)).convert('RGB')
        # Fallback for multipart/form-data (if needed later)
        elif 'image' in request.files:
            file = request.files['image']
            image = Image.open(file).convert('RGB')
        else:
            return jsonify({'error': 'No image data provided'}), 400

        # Process the image
        result_image = remove_background(image)

        # Save the result to a bytes buffer
        img_io = io.BytesIO()
        result_image.save(img_io, format='PNG')
        img_io.seek(0)

        # Return the image as a response
        return send_file(img_io, mimetype='image/png', as_attachment=True, download_name='result.png')
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/')
def index():
    return jsonify({'status': 'Background Removal API is running'})

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860)