mirror of
https://github.com/RapidAI/RapidOCR.git
synced 2026-11-11 05:20:49 +08:00
feat: add PP-OCRv5 cls module
This commit is contained in:
2 files changed
+11
-3
No files matched your search
@@ -20,13 +20,19 @@ import cv2
|
||||
import numpy as np
|
||||
|
||||
from rapidocr.inference_engine.base import get_engine
|
||||
from rapidocr.utils.typings import OCRVersion
|
||||
|
||||
from .utils import ClsPostProcess, TextClsOutput
|
||||
|
||||
CLS_SHAPE_BY_OCR_VERSION = {
|
||||
OCRVersion.PPOCRV4: [3, 48, 192],
|
||||
OCRVersion.PPOCRV5: [3, 80, 160],
|
||||
}
|
||||
|
||||
|
||||
class TextClassifier:
|
||||
def __init__(self, cfg: Dict[str, Any]):
|
||||
self.cls_image_shape = cfg["cls_image_shape"]
|
||||
self.cls_image_shape = CLS_SHAPE_BY_OCR_VERSION[cfg["ocr_version"]]
|
||||
self.cls_batch_num = cfg["cls_batch_num"]
|
||||
self.cls_thresh = cfg["cls_thresh"]
|
||||
self.postprocess_op = ClsPostProcess(cfg["label_list"])
|
||||
@@ -83,7 +89,9 @@ class TextClassifier:
|
||||
else:
|
||||
resized_w = int(math.ceil(img_h * ratio))
|
||||
|
||||
resized_image = cv2.resize(img, (resized_w, img_h))
|
||||
resized_image = cv2.resize(
|
||||
img, (resized_w, img_h), interpolation=cv2.INTER_LINEAR
|
||||
)
|
||||
resized_image = resized_image.astype("float32")
|
||||
if img_c == 1:
|
||||
resized_image = resized_image / 255
|
||||
|
||||
@@ -136,7 +136,7 @@ class InferSession(abc.ABC):
|
||||
return model_dict[k]
|
||||
|
||||
for k in model_dict:
|
||||
if k.startswith(lang_type):
|
||||
if k.startswith(lang_type) and model_type in k:
|
||||
return model_dict[k]
|
||||
|
||||
logger.error(
|
||||
|
||||
Reference in new issue
Block a user