Parameters: img: RGB 0~255.
(self, img, bg=None, sf=4, is_aligned=False, lq_label=None)
| 87 | |
| 88 | |
| 89 | def handle_texts(self, img, bg=None, sf=4, is_aligned=False, lq_label=None): |
| 90 | ''' |
| 91 | Parameters: |
| 92 | img: RGB 0~255. |
| 93 | ''' |
| 94 | |
| 95 | height, width = img.shape[:2] |
| 96 | bg_height, bg_width = bg.shape[:2] |
| 97 | print(' ' * 25 + f' ... The input->output image size is {bg_height//sf}*{bg_width//sf}->{bg_height}*{bg_width}') |
| 98 | |
| 99 | full_mask_blur = np.zeros(bg.shape, dtype=np.float32) |
| 100 | full_mask_noblur = np.zeros(bg.shape, dtype=np.float32) |
| 101 | full_text_img = np.zeros(bg.shape, dtype=np.float32) #+255 |
| 102 | |
| 103 | orig_texts, enhanced_texts, debug_texts, pred_texts = [], [], [], [] |
| 104 | ocr_scores = [] |
| 105 | |
| 106 | if not is_aligned: |
| 107 | box_infos = self.cnstd.detect(img) |
| 108 | for iix, box_info in enumerate(box_infos['detected_texts']): |
| 109 | box = box_info['box'].astype(int)# left top, right top, right bottom, left bottom, [width, height] |
| 110 | score = box_info['score'] |
| 111 | if score < 0.5: |
| 112 | continue |
| 113 | |
| 114 | extend_box = box.copy() |
| 115 | w = int(np.linalg.norm(box[0] - box[1])) |
| 116 | h = int(np.linalg.norm(box[0] - box[3])) |
| 117 | |
| 118 | # extend the bounding box |
| 119 | extend_lr = 0.15 * h |
| 120 | extend_tb = 0.05 * h |
| 121 | vec_w = (box[1] - box[0]) / w |
| 122 | vec_h = (box[3] - box[0]) / h |
| 123 | |
| 124 | extend_box[0] = box[0] - vec_w * extend_lr - vec_h * extend_tb |
| 125 | extend_box[1] = box[1] + vec_w * extend_lr - vec_h * extend_tb |
| 126 | extend_box[2] = box[2] + vec_w * extend_lr + vec_h * extend_tb |
| 127 | extend_box[3] = box[3] - vec_w * extend_lr + vec_h * extend_tb |
| 128 | extend_box = extend_box.astype(int) |
| 129 | |
| 130 | w = int(np.linalg.norm(extend_box[0] - extend_box[1])) |
| 131 | h = int(np.linalg.norm(extend_box[0] - extend_box[3])) |
| 132 | |
| 133 | if w > h: |
| 134 | ref_h = self.insize |
| 135 | ref_w = int(ref_h * w / h) |
| 136 | else: |
| 137 | print(' ' * 25 + ' ... Can not handle vertical text temporarily') |
| 138 | continue |
| 139 | |
| 140 | ref_point = np.float32([[0,0], [ref_w, 0], [ref_w, ref_h], [0, ref_h]]) |
| 141 | det_point = np.float32(extend_box) |
| 142 | |
| 143 | matrix = cv2.getPerspectiveTransform(det_point, ref_point) |
| 144 | inv_matrix = cv2.getPerspectiveTransform(ref_point*sf, det_point*sf) |
| 145 | |
| 146 | cropped_img = cv2.warpPerspective(img, matrix, (ref_w, ref_h), borderMode=cv2.BORDER_REPLICATE, flags=cv2.INTER_LINEAR) |