File size: 813 Bytes
0b4516f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
# Copyright (c) OpenMMLab. All rights reserved.
import unittest

import numpy as np
import torch

from mmocr.utils import fill_hole


class TestFillHole(unittest.TestCase):

    def setUp(self) -> None:
        self.input_mask_list = [[0, 1, 1, 1, 0], [0, 1, 0, 1, 0],
                                [0, 1, 1, 1, 0]]
        self.input_mask_array = np.array(self.input_mask_list)
        self.input_mask_tensor = torch.tensor(self.input_mask_list)
        self.gt = np.array([[0, 1, 1, 1, 0], [0, 1, 1, 1, 0], [0, 1, 1, 1, 0]])

    def test_fill_hole(self):
        self.assertTrue(np.allclose(fill_hole(self.input_mask_list), self.gt))
        self.assertTrue(np.allclose(fill_hole(self.input_mask_array), self.gt))
        self.assertTrue(
            np.allclose(fill_hole(self.input_mask_tensor), self.gt))