File size: 4,744 Bytes
948d40c | 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 | import math
import cv2
import numpy as np
import triton_python_backend_utils as pb_utils # type: ignore
class TritonPythonModel:
def initialize(self, args):
self.logger = pb_utils.Logger
self.text_threshold = 0.7
self.link_threshold = 0.35
self.low_text = 0.4
def _getDetBoxes_core(self, textmap, linkmap, text_threshold, link_threshold, low_text):
# prepare data
linkmap = linkmap.copy()
textmap = textmap.copy()
img_h, img_w = textmap.shape
""" labeling method """
_, text_score = cv2.threshold(textmap, low_text, 1, 0)
_, link_score = cv2.threshold(linkmap, link_threshold, 1, 0)
text_score_comb = np.clip(text_score + link_score, 0, 1)
nLabels, labels, stats, _ = cv2.connectedComponentsWithStats(
text_score_comb.astype(np.uint8),
connectivity=4
)
detected_text_box_list = []
for k in range(1, nLabels):
# size filtering
size = stats[k, cv2.CC_STAT_AREA]
if size < 10: continue
# thresholding
if np.max(textmap[labels == k]) < text_threshold: continue
# make segmentation map
segmap = np.zeros(textmap.shape, dtype=np.uint8)
segmap[labels == k] = 255
segmap[np.logical_and(link_score == 1, text_score == 0)] = 0 # remove link area
x, y = stats[k, cv2.CC_STAT_LEFT], stats[k, cv2.CC_STAT_TOP]
w, h = stats[k, cv2.CC_STAT_WIDTH], stats[k, cv2.CC_STAT_HEIGHT]
niter = int(math.sqrt(size * min(w, h) / (w * h)) * 2)
sx, ex, sy, ey = x - niter, x + w + niter + 1, y - niter, y + h + niter + 1
# boundary check
if sx < 0: sx = 0
if sy < 0: sy = 0
if ex >= img_w: ex = img_w
if ey >= img_h: ey = img_h
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1 + niter, 1 + niter))
segmap[sy:ey, sx:ex] = cv2.dilate(segmap[sy:ey, sx:ex], kernel)
# make box
np_contours = np.roll(np.array(np.where(segmap != 0)), 1, axis=0).transpose().reshape(-1, 2)
rectangle = cv2.minAreaRect(np_contours)
box = cv2.boxPoints(rectangle)
# align diamond-shape
w, h = np.linalg.norm(box[0] - box[1]), np.linalg.norm(box[1] - box[2])
box_ratio = max(w, h) / (min(w, h) + 1e-5)
if abs(1 - box_ratio) <= 0.1:
l, r = min(np_contours[:, 0]), max(np_contours[:, 0])
t, b = min(np_contours[:, 1]), max(np_contours[:, 1])
box = np.array([[l, t], [r, t], [r, b], [l, b]], dtype=np.float32)
# make clock-wise order
startidx = box.sum(axis=1).argmin()
box = np.roll(box, 4 - startidx, 0)
box = np.array(box)
margin = int(min(w, h)* 0.02)
box[0][1] = max(0, box[0][1] - margin)
box[1][1] = max(0, box[1][1] - margin)
detected_text_box_list.append(box)
return detected_text_box_list
def _postprocess(self, td_output, resize_ratio):
out = td_output[0]
score_text = np.float32(out[:, :, 0])
score_link = np.float32(out[:, :, 1])
detected_text_box_list = self._getDetBoxes_core(
textmap=score_text,
linkmap=score_link,
text_threshold=self.text_threshold,
link_threshold=self.link_threshold,
low_text=self.low_text
)
for i in range(len(detected_text_box_list)):
detected_text_box_list[i] = detected_text_box_list[i] * 1/resize_ratio * 2
if len(detected_text_box_list) == 0:
output_np = np.zeros((0, 4, 2), dtype=np.float32)
else:
output_np = np.array(detected_text_box_list, dtype=np.float32)
out_tensor_0 = pb_utils.Tensor("detected_text_box_list", output_np)
return out_tensor_0
def execute(self, requests):
responses = []
for request in requests:
try:
td_output = pb_utils.get_input_tensor_by_name(request, "td_output").as_numpy()
resize_ratio = pb_utils.get_input_tensor_by_name(request, "resize_ratio").as_numpy()
resize_ratio = resize_ratio[0]
out_tensor_0 = self._postprocess(td_output, resize_ratio)
except pb_utils.TritonModelException as e:
responses.append(pb_utils.InferenceResponse(error=pb_utils.TritonError(str(e), pb_utils.TritonError.BAD_REQUEST)))
continue
responses.append(pb_utils.InferenceResponse(output_tensors=[out_tensor_0]))
return responses
|