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