File size: 2,006 Bytes
6c2a15f
 
 
 
 
 
 
 
 
 
 
4dc2bc3
6c2a15f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4dc2bc3
6c2a15f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
""" 
Based on https://pypi.org/project/pyEdgeEval/0.2.6
Cite: http://arxiv.org/abs/2304.09427
"""

from typing import List

import cv2
import numpy as np


def mask2bdry(mask: np.ndarray, ignore_mask: np.ndarray, thickness: int, quality: int = 5) -> np.ndarray:
    """
    Convert binary mask to boundaries.

    Args:
        mask (np.ndarray): 2D binary mask
        ignore_mask (np.ndarray): 2D binary mask
        thickness (int): boundary thickness
        quality (int): distance transform quality

    Returns:
        bdry (np.ndarray): 2D binary boundary mask
    """
    inner = cv2.distanceTransform(((mask + ignore_mask) > 0).astype(np.uint8), cv2.DIST_L2, quality)
    outer = cv2.distanceTransform(((1.0 - mask) > 0).astype(np.uint8), cv2.DIST_L2, quality)
    dist = outer + inner

    dist[dist > thickness] = 0
    bdry = (dist > 0).astype(np.uint8)
    return bdry


def mask2sbd(mask: np.ndarray, ignore_indices: List, thickness: int = 4, quality: int = 5) -> np.ndarray:
    """
    Convert Segmentation Mask to Semantic Boundaries.

    Args:
        mask (np.ndarray): segmentation mask
        ignore_indicies (List[int]): list of indices to ignore
        thickness (int): boundary thickness
        quality (int): distance transform quality

    Returns:
        bdrys (np.ndarray): 3D array containing boundaries
    """
    assert mask.ndim == 3
    num_labels, h, w = mask.shape

    # make ignore mask
    ignore_mask = np.zeros((h, w), dtype=np.uint8)
    for i in ignore_indices:
        ignore_mask += mask[i]

    bdrys = np.zeros_like(mask)
    for label in range(num_labels):
        m = mask[label]

        if label in ignore_indices:
            continue

        # if there are no class labels in the mask
        if not np.count_nonzero(m):
            continue

        edge = mask2bdry(
            mask=m,
            ignore_mask=ignore_mask,
            thickness=thickness,
            quality=quality,
        )

        bdrys[label] = edge

    return bdrys