update code

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, process_image
+from utils.helpers import combine_images, image_to_base64, calculate_mask_area, process_image
 
 
 dotenv.load_dotenv()
@@ -74,9 +74,14 @@ async def predict_image(file: UploadFile = File(...)):
         # add original image base 64 as original image:
         image_bg_removed = image_bg_removed.convert("RGB")
         
+        # Combine the original image and masks
+        combined_image = combine_images(pil_image, results)
+        combined_image_base64 = image_to_base64(combined_image, "PNG")
+        
         response = {
             "original_image": image_to_base64(image_bg_removed),
-            "segmentation_results":results
+            "segmentation_results":results,
+            "combined_image": combined_image_base64
         }
         
         # Kembalikan hasil klasifikasi

utils/helpers.py

@@ -1,4 +1,4 @@
-from PIL import Image
+from PIL import Image, ImageChops, ImageEnhance
 import io
 import base64
 import numpy as np
@@ -28,4 +28,15 @@ def process_image(input_image: Image.Image) -> Image.Image:
     # Apply the mask to the alpha channel
     data[..., 3][black_areas] = 0
 
-    return Image.fromarray(data)
+    return Image.fromarray(data)
+
+def combine_images(original_image: Image.Image, masks: list) -> Image.Image:
+    combined = original_image.copy()
+    for mask in masks:
+        if mask['label'] == 'background':
+            continue
+        mask_image = Image.open(io.BytesIO(base64.b64decode(mask['mask'])))
+        mask_image = mask_image.convert("L")  # Convert mask to grayscale
+        mask_image = ImageEnhance.Brightness(mask_image).enhance(0.5)  # Adjust the brightness to make it more visible
+        combined = ImageChops.add(combined, mask_image)
+    return combined