Spaces:
Running
Running
File size: 1,625 Bytes
0271b14 7941a82 901bf30 0271b14 93700b4 0271b14 7941a82 901bf30 0271b14 12e9f72 0271b14 12e9f72 0271b14 12e9f72 0271b14 93700b4 0271b14 |
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 |
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) |