| |
| 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 |
|
|
