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