src/plot_utils.py

from typing import Optional

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.pyplot import Figure


def render_masks(
    image,
    masks,
    random_color: Optional[bool] = True,
    smoothen_contours: Optional[bool] = True,
) -> "Figure":
    h, w = image.shape[:2]
    fig, ax = plt.subplots(figsize=(w / 100, h / 100), dpi=100)
    ax.axis("off")
    ax.imshow(image)

    for mask in masks:
        if random_color:
            color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
        else:
            color = np.array([30 / 255, 144 / 255, 255 / 255, 0.6])

        mask = mask.astype(np.uint8)
        mask = mask.reshape(h, w)
        mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)

        if smoothen_contours:
            import cv2

            contours, _ = cv2.findContours(
                mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE
            )
            contours = [
                cv2.approxPolyDP(contour, epsilon=0.01, closed=True)
                for contour in contours
            ]
            mask_image = cv2.drawContours(
                mask_image, contours, -1, (1, 1, 1, 0.5), thickness=2
            )

        ax.imshow(mask_image, alpha=0.6)

    # Make image occupy the whole figure
    ax.set_xlim(0, w)
    ax.set_ylim(h, 0)
    plt.subplots_adjust(left=0, right=1, top=1, bottom=0)

    return fig