diff --git a/tools/infer/predict_system.py b/tools/infer/predict_system.py index affd0d1bcd1283be02ead3cd61c01c375b49bdf9..19e0525a2a72d9148ef31405d51e4364cec36848 100755 --- a/tools/infer/predict_system.py +++ b/tools/infer/predict_system.py @@ -34,7 +34,7 @@ import tools.infer.predict_det as predict_det import tools.infer.predict_cls as predict_cls from ppocr.utils.utility import get_image_file_list, check_and_read from ppocr.utils.logging import get_logger -from tools.infer.utility import draw_ocr_box_txt, get_rotate_crop_image +from tools.infer.utility import draw_ocr_box_txt, get_rotate_crop_image, get_minarea_rect_crop logger = get_logger() @@ -80,7 +80,10 @@ class TextSystem(object): for bno in range(len(dt_boxes)): tmp_box = copy.deepcopy(dt_boxes[bno]) - img_crop = get_rotate_crop_image(ori_im, tmp_box) + if self.args.det_box_type == "quad": + img_crop = get_rotate_crop_image(ori_im, tmp_box) + else: + img_crop = get_minarea_rect_crop(ori_im, tmp_box) img_crop_list.append(img_crop) if self.use_angle_cls and cls: img_crop_list, angle_list, elapse = self.text_classifier( diff --git a/tools/infer/utility.py b/tools/infer/utility.py index 34cad2590f2904f79709530acf841033c89088e0..207b2e5f35dec8934fd2875748e8bcc21a94b7bf 100644 --- a/tools/infer/utility.py +++ b/tools/infer/utility.py @@ -629,6 +629,29 @@ def get_rotate_crop_image(img, points): return dst_img +def get_minarea_rect_crop(img, points): + bounding_box = cv2.minAreaRect(np.array(points).astype(np.int32)) + points = sorted(list(cv2.boxPoints(bounding_box)), key=lambda x: x[0]) + + index_a, index_b, index_c, index_d = 0, 1, 2, 3 + if points[1][1] > points[0][1]: + index_a = 0 + index_d = 1 + else: + index_a = 1 + index_d = 0 + if points[3][1] > points[2][1]: + index_b = 2 + index_c = 3 + else: + index_b = 3 + index_c = 2 + + box = [points[index_a], points[index_b], points[index_c], points[index_d]] + crop_img = get_rotate_crop_image(img, np.array(box)) + return crop_img + + def check_gpu(use_gpu): if use_gpu and not paddle.is_compiled_with_cuda(): use_gpu = False