from flask import Flask, request, send_file, render_template_string
import numpy as np
import torch
import torch.nn.functional as F
from torchvision.transforms.functional import normalize
from briarmbg import BriaRMBG
import io
from PIL import Image
# --- Model Loading and Processing Functions ---
# يتم تحميل النموذج مرة واحدة عند بدء تشغيل التطبيق
net = BriaRMBG.from_pretrained("briaai/RMBG-1.4")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
net.to(device)
net.eval()
def resize_image(image):
image = image.convert('RGB')
model_input_size = (1024, 1024)
image = image.resize(model_input_size, Image.BILINEAR)
return image
def process(image_np):
# prepare input
orig_image = Image.fromarray(image_np)
w, h = orig_im_size = orig_image.size
image = resize_image(orig_image)
im_np = np.array(image)
im_tensor = torch.tensor(im_np, dtype=torch.float32).permute(2, 0, 1)
im_tensor = torch.unsqueeze(im_tensor, 0)
im_tensor = torch.divide(im_tensor, 255.0)
im_tensor = normalize(im_tensor, [0.5, 0.5, 0.5], [1.0, 1.0, 1.0])
if torch.cuda.is_available():
im_tensor = im_tensor.cuda()
# inference
result = net(im_tensor)
# post process
result = torch.squeeze(F.interpolate(result[0][0], size=(h, w), mode='bilinear'), 0)
ma = torch.max(result)
mi = torch.min(result)
result = (result - mi) / (ma - mi)
# image to pil
result_array = (result * 255).cpu().data.numpy().astype(np.uint8)
pil_mask = Image.fromarray(np.squeeze(result_array))
# add the mask on the original image as alpha channel
new_im = orig_image.copy()
new_im.putalpha(pil_mask)
return new_im
# --- Flask App Setup ---
app = Flask(__name__)
@app.route('/')
def index():
return render_template_string('''
Background Remover
Background Remover API
Original Image
Processed Image
''')
# API Route for background removal
@app.route('/remove_bg', methods=['POST'])
def remove_background():
try:
if 'file' not in request.files:
return "No file part in the request", 400
file = request.files['file']
if file.filename == '':
return "No selected file", 400
# Read the file and convert to a NumPy array
image_bytes = file.read()
pil_image = Image.open(io.BytesIO(image_bytes))
image_np = np.array(pil_image)
# استدعاء دالة المعالجة
processed_image = process(image_np)
# تحويل الصورة الناتجة إلى bytes وإرسالها
# ملاحظة: يجب استخدام PNG لدعم الشفافية
img_io = io.BytesIO()
processed_image.save(img_io, format='PNG')
img_io.seek(0)
return send_file(img_io, mimetype='image/png')
except Exception as e:
import traceback
traceback.print_exc()
return f"Error processing image: {str(e)}", 500
if __name__ == '__main__':
app.run(host='0.0.0.0', port=7860)