change mask color according to their labels

main.py

@@ -64,11 +64,17 @@ async def predict_image(file: UploadFile = File(...)):
         else:
             background_area = 0
         
+        mask_dict = {
+            'acne': (255, 0, 0),
+            'dry': (0, 255, 0),
+            'oily': (0,0,255),
+            'background': (255,255,255)
+        }
         # change the mask to base64 and calculate the score
         for i in range(len(results)):
             mask_area = calculate_mask_area(results[i]["mask"])
             print(results[i]["mask"].size)
-            processed_image = process_image(results[i]["mask"])
+            processed_image = process_image(results[i]["mask"], mask_dict[results[i]['label']])
             results[i]["mask"] = image_to_base64(processed_image, "PNG")
             if results[i]["label"] == "background":
                 continue

utils/helpers.py

@@ -14,22 +14,22 @@ def calculate_mask_area(mask: Image.Image, background=False) -> int:
     non_zero_pixels = np.count_nonzero(mask_array)
     return non_zero_pixels
 
-def process_image(input_image: Image.Image) -> Image.Image:
+def process_image(input_image: Image.Image, fill_color:tuple = (255, 255, 255)) -> Image.Image:
 
 
     data = np.array(input_image)
     # Split the image into its component channels
 
-    # Create a mask where all pixels that are black (0,0,0) will have 0 alpha
+    # Create a mask where all pixels that are black (0) will have 0 alpha
     black_areas = data == 0
     
     rgba_image = Image.new('RGBA', input_image.size)
     rgba_data = np.array(rgba_image)
 
     # Copy the grayscale data to all RGB channels
-    rgba_data[..., 0] = data
-    rgba_data[..., 1] = data
-    rgba_data[..., 2] = data
+    rgba_data[..., 0] = fill_color[0]
+    rgba_data[..., 1] = fill_color[1]
+    rgba_data[..., 2] = fill_color[2]
 
     # Set alpha channel to 255 (fully opaque) for all pixels
     rgba_data[..., 3] = 255