| |
| import numpy as np |
| import torch |
|
|
| from ..utils import ext_loader |
|
|
| ext_module = ext_loader.load_ext('_ext', ['pixel_group']) |
|
|
|
|
| def pixel_group(score, mask, embedding, kernel_label, kernel_contour, |
| kernel_region_num, distance_threshold): |
| """Group pixels into text instances, which is widely used text detection |
| methods. |
| |
| Arguments: |
| score (np.array or Tensor): The foreground score with size hxw. |
| mask (np.array or Tensor): The foreground mask with size hxw. |
| embedding (np.array or Tensor): The embedding with size hxwxc to |
| distinguish instances. |
| kernel_label (np.array or Tensor): The instance kernel index with |
| size hxw. |
| kernel_contour (np.array or Tensor): The kernel contour with size hxw. |
| kernel_region_num (int): The instance kernel region number. |
| distance_threshold (float): The embedding distance threshold between |
| kernel and pixel in one instance. |
| |
| Returns: |
| pixel_assignment (List[List[float]]): The instance coordinate list. |
| Each element consists of averaged confidence, pixel number, and |
| coordinates (x_i, y_i for all pixels) in order. |
| """ |
| assert isinstance(score, (torch.Tensor, np.ndarray)) |
| assert isinstance(mask, (torch.Tensor, np.ndarray)) |
| assert isinstance(embedding, (torch.Tensor, np.ndarray)) |
| assert isinstance(kernel_label, (torch.Tensor, np.ndarray)) |
| assert isinstance(kernel_contour, (torch.Tensor, np.ndarray)) |
| assert isinstance(kernel_region_num, int) |
| assert isinstance(distance_threshold, float) |
|
|
| if isinstance(score, np.ndarray): |
| score = torch.from_numpy(score) |
| if isinstance(mask, np.ndarray): |
| mask = torch.from_numpy(mask) |
| if isinstance(embedding, np.ndarray): |
| embedding = torch.from_numpy(embedding) |
| if isinstance(kernel_label, np.ndarray): |
| kernel_label = torch.from_numpy(kernel_label) |
| if isinstance(kernel_contour, np.ndarray): |
| kernel_contour = torch.from_numpy(kernel_contour) |
|
|
| if torch.__version__ == 'parrots': |
| label = ext_module.pixel_group( |
| score, |
| mask, |
| embedding, |
| kernel_label, |
| kernel_contour, |
| kernel_region_num=kernel_region_num, |
| distance_threshold=distance_threshold) |
| label = label.tolist() |
| label = label[0] |
| list_index = kernel_region_num |
| pixel_assignment = [] |
| for x in range(kernel_region_num): |
| pixel_assignment.append( |
| np.array( |
| label[list_index:list_index + int(label[x])], |
| dtype=np.float)) |
| list_index = list_index + int(label[x]) |
| else: |
| pixel_assignment = ext_module.pixel_group(score, mask, embedding, |
| kernel_label, kernel_contour, |
| kernel_region_num, |
| distance_threshold) |
| return pixel_assignment |
|
|
