File size: 1,028 Bytes
d3752e4 |
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 |
import base64
from io import BytesIO
from PIL import Image
import time
# Import your utility functions
from src.utils import change_background, matte
class BackgroundRemover:
def __init__(self):
self.hexmap = {
"Transparent (PNG)": "#000000",
"Black": "#000000",
"White": "#FFFFFF",
"Green": "#22EE22",
"Red": "#EE2222",
"Blue": "#2222EE",
}
def process_image(self, image_base64, background_color):
img_input = Image.open(BytesIO(base64.b64decode(image_base64)))
alpha = 0.0 if background_color == "Transparent (PNG)" else 1.0
img_matte = matte(img_input)
img_output = change_background(img_input, img_matte, background_alpha=alpha, background_hex=self.hexmap[background_color])
buffered = BytesIO()
img_output.save(buffered, format="PNG")
img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
return img_str
model = BackgroundRemover() |