haydpw's picture
change mask color according to their labels
9d7fc49
raw
history blame
1.88 kB
from PIL import Image, ImageChops, ImageEnhance
import io
import base64
import numpy as np
def image_to_base64(image: Image.Image,format="JPEG") -> str:
buffered = io.BytesIO()
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
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) 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] = 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
# Apply the mask to the alpha channel
rgba_data[..., 3][black_areas] = 0
return Image.fromarray(rgba_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
color_mask = Image.new("RGBA", original_image.size, (255, 0, 0, 128)) # Red color with transparency
combined.paste(color_mask, (0, 0), mask_image)
return combined