Spaces:
Build error
Build error
| # Copyright (c) Meta Platforms, Inc. and affiliates. | |
| # All rights reserved. | |
| # | |
| # This source code is licensed under the license found in the | |
| # LICENSE file in the root directory of this source tree. | |
| from typing import Union | |
| 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: Union[np.array, torch.Tensor], | |
| internal_kernel_label: Union[np.array, torch.Tensor], | |
| min_kernel_area: int, kernel_num: int) -> list: | |
| """Expand kernel contours so that foreground pixels are assigned into | |
| instances. | |
| Args: | |
| kernel_mask (np.array or torch.Tensor): The instance kernel mask with | |
| size hxw. | |
| internal_kernel_label (np.array or torch.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: | |
| 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() # type: ignore | |
| else: | |
| label = ext_module.contour_expand(kernel_mask, internal_kernel_label, | |
| min_kernel_area, kernel_num) | |
| return label | |