use PNG for mask

main.py

@@ -10,7 +10,7 @@ import models.face_classifier as classifier
 from fastapi.middleware.cors import CORSMiddleware
 from PIL import Image
 from rembg import remove
-from utils.helpers import image_to_base64, calculate_mask_area
+from utils.helpers import image_to_base64, calculate_mask_area, process_image
 
 
 dotenv.load_dotenv()
@@ -62,7 +62,8 @@ async def predict_image(file: UploadFile = File(...)):
         # change the mask to base64 and calculate the score
         for i in range(len(results)):
             mask_area = calculate_mask_area(results[i]["mask"])
-            results[i]["mask"] = image_to_base64(results[i]["mask"])
+            processed_image = process_image(results[i]["mask"])
+            results[i]["mask"] = image_to_base64(processed_image, "PNG")
             if results[i]["label"] == "background":
                 continue
             print(f"{results[i]['label']} area: {mask_area}")

utils/helpers.py

@@ -3,13 +3,29 @@ import io
 import base64
 import numpy as np
 
-def image_to_base64(image: Image.Image) -> str:
+def image_to_base64(image: Image.Image,format="JPEG") -> str:
     buffered = io.BytesIO()
 
-    image.save(buffered, format="JPEG")
+    image.save(buffered, format=format)
     return base64.b64encode(buffered.getvalue()).decode("utf-8")
 
 def calculate_mask_area(mask: Image.Image, background=False) -> int:
     mask_array = np.array(mask)
     non_zero_pixels = np.count_nonzero(mask_array)
-    return non_zero_pixels
+    return non_zero_pixels
+
+def process_image(input_image: Image.Image) -> Image.Image:
+    if input_image.mode != 'RGBA':
+        input_image = input_image.convert('RGBA')
+
+    data = np.array(input_image)
+    # Split the image into its component channels
+    r, g, b, a = data.T
+
+    # Create a mask where all pixels that are black (0,0,0) will have 0 alpha
+    black_areas = (r == 0) & (g == 0) & (b == 0)
+
+    # Apply the mask to the alpha channel
+    data[..., 3][black_areas] = 0
+
+    return Image.fromarray(data)