File size: 3,777 Bytes
d7e58f0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
77
78
79
# Copyright (c) OpenMMLab. All rights reserved.
import numpy as np
import torch


def test_pixel_group():
    from mmcv.ops import pixel_group
    np_score = np.array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                         [0, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0],
                         [0, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0],
                         [0, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0],
                         [0, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0],
                         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]).astype(np.float32)
    np_mask = (np_score > 0.5)
    np_embedding = np.zeros((10, 10, 8)).astype(np.float32)
    np_embedding[:, :7] = 0.9
    np_embedding[:, 7:] = 10.0
    np_kernel_label = np.array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                                [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                                [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                                [0, 0, 1, 1, 1, 0, 0, 0, 2, 0],
                                [0, 0, 1, 1, 1, 0, 0, 0, 2, 0],
                                [0, 0, 1, 1, 1, 0, 0, 0, 2, 0],
                                [0, 0, 1, 1, 1, 0, 0, 0, 2, 0],
                                [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                                [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                                [0, 0, 0, 0, 0, 0, 0, 0, 0,
                                 0]]).astype(np.int32)
    np_kernel_contour = np.array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                                  [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                                  [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                                  [0, 0, 1, 1, 1, 0, 0, 0, 1, 0],
                                  [0, 0, 1, 0, 1, 0, 0, 0, 1, 0],
                                  [0, 0, 1, 0, 1, 0, 0, 0, 1, 0],
                                  [0, 0, 1, 1, 1, 0, 0, 0, 1, 0],
                                  [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                                  [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                                  [0, 0, 0, 0, 0, 0, 0, 0, 0,
                                   0]]).astype(np.uint8)
    kernel_region_num = 3
    distance_threshold = float(0.8)
    result = pixel_group(np_score, np_mask, np_embedding, np_kernel_label,
                         np_kernel_contour, kernel_region_num,
                         distance_threshold)
    gt_1 = [
        0.8999997973442078, 24.0, 1.0, 3.0, 2.0, 3.0, 3.0, 3.0, 4.0, 3.0, 5.0,
        3.0, 6.0, 3.0, 1.0, 4.0, 2.0, 4.0, 3.0, 4.0, 4.0, 4.0, 5.0, 4.0, 6.0,
        4.0, 1.0, 5.0, 2.0, 5.0, 3.0, 5.0, 4.0, 5.0, 5.0, 5.0, 6.0, 5.0, 1.0,
        6.0, 2.0, 6.0, 3.0, 6.0, 4.0, 6.0, 5.0, 6.0, 6.0, 6.0
    ]

    gt_2 = [
        0.9000000357627869, 8.0, 7.0, 3.0, 8.0, 3.0, 7.0, 4.0, 8.0, 4.0, 7.0,
        5.0, 8.0, 5.0, 7.0, 6.0, 8.0, 6.0
    ]

    assert np.allclose(result[0], [0, 0])
    assert np.allclose(result[1], gt_1)
    assert np.allclose(result[2], gt_2)

    # test torch Tensor
    np_score_t = torch.from_numpy(np_score)
    np_mask_t = torch.from_numpy(np_mask)
    np_embedding_t = torch.from_numpy(np_embedding)
    np_kernel_label_t = torch.from_numpy(np_kernel_label)
    np_kernel_contour_t = torch.from_numpy(np_kernel_contour)

    result = pixel_group(np_score_t, np_mask_t, np_embedding_t,
                         np_kernel_label_t, np_kernel_contour_t,
                         kernel_region_num, distance_threshold)

    assert np.allclose(result[0], [0, 0])
    assert np.allclose(result[1], gt_1)
    assert np.allclose(result[2], gt_2)