Update handler.py

handler.py

@@ -7,12 +7,11 @@ from typing import Dict, List, Any
 from io import BytesIO
 import base64
 
-class EndpointHandler():
-    def init(self, path=""):
+class EndpointHandler:
+    def __init__(self, path=""):
         # Initialize the image segmentation pipeline
         self.pipeline = AutoModelForImageSegmentation.from_pretrained('briaai/RMBG-2.0', trust_remote_code=True)
         torch.set_float32_matmul_precision(['high', 'highest'][0])
-    
         self.image_size = (1024, 1024)
     
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
@@ -27,7 +26,7 @@ class EndpointHandler():
         # Extract the image path from the input data
         image_b64 = data.get("inputs", "")
         image_data = base64.b64decode(image_b64)
-        image = Image.open(BytesIO(image_data)).convert("RGB")
+        image = Image.open(BytesIO(image_data)).convert("RGB")  # Convert to RGB instead of RGBA
         input_images = transform_image(image).unsqueeze(0)
         print(2)
         # Prediction
@@ -35,13 +34,16 @@ class EndpointHandler():
             preds = self.pipeline(input_images)[-1].sigmoid()
         pred = preds[0].squeeze()
         pred_pil = transforms.ToPILImage()(pred)
-        mask = pred_pil.resize(image.size)
-        image.putalpha(mask)
+        
+        # Resize the mask to the original image size
+        original_image = Image.open(BytesIO(image_data)).convert("RGBA")  # Load original RGBA image for alpha
+        mask = pred_pil.resize(original_image.size)
+        original_image.putalpha(mask)
         print(3)
     
         # Convert the image with alpha mask to base64
         buffered = BytesIO()
-        image.save(buffered, format="PNG")
+        original_image.save(buffered, format="PNG")
         base64_image = base64.b64encode(buffered.getvalue()).decode("utf-8")
         
         # Return the result as a list of dictionaries