mmocr/utils/mask_utils.py

# Copyright (c) OpenMMLab. All rights reserved.
import cv2
import numpy as np

from mmocr.utils.typing_utils import ArrayLike


def fill_hole(input_mask: ArrayLike) -> np.array:
    """Fill holes in matrix.

        Input:
            [[0, 0, 0, 0, 0, 0, 0],
             [0, 1, 1, 1, 1, 1, 0],
             [0, 1, 0, 0, 0, 1, 0],
             [0, 1, 1, 1, 1, 1, 0],
             [0, 0, 0, 0, 0, 0, 0]]
        Output:
            [[0, 0, 0, 0, 0, 0, 0],
             [0, 1, 1, 1, 1, 1, 0],
             [0, 1, 1, 1, 1, 1, 0],
             [0, 1, 1, 1, 1, 1, 0],
             [0, 0, 0, 0, 0, 0, 0]]

    Args:
        input_mask (ArrayLike): The input mask.

    Returns:
        np.array: The output mask that has been filled.
    """
    input_mask = np.array(input_mask)
    h, w = input_mask.shape
    canvas = np.zeros((h + 2, w + 2), np.uint8)
    canvas[1:h + 1, 1:w + 1] = input_mask.copy()

    mask = np.zeros((h + 4, w + 4), np.uint8)

    cv2.floodFill(canvas, mask, (0, 0), 1)
    canvas = canvas[1:h + 1, 1:w + 1].astype(np.bool_)

    return ~canvas | input_mask