Add type hints in mmcv/ops/pixel_group.py (#2016)

* add typehint in mmcv/ops/pixel_group.py

* fix indent

* fix indent

* improve format

Co-authored-by: Zaida Zhou <58739961+zhouzaida@users.noreply.github.com>
pull/2032/head
gy77 2022-05-31 10:06:05 +08:00 committed by GitHub
parent 9ea55badcb
commit d499abf2fd
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 12 additions and 2 deletions

View File

@ -1,14 +1,24 @@
# Copyright (c) OpenMMLab. All rights reserved.
from typing import List, Union
import numpy as np
import torch
from torch import Tensor
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):
def pixel_group(
score: Union[np.ndarray, Tensor],
mask: Union[np.ndarray, Tensor],
embedding: Union[np.ndarray, Tensor],
kernel_label: Union[np.ndarray, Tensor],
kernel_contour: Union[np.ndarray, Tensor],
kernel_region_num: int,
distance_threshold: float,
) -> List[List[float]]:
"""Group pixels into text instances, which is widely used text detection
methods.