File size: 1,400 Bytes
8260e47
bf29adc
8260e47
 
 
bf29adc
 
8260e47
bf29adc
 
8260e47
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
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