Spaces:
Running
on
A10G
Running
on
A10G
File size: 1,795 Bytes
251e479 |
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 |
# Copyright (c) OpenMMLab. All rights reserved.
import numpy as np
import torch
from ..utils import ext_loader
ext_module = ext_loader.load_ext('_ext', ['contour_expand'])
def contour_expand(kernel_mask, internal_kernel_label, min_kernel_area,
kernel_num):
"""Expand kernel contours so that foreground pixels are assigned into
instances.
Arguments:
kernel_mask (np.array or Tensor): The instance kernel mask with
size hxw.
internal_kernel_label (np.array or Tensor): The instance internal
kernel label with size hxw.
min_kernel_area (int): The minimum kernel area.
kernel_num (int): The instance kernel number.
Returns:
label (list): The instance index map with size hxw.
"""
assert isinstance(kernel_mask, (torch.Tensor, np.ndarray))
assert isinstance(internal_kernel_label, (torch.Tensor, np.ndarray))
assert isinstance(min_kernel_area, int)
assert isinstance(kernel_num, int)
if isinstance(kernel_mask, np.ndarray):
kernel_mask = torch.from_numpy(kernel_mask)
if isinstance(internal_kernel_label, np.ndarray):
internal_kernel_label = torch.from_numpy(internal_kernel_label)
if torch.__version__ == 'parrots':
if kernel_mask.shape[0] == 0 or internal_kernel_label.shape[0] == 0:
label = []
else:
label = ext_module.contour_expand(
kernel_mask,
internal_kernel_label,
min_kernel_area=min_kernel_area,
kernel_num=kernel_num)
label = label.tolist()
else:
label = ext_module.contour_expand(kernel_mask, internal_kernel_label,
min_kernel_area, kernel_num)
return label
|