mirror of
https://github.com/RapidAI/RapidOCR.git
synced 2026-11-10 21:10:55 +08:00
fix(cls): classify the middle of long lines with the v2.0 model (#748)
Fixes #744.
This commit is contained in:
1 parent
7cfd3d9ef7
commit
8ac45f4e73
2 files changed
+30
No files matched your search
@@ -33,6 +33,10 @@ CLS_SHAPE_BY_OCR_VERSION = {
|
||||
class TextClassifier:
|
||||
def __init__(self, cfg: Dict[str, Any]):
|
||||
self.cls_image_shape = CLS_SHAPE_BY_OCR_VERSION[cfg["ocr_version"]]
|
||||
# The v2.0 classifier cannot read a line squeezed to its input width.
|
||||
# PP-LCNet textline_ori (PP-OCRv5) is trained on squeezed lines and
|
||||
# must keep receiving them.
|
||||
self.crop_long_lines = cfg["ocr_version"] == OCRVersion.PPOCRV4
|
||||
self.cls_batch_num = cfg["cls_batch_num"]
|
||||
self.cls_thresh = cfg["cls_thresh"]
|
||||
self.postprocess_op = ClsPostProcess(cfg["label_list"])
|
||||
@@ -83,6 +87,14 @@ class TextClassifier:
|
||||
def resize_norm_img(self, img: np.ndarray) -> np.ndarray:
|
||||
img_c, img_h, img_w = self.cls_image_shape
|
||||
h, w = img.shape[:2]
|
||||
if self.crop_long_lines:
|
||||
# Classify the middle of a long line at the model's own aspect
|
||||
# ratio instead of the whole line squeezed into img_w.
|
||||
keep_w = int(math.ceil(h * img_w / float(img_h)))
|
||||
if w > keep_w:
|
||||
left = (w - keep_w) // 2
|
||||
img = img[:, left : left + keep_w]
|
||||
w = keep_w
|
||||
ratio = w / float(h)
|
||||
if math.ceil(img_h * ratio) > img_w:
|
||||
resized_w = img_w
|
||||
|
||||
@@ -4,6 +4,7 @@
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
from pytest import mark
|
||||
|
||||
root_dir = Path(__file__).resolve().parent.parent
|
||||
@@ -52,3 +53,20 @@ def test_ppocrv5_cls_server():
|
||||
|
||||
assert result.cls_res is not None
|
||||
assert result.cls_res[0][0] == "180"
|
||||
|
||||
|
||||
def test_ppocrv4_cls_long_line():
|
||||
# A long line used to be squeezed into the 192x48 input of the v2.0 model,
|
||||
# which then missed that it was upside down (#744).
|
||||
engine = RapidOCR()
|
||||
img = cv2.imread(str(tests_dir / "test_letterbox_like.jpg"))
|
||||
thresh = engine.cfg.Cls.cls_thresh
|
||||
|
||||
result = engine(img, use_det=False, use_cls=True, use_rec=False)
|
||||
label, score = result.cls_res[0]
|
||||
assert not ("180" in label and score > thresh)
|
||||
|
||||
rotated = cv2.rotate(img, cv2.ROTATE_180)
|
||||
result = engine(rotated, use_det=False, use_cls=True, use_rec=False)
|
||||
label, score = result.cls_res[0]
|
||||
assert "180" in label and score > thresh
|
||||
Reference in new issue
Block a user