| import math |
|
|
| import cv2 |
| import numpy as np |
| import triton_python_backend_utils as pb_utils |
|
|
|
|
| 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): |
| |
| 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 = stats[k, cv2.CC_STAT_AREA] |
| if size < 10: continue |
|
|
| |
| if np.max(textmap[labels == k]) < text_threshold: continue |
|
|
| |
| segmap = np.zeros(textmap.shape, dtype=np.uint8) |
| segmap[labels == k] = 255 |
| segmap[np.logical_and(link_score == 1, text_score == 0)] = 0 |
| 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 |
| |
| 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) |
|
|
| |
| 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) |
|
|
| |
| 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) |
|
|
| |
| 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 |
|
|
