43 lines
1.7 KiB
Python
43 lines
1.7 KiB
Python
from flask import Flask, render_template, request, send_file, flash, redirect, url_for
|
|
from rembg import remove
|
|
from PIL import Image
|
|
from io import BytesIO
|
|
|
|
app = Flask(__name__)
|
|
|
|
# Update allowed extensions to include additional file types
|
|
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif', 'bmp', 'tiff', 'webp'}
|
|
# Set max content length to 1GB (expressed in bytes)
|
|
app.config['MAX_CONTENT_LENGTH'] = 1 * 1024 * 1024 * 1024 # 1 GB limit
|
|
|
|
def file_allowed(filename):
|
|
"""
|
|
Check if the file has an allowed extension.
|
|
"""
|
|
return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
|
|
|
|
@app.route('/', methods=['GET', 'POST'])
|
|
def upload_file():
|
|
"""
|
|
Handle file upload via POST, process image to remove background, and return the processed image.
|
|
"""
|
|
if request.method == 'POST':
|
|
file = request.files.get('file', None)
|
|
if file and file_allowed(file.filename):
|
|
try:
|
|
input_image = Image.open(file.stream)
|
|
output_image = remove(input_image, post_process_mask=True)
|
|
img_io = BytesIO()
|
|
output_image.save(img_io, 'PNG')
|
|
img_io.seek(0)
|
|
return send_file(img_io, mimetype='image/png', as_attachment=True, download_name='processed_image.png')
|
|
except Exception as e:
|
|
flash('Failed to process the image. Error: {}'.format(e), 'error')
|
|
return redirect(url_for('upload_file'))
|
|
else:
|
|
flash('Please upload a valid image file (PNG, JPG, JPEG, GIF, BMP, TIFF, WEBP).', 'error')
|
|
return redirect(url_for('upload_file'))
|
|
return render_template('index.html')
|
|
|
|
if __name__ == '__main__':
|
|
app.run(host='0.0.0.0', debug=True, port=5124) |