diff --git a/python/rapidocr/cal_rec_boxes/main.py b/python/rapidocr/cal_rec_boxes/main.py index 2b8728f..ed905a5 100644 --- a/python/rapidocr/cal_rec_boxes/main.py +++ b/python/rapidocr/cal_rec_boxes/main.py @@ -4,7 +4,7 @@ import copy import math from enum import Enum -from typing import List, Tuple +from typing import List, Optional, Sequence, Tuple import cv2 import numpy as np @@ -28,6 +28,7 @@ class CalRecBoxes: dt_boxes: np.ndarray, rec_res: TextRecOutput, return_single_char_box: bool = False, + rotated_180: Optional[Sequence[bool]] = None, ) -> TextRecOutput: word_results = [] for idx, (img, box) in enumerate(zip(imgs, dt_boxes)): @@ -43,6 +44,10 @@ class CalRecBoxes: return_single_char_box, ) word_box_list = self.adjust_box_overlap(copy.deepcopy(word_box_list)) + if rotated_180 is not None and rotated_180[idx]: + # The text was recognized on the crop after the classifier turned + # it by 180 degrees, so map the boxes back onto the unrotated crop. + word_box_list = self.rotate_boxes_180(word_box_list, w, h) direction = self.get_box_direction(box) word_box_list = self.reverse_rotate_crop_image( copy.deepcopy(box), word_box_list, direction @@ -54,6 +59,12 @@ class CalRecBoxes: rec_res.word_results = tuple(word_results) return rec_res + @staticmethod + def rotate_boxes_180( + word_box_list: List[List[List[float]]], w: int, h: int + ) -> List[List[List[float]]]: + return [[[w - x, h - y] for x, y in box] for box in word_box_list] + @staticmethod def get_box_direction(box: np.ndarray) -> Direction: edge_lengths = [ diff --git a/python/rapidocr/main.py b/python/rapidocr/main.py index 8ac85f1..12472c6 100644 --- a/python/rapidocr/main.py +++ b/python/rapidocr/main.py @@ -204,6 +204,8 @@ class RapidOCR: ratio_w = op_record["preprocess"]["ratio_w"] cropped_img_list = map_img_to_original(cropped_img_list, ratio_h, ratio_w) + rotated_180 = self.get_rotated_180_flags(cls_res) + # 过滤识别结果为空的值 if ( rec_res.txts is not None @@ -221,6 +223,8 @@ class RapidOCR: rec_res.scores = filter_by_indices(rec_res.scores, valid_ids) cropped_img_list = filter_by_indices(cropped_img_list, valid_ids) + if rotated_180 is not None: + rotated_180 = filter_by_indices(rotated_180, valid_ids) # 仅分类结果 if ( @@ -248,7 +252,7 @@ class RapidOCR: and all(rec_res.word_results) ): rec_res.word_results = self.calc_word_boxes( - cropped_img_list, det_res.boxes, rec_res + cropped_img_list, det_res.boxes, rec_res, rotated_180 ) ocr_res = RapidOCROutput( @@ -268,11 +272,25 @@ class RapidOCR: ocr_res = self.filter_by_text_score(ocr_res) return ocr_res if len(ocr_res) > 0 else RapidOCROutput() + def get_rotated_180_flags(self, cls_res: TextClsOutput) -> Optional[List[bool]]: + """Which crops the text classifier turned by 180 degrees before recognition.""" + if cls_res.cls_res is None: + return None + + cls_thresh = self.cfg.Cls.cls_thresh + return [ + "180" in label and score > cls_thresh for label, score in cls_res.cls_res + ] + def calc_word_boxes( - self, img: List[np.ndarray], dt_boxes: np.ndarray, rec_res: TextRecOutput + self, + img: List[np.ndarray], + dt_boxes: np.ndarray, + rec_res: TextRecOutput, + rotated_180: Optional[List[bool]] = None, ) -> Any: rec_res = self.cal_rec_boxes( - img, dt_boxes, rec_res, self.return_single_char_box + img, dt_boxes, rec_res, self.return_single_char_box, rotated_180 ) origin_words = [] diff --git a/python/tests/test_word_box_cls_rotation.py b/python/tests/test_word_box_cls_rotation.py new file mode 100644 index 0000000..6315a4b --- /dev/null +++ b/python/tests/test_word_box_cls_rotation.py @@ -0,0 +1,110 @@ +# -*- encoding: utf-8 -*- +# @Author: Mohammad Hijjawi +from typing import List, Optional, Tuple + +import numpy as np +import pytest +from rapidocr import RapidOCR +from rapidocr.ch_ppocr_cls import TextClsOutput +from rapidocr.ch_ppocr_det import TextDetOutput +from rapidocr.ch_ppocr_rec import TextRecOutput +from rapidocr.ch_ppocr_rec.typings import WordInfo, WordType + +LINE_TXT = "我是人" + +# A 200x20 horizontal line and a 20x200 vertical line in a 400x400 image. +HORIZONTAL_BOX = [[100, 50], [300, 50], [300, 70], [100, 70]] +VERTICAL_BOX = [[100, 100], [120, 100], [120, 300], [100, 300]] + + +def make_word_info() -> WordInfo: + return WordInfo( + words=[list(LINE_TXT)], + word_cols=[[2, 15, 28]], + word_types=[WordType.CN], + line_txt_len=30, + confs=[0.9, 0.9, 0.9], + ) + + +def run_word_boxes( + boxes: List[List[List[int]]], + txts: Tuple[str, ...], + cls_res: Optional[List[Tuple[str, float]]], +): + engine = RapidOCR(params={"Global.return_word_box": True}) + + # The crop handed to recognition is always laid out horizontally. + crops = [np.zeros((20, 200, 3), dtype=np.uint8) for _ in boxes] + det_res = TextDetOutput( + boxes=np.array(boxes, dtype=np.float32), scores=[0.9] * len(boxes) + ) + cls_out = TextClsOutput(img_list=crops, cls_res=cls_res) + rec_res = TextRecOutput( + imgs=crops, + txts=txts, + scores=tuple(0.9 for _ in txts), + word_results=tuple(make_word_info() for _ in txts), + ) + return engine.build_final_output( + np.zeros((400, 400, 3), dtype=np.uint8), + det_res, + cls_out, + rec_res, + crops, + {"preprocess": {"ratio_h": 1.0, "ratio_w": 1.0}}, + ) + + +def word_centers(word_line) -> List[Tuple[float, float]]: + return [tuple(np.array(points).mean(axis=0)) for _, _, points in word_line] + + +@pytest.mark.parametrize( + "box,cls_res,axis,reading_direction", + [ + (HORIZONTAL_BOX, None, 0, 1), + (HORIZONTAL_BOX, [("0", 0.99)], 0, 1), + # below cls_thresh the crop is not rotated, so neither are the boxes + (HORIZONTAL_BOX, [("180", 0.5)], 0, 1), + # issue #328: upside-down text is read right to left in the image + (HORIZONTAL_BOX, [("180", 0.99)], 0, -1), + (VERTICAL_BOX, [("0", 0.99)], 1, 1), + (VERTICAL_BOX, [("180", 0.99)], 1, -1), + ], +) +def test_word_boxes_follow_cls_rotation(box, cls_res, axis, reading_direction): + result = run_word_boxes([box], (LINE_TXT,), cls_res) + + word_line = result.word_results[0] + assert "".join(txt for txt, _, _ in word_line) == LINE_TXT + + centers = [center[axis] for center in word_centers(word_line)] + assert np.all(reading_direction * np.diff(centers) > 0) + + +def test_rotated_word_boxes_are_mirrored_within_line(): + upright = run_word_boxes([HORIZONTAL_BOX], (LINE_TXT,), [("0", 0.99)]) + flipped = run_word_boxes([HORIZONTAL_BOX], (LINE_TXT,), [("180", 0.99)]) + + upright_xs = [ + sorted({p[0] for p in points}) for _, _, points in upright.word_results[0] + ] + flipped_xs = [ + sorted({p[0] for p in points}) for _, _, points in flipped.word_results[0] + ] + # Each character's box is mirrored about the centre of the line. + for (x0, x1), (fx0, fx1) in zip(upright_xs, flipped_xs): + assert (fx0, fx1) == (400 - x1, 400 - x0) + + +def test_rotation_flags_stay_aligned_after_empty_lines_are_dropped(): + result = run_word_boxes( + [HORIZONTAL_BOX, VERTICAL_BOX], + ("", LINE_TXT), + [("0", 0.99), ("180", 0.99)], + ) + + assert result.txts == (LINE_TXT,) + ys = [center[1] for center in word_centers(result.word_results[0])] + assert np.all(np.diff(ys) < 0)