Create model.py

model.py

@@ -0,0 +1,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()