Spaces:
Runtime error
Runtime error
File size: 6,549 Bytes
fc8c192 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 |
from __future__ import absolute_import, division, print_function
import os
import sys
import paddle
from .extract_textpoint_fast import *
from .extract_textpoint_slow import *
__dir__ = os.path.dirname(__file__)
sys.path.append(__dir__)
sys.path.append(os.path.join(__dir__, ".."))
class PGNet_PostProcess(object):
# two different post-process
def __init__(
self, character_dict_path, valid_set, score_thresh, outs_dict, shape_list
):
self.Lexicon_Table = get_dict(character_dict_path)
self.valid_set = valid_set
self.score_thresh = score_thresh
self.outs_dict = outs_dict
self.shape_list = shape_list
def pg_postprocess_fast(self):
p_score = self.outs_dict["f_score"]
p_border = self.outs_dict["f_border"]
p_char = self.outs_dict["f_char"]
p_direction = self.outs_dict["f_direction"]
if isinstance(p_score, paddle.Tensor):
p_score = p_score[0].numpy()
p_border = p_border[0].numpy()
p_direction = p_direction[0].numpy()
p_char = p_char[0].numpy()
else:
p_score = p_score[0]
p_border = p_border[0]
p_direction = p_direction[0]
p_char = p_char[0]
src_h, src_w, ratio_h, ratio_w = self.shape_list[0]
instance_yxs_list, seq_strs = generate_pivot_list_fast(
p_score,
p_char,
p_direction,
self.Lexicon_Table,
score_thresh=self.score_thresh,
)
poly_list, keep_str_list = restore_poly(
instance_yxs_list,
seq_strs,
p_border,
ratio_w,
ratio_h,
src_w,
src_h,
self.valid_set,
)
data = {
"points": poly_list,
"texts": keep_str_list,
}
return data
def pg_postprocess_slow(self):
p_score = self.outs_dict["f_score"]
p_border = self.outs_dict["f_border"]
p_char = self.outs_dict["f_char"]
p_direction = self.outs_dict["f_direction"]
if isinstance(p_score, paddle.Tensor):
p_score = p_score[0].numpy()
p_border = p_border[0].numpy()
p_direction = p_direction[0].numpy()
p_char = p_char[0].numpy()
else:
p_score = p_score[0]
p_border = p_border[0]
p_direction = p_direction[0]
p_char = p_char[0]
src_h, src_w, ratio_h, ratio_w = self.shape_list[0]
is_curved = self.valid_set == "totaltext"
char_seq_idx_set, instance_yxs_list = generate_pivot_list_slow(
p_score,
p_char,
p_direction,
score_thresh=self.score_thresh,
is_backbone=True,
is_curved=is_curved,
)
seq_strs = []
for char_idx_set in char_seq_idx_set:
pr_str = "".join([self.Lexicon_Table[pos] for pos in char_idx_set])
seq_strs.append(pr_str)
poly_list = []
keep_str_list = []
all_point_list = []
all_point_pair_list = []
for yx_center_line, keep_str in zip(instance_yxs_list, seq_strs):
if len(yx_center_line) == 1:
yx_center_line.append(yx_center_line[-1])
offset_expand = 1.0
if self.valid_set == "totaltext":
offset_expand = 1.2
point_pair_list = []
for batch_id, y, x in yx_center_line:
offset = p_border[:, y, x].reshape(2, 2)
if offset_expand != 1.0:
offset_length = np.linalg.norm(offset, axis=1, keepdims=True)
expand_length = np.clip(
offset_length * (offset_expand - 1), a_min=0.5, a_max=3.0
)
offset_detal = offset / offset_length * expand_length
offset = offset + offset_detal
ori_yx = np.array([y, x], dtype=np.float32)
point_pair = (
(ori_yx + offset)[:, ::-1]
* 4.0
/ np.array([ratio_w, ratio_h]).reshape(-1, 2)
)
point_pair_list.append(point_pair)
all_point_list.append(
[int(round(x * 4.0 / ratio_w)), int(round(y * 4.0 / ratio_h))]
)
all_point_pair_list.append(point_pair.round().astype(np.int32).tolist())
detected_poly, pair_length_info = point_pair2poly(point_pair_list)
detected_poly = expand_poly_along_width(
detected_poly, shrink_ratio_of_width=0.2
)
detected_poly[:, 0] = np.clip(detected_poly[:, 0], a_min=0, a_max=src_w)
detected_poly[:, 1] = np.clip(detected_poly[:, 1], a_min=0, a_max=src_h)
if len(keep_str) < 2:
continue
keep_str_list.append(keep_str)
detected_poly = np.round(detected_poly).astype("int32")
if self.valid_set == "partvgg":
middle_point = len(detected_poly) // 2
detected_poly = detected_poly[
[0, middle_point - 1, middle_point, -1], :
]
poly_list.append(detected_poly)
elif self.valid_set == "totaltext":
poly_list.append(detected_poly)
else:
print("--> Not supported format.")
exit(-1)
data = {
"points": poly_list,
"texts": keep_str_list,
}
return data
class PGPostProcess(object):
"""
The post process for PGNet.
"""
def __init__(self, character_dict_path, valid_set, score_thresh, mode, **kwargs):
self.character_dict_path = character_dict_path
self.valid_set = valid_set
self.score_thresh = score_thresh
self.mode = mode
# c++ la-nms is faster, but only support python 3.5
self.is_python35 = False
if sys.version_info.major == 3 and sys.version_info.minor == 5:
self.is_python35 = True
def __call__(self, outs_dict, shape_list):
post = PGNet_PostProcess(
self.character_dict_path,
self.valid_set,
self.score_thresh,
outs_dict,
shape_list,
)
if self.mode == "fast":
data = post.pg_postprocess_fast()
else:
data = post.pg_postprocess_slow()
return data
|