Spaces:
Running
on
L40S
Running
on
L40S
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)
|