src/plot_utils.py

from typing import Optional

import numpy as np
from PIL import Image


def export_mask(
    masks: np.ndarray,
    random_color: Optional[bool] = True,
    smoothen_contours: Optional[bool] = True,
) -> Image:
    num_masks, _, h, w = masks.shape
    num_masks = len(masks)

    # Ensure masks are 2D by squeezing channel dimension
    masks = masks.squeeze(axis=1)

    # Create a single uint8 image with unique values for each mask
    combined_mask = np.zeros((h, w), dtype=np.uint8)

    for i in range(num_masks):
        mask = masks[i]
        mask = mask.astype(np.uint8)
        combined_mask[mask > 0] = i + 1

    # Create color map for visualization
    if random_color:
        colors = np.random.rand(num_masks, 3)  # Random colors for each mask
    else:
        colors = np.array(
            [[30 / 255, 144 / 255, 255 / 255]] * num_masks
        )  # Use fixed color

    # Create an RGB image where each mask has its own color
    color_image = np.zeros((h, w, 3), dtype=np.uint8)

    for i in range(1, num_masks + 1):
        mask_color = colors[i - 1] * 255
        color_image[combined_mask == i] = mask_color

    # Convert the NumPy array to a PIL Image
    pil_image = Image.fromarray(color_image)

    # Optional: Add contours to the mask image
    if smoothen_contours:
        import cv2

        contours_image = np.zeros((h, w, 4), dtype=np.float32)

        for i in range(1, num_masks + 1):
            mask = (combined_mask == i).astype(np.uint8)
            contours, _ = cv2.findContours(
                mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
            )
            contours = [
                cv2.approxPolyDP(contour, epsilon=0.01, closed=True)
                for contour in contours
            ]
            contours_image = cv2.drawContours(
                contours_image, contours, -1, (0, 0, 0, 0.5), thickness=2
            )

        # Convert contours to PIL image and blend with the color image
        contours_image = (contours_image[:, :, :3] * 255).astype(np.uint8)
        contours_pil_image = Image.fromarray(contours_image)
        pil_image = Image.blend(pil_image, contours_pil_image, alpha=0.6)

    return pil_image