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()