mtfyun/rapidocr
1
0 Fork 0
mirror of https://github.com/RapidAI/RapidOCR.git synced 2026-11-11 05:20:49 +08:00
Code Issues Packages Projects Releases Wiki Activity

Format by black

Browse Source
This commit is contained in:
SWHL committed 2023-06-24 10:27:12 +08:00
1 parent ad4631fdd4
commit f4fd5cdcba
40 files changed
+908 -854
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

python/demo.py
+40 -34
View File
@@ -10,18 +10,20 @@ import numpy as np
from PIL import Image, ImageDraw, ImageFont
from rapidocr_onnxruntime import RapidOCR
# from rapidocr_openvino import RapidOCR
def draw_ocr_box_txt(image, boxes, txts, font_path,
scores=None, text_score=0.5):
def draw_ocr_box_txt(image, boxes, txts, font_path, scores=None, text_score=0.5):
if not Path(font_path).exists():
raise FileNotFoundError(f'The {font_path} does not exists! \n'
f'Please download the file in the https://drive.google.com/file/d/1evWVX38EFNwTq_n5gTFgnlv8tdaNcyIA/view?usp=sharing')
raise FileNotFoundError(
f"The {font_path} does not exists! \n"
f"Please download the file in the https://drive.google.com/file/d/1evWVX38EFNwTq_n5gTFgnlv8tdaNcyIA/view?usp=sharing"
)
h, w = image.height, image.width
img_left = image.copy()
img_right = Image.new('RGB', (w, h), (255, 255, 255))
img_right = Image.new("RGB", (w, h), (255, 255, 255))
random.seed(0)
draw_left = ImageDraw.Draw(img_left)
@@ -30,40 +32,45 @@ def draw_ocr_box_txt(image, boxes, txts, font_path,
if scores is not None and float(scores[idx]) < text_score:
continue
color = (random.randint(0, 255),
random.randint(0, 255),
random.randint(0, 255))
color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
draw_left.polygon(box, fill=color)
draw_right.polygon([box[0][0], box[0][1],
box[1][0], box[1][1],
box[2][0], box[2][1],
box[3][0], box[3][1]],
outline=color)
draw_right.polygon(
[
box[0][0],
box[0][1],
box[1][0],
box[1][1],
box[2][0],
box[2][1],
box[3][0],
box[3][1],
],
outline=color,
)
box_height = math.sqrt((box[0][0] - box[3][0])**2
+ (box[0][1] - box[3][1])**2)
box_height = math.sqrt(
(box[0][0] - box[3][0]) ** 2 + (box[0][1] - box[3][1]) ** 2
)
box_width = math.sqrt((box[0][0] - box[1][0])**2
+ (box[0][1] - box[1][1])**2)
box_width = math.sqrt(
(box[0][0] - box[1][0]) ** 2 + (box[0][1] - box[1][1]) ** 2
)
if box_height > 2 * box_width:
font_size = max(int(box_width * 0.9), 10)
font = ImageFont.truetype(font_path, font_size,
encoding="utf-8")
font = ImageFont.truetype(font_path, font_size, encoding="utf-8")
cur_y = box[0][1]
for c in txt:
char_size = font.getsize(c)
draw_right.text((box[0][0] + 3, cur_y), c,
fill=(0, 0, 0), font=font)
draw_right.text((box[0][0] + 3, cur_y), c, fill=(0, 0, 0), font=font)
cur_y += char_size[1]
else:
font_size = max(int(box_height * 0.8), 10)
font = ImageFont.truetype(font_path, font_size, encoding="utf-8")
draw_right.text([box[0][0], box[0][1]], txt,
fill=(0, 0, 0), font=font)
draw_right.text([box[0][0], box[0][1]], txt, fill=(0, 0, 0), font=font)
img_left = Image.blend(image, img_left, 0.5)
img_show = Image.new('RGB', (w * 2, h), (255, 255, 255))
img_show = Image.new("RGB", (w * 2, h), (255, 255, 255))
img_show.paste(img_left, (0, 0, w, h))
img_show.paste(img_right, (w, 0, w * 2, h))
return np.array(img_show)
@@ -73,29 +80,28 @@ def visualize(image_path, result, font_path="resources/fonts/FZYTK.TTF"):
image = Image.open(image_path)
boxes, txts, scores = list(zip(*result))
draw_img = draw_ocr_box_txt(image, np.array(boxes),
txts, font_path,
scores,
text_score=0.5)
draw_img = draw_ocr_box_txt(
image, np.array(boxes), txts, font_path, scores, text_score=0.5
)
draw_img_save = Path("./inference_results/")
if not draw_img_save.exists():
draw_img_save.mkdir(parents=True, exist_ok=True)
image_save = str(draw_img_save / f'infer_{Path(image_path).name}')
image_save = str(draw_img_save / f"infer_{Path(image_path).name}")
cv2.imwrite(image_save, draw_img[:, :, ::-1])
print(f'The infer result has saved in {image_save}')
print(f"The infer result has saved in {image_save}")
if __name__ == '__main__':
if __name__ == "__main__":
rapid_ocr = RapidOCR()
image_path = 'tests/test_files/ch_en_num.jpg'
with open(image_path, 'rb') as f:
image_path = "tests/test_files/ch_en_num.jpg"
with open(image_path, "rb") as f:
img = f.read()
result, elapse_list = rapid_ocr(img)
print(result)
print(elapse_list)
if result:
visualize(image_path, result, font_path='resources/fonts/FZYTK.TTF')
visualize(image_path, result, font_path="resources/fonts/FZYTK.TTF")
python/rapidocr_onnxruntime/ch_ppocr_v2_cls/text_cls.py
+12 -11
View File
@@ -25,12 +25,12 @@ from rapidocr_onnxruntime.utils import OrtInferSession, read_yaml
from .utils import ClsPostProcess
class TextClassifier():
class TextClassifier:
def __init__(self, config):
self.cls_image_shape = config['cls_image_shape']
self.cls_batch_num = config['cls_batch_num']
self.cls_thresh = config['cls_thresh']
self.postprocess_op = ClsPostProcess(config['label_list'])
self.cls_image_shape = config["cls_image_shape"]
self.cls_batch_num = config["cls_batch_num"]
self.cls_thresh = config["cls_thresh"]
self.postprocess_op = ClsPostProcess(config["label_list"])
self.infer = OrtInferSession(config)
@@ -47,7 +47,7 @@ class TextClassifier():
indices = np.argsort(np.array(width_list))
img_num = len(img_list)
cls_res = [['', 0.0]] * img_num
cls_res = [["", 0.0]] * img_num
batch_num = self.cls_batch_num
elapse = 0
for beg_img_no in range(0, img_num, batch_num):
@@ -68,9 +68,10 @@ class TextClassifier():
for rno in range(len(cls_result)):
label, score = cls_result[rno]
cls_res[indices[beg_img_no + rno]] = [label, score]
if '180' in label and score > self.cls_thresh:
if "180" in label and score > self.cls_thresh:
img_list[indices[beg_img_no + rno]] = cv2.rotate(
img_list[indices[beg_img_no + rno]], 1)
img_list[indices[beg_img_no + rno]], 1
)
return img_list, cls_res, elapse
def resize_norm_img(self, img):
@@ -83,7 +84,7 @@ class TextClassifier():
resized_w = int(math.ceil(img_h * ratio))
resized_image = cv2.resize(img, (resized_w, img_h))
resized_image = resized_image.astype('float32')
resized_image = resized_image.astype("float32")
if img_c == 1:
resized_image = resized_image / 255
resized_image = resized_image[np.newaxis, :]
@@ -99,8 +100,8 @@ class TextClassifier():
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--image_path', type=str, help='image_dir|image_path')
parser.add_argument('--config_path', type=str, default='config.yaml')
parser.add_argument("--image_path", type=str, help="image_dir|image_path")
parser.add_argument("--config_path", type=str, default="config.yaml")
args = parser.parse_args()
config = read_yaml(args.config_path)
python/rapidocr_onnxruntime/ch_ppocr_v2_cls/utils.py
+5 -4
View File
@@ -11,8 +11,8 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
class ClsPostProcess():
""" Convert between text-label and text-index """
class ClsPostProcess:
"""Convert between text-label and text-index"""
def __init__(self, label_list):
super(ClsPostProcess, self).__init__()
@@ -20,8 +20,9 @@ class ClsPostProcess():
def __call__(self, preds, label=None):
pred_idxs = preds.argmax(axis=1)
decode_out = [(self.label_list[idx], preds[i, idx])
for i, idx in enumerate(pred_idxs)]
decode_out = [
(self.label_list[idx], preds[i, idx]) for i, idx in enumerate(pred_idxs)
]
if label is None:
return decode_out
python/rapidocr_onnxruntime/ch_ppocr_v3_det/text_detect.py
+25 -26
View File
@@ -25,33 +25,31 @@ from rapidocr_onnxruntime.utils import OrtInferSession, read_yaml
from .utils import DBPostProcess, create_operators, transform
class TextDetector():
class TextDetector:
def __init__(self, config):
pre_process_list = {
'DetResizeForTest': {
'limit_side_len': config.get('limit_side_len', 736),
'limit_type': config.get('limit_type', 'min')
"DetResizeForTest": {
"limit_side_len": config.get("limit_side_len", 736),
"limit_type": config.get("limit_type", "min"),
},
'NormalizeImage': {
'std': [0.229, 0.224, 0.225],
'mean': [0.485, 0.456, 0.406],
'scale': '1./255.',
'order': 'hwc'
"NormalizeImage": {
"std": [0.229, 0.224, 0.225],
"mean": [0.485, 0.456, 0.406],
"scale": "1./255.",
"order": "hwc",
},
'ToCHWImage': None,
'KeepKeys': {
'keep_keys': ['image', 'shape']
}
"ToCHWImage": None,
"KeepKeys": {"keep_keys": ["image", "shape"]},
}
self.preprocess_op = create_operators(pre_process_list)
post_process = {
'thresh': config.get('thresh', 0.3),
'box_thresh': config.get('box_thresh', 0.5),
'max_candidates': config.get('max_candidates', 1000),
'unclip_ratio': config.get('unclip_ratio', 1.6),
'use_dilation': config.get('use_dilation', True),
'score_mode': config.get('score_mode', 'fast'),
"thresh": config.get("thresh", 0.3),
"box_thresh": config.get("box_thresh", 0.5),
"max_candidates": config.get("max_candidates", 1000),
"unclip_ratio": config.get("unclip_ratio", 1.6),
"use_dilation": config.get("use_dilation", True),
"score_mode": config.get("score_mode", "fast"),
}
self.postprocess_op = DBPostProcess(**post_process)
@@ -59,11 +57,11 @@ class TextDetector():
def __call__(self, img):
if img is None:
raise ValueError('img is None')
raise ValueError("img is None")
ori_im_shape = img.shape[:2]
data = {'image': img}
data = {"image": img}
data = transform(data, self.preprocess_op)
img, shape_list = data
if img is None:
@@ -76,7 +74,7 @@ class TextDetector():
preds = self.infer(img)[0]
post_result = self.postprocess_op(preds, shape_list)
dt_boxes = post_result[0]['points']
dt_boxes = post_result[0]["points"]
dt_boxes = self.filter_tag_det_res(dt_boxes, ori_im_shape)
elapse = time.time() - starttime
return dt_boxes, elapse
@@ -129,8 +127,8 @@ class TextDetector():
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--config_path', type=str, default='config.yaml')
parser.add_argument('--image_path', type=str, default=None)
parser.add_argument("--config_path", type=str, default="config.yaml")
parser.add_argument("--image_path", type=str, default=None)
args = parser.parse_args()
config = read_yaml(args.config_path)
@@ -141,6 +139,7 @@ if __name__ == "__main__":
dt_boxes, elapse = text_detector(img)
from utils import draw_text_det_res
src_im = draw_text_det_res(dt_boxes, args.image_path)
cv2.imwrite('det_results.jpg', src_im)
print('The det_results.jpg has been saved in the current directory.')
cv2.imwrite("det_results.jpg", src_im)
print("The det_results.jpg has been saved in the current directory.")
python/rapidocr_onnxruntime/ch_ppocr_v3_det/utils.py
+76 -69
View File
@@ -25,69 +25,74 @@ import six
from shapely.geometry import Polygon
class DecodeImage():
""" decode image """
class DecodeImage:
"""decode image"""
def __init__(self, img_mode='RGB', channel_first=False):
def __init__(self, img_mode="RGB", channel_first=False):
self.img_mode = img_mode
self.channel_first = channel_first
def __call__(self, data):
img = data['image']
img = data["image"]
if six.PY2:
assert type(img) is str and len(img) > 0, "invalid input 'img' in DecodeImage"
assert (
type(img) is str and len(img) > 0
), "invalid input 'img' in DecodeImage"
else:
assert type(img) is bytes and len(img) > 0, "invalid input 'img' in DecodeImage"
assert (
type(img) is bytes and len(img) > 0
), "invalid input 'img' in DecodeImage"
img = np.frombuffer(img, dtype='uint8')
img = np.frombuffer(img, dtype="uint8")
img = cv2.imdecode(img, 1)
if img is None:
return None
if self.img_mode == 'GRAY':
if self.img_mode == "GRAY":
img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
elif self.img_mode == 'RGB':
assert img.shape[2] == 3, f'invalid shape of image[{img.shape}]'
elif self.img_mode == "RGB":
assert img.shape[2] == 3, f"invalid shape of image[{img.shape}]"
img = img[:, :, ::-1]
if self.channel_first:
img = img.transpose((2, 0, 1))
data['image'] = img
data["image"] = img
return data
class NormalizeImage():
""" normalize image such as substract mean, divide std"""
class NormalizeImage:
"""normalize image such as substract mean, divide std"""
def __init__(self, scale=None, mean=None, std=None, order='chw'):
def __init__(self, scale=None, mean=None, std=None, order="chw"):
if isinstance(scale, str):
scale = eval(scale)
self.scale = np.float32(scale if scale is not None else 1.0 / 255.0)
mean = mean if mean is not None else [0.485, 0.456, 0.406]
std = std if std is not None else [0.229, 0.224, 0.225]
shape = (3, 1, 1) if order == 'chw' else (1, 1, 3)
self.mean = np.array(mean).reshape(shape).astype('float32')
self.std = np.array(std).reshape(shape).astype('float32')
shape = (3, 1, 1) if order == "chw" else (1, 1, 3)
self.mean = np.array(mean).reshape(shape).astype("float32")
self.std = np.array(std).reshape(shape).astype("float32")
def __call__(self, data):
img = np.array(data['image']).astype(np.float32)
data['image'] = (img * self.scale - self.mean) / self.std
img = np.array(data["image"]).astype(np.float32)
data["image"] = (img * self.scale - self.mean) / self.std
return data
class ToCHWImage():
""" convert hwc image to chw image"""
class ToCHWImage:
"""convert hwc image to chw image"""
def __init__(self):
pass
def __call__(self, data):
img = np.array(data['image'])
data['image'] = img.transpose((2, 0, 1))
img = np.array(data["image"])
data["image"] = img.transpose((2, 0, 1))
return data
class KeepKeys():
class KeepKeys:
def __init__(self, keep_keys):
self.keep_keys = keep_keys
@@ -98,26 +103,26 @@ class KeepKeys():
return data_list
class DetResizeForTest():
class DetResizeForTest:
def __init__(self, **kwargs):
super(DetResizeForTest, self).__init__()
self.resize_type = 0
if 'image_shape' in kwargs:
self.image_shape = kwargs['image_shape']
if "image_shape" in kwargs:
self.image_shape = kwargs["image_shape"]
self.resize_type = 1
elif 'limit_side_len' in kwargs:
self.limit_side_len = kwargs.get('limit_side_len', 736)
self.limit_type = kwargs.get('limit_type', 'min')
elif "limit_side_len" in kwargs:
self.limit_side_len = kwargs.get("limit_side_len", 736)
self.limit_type = kwargs.get("limit_type", "min")
if 'resize_long' in kwargs:
if "resize_long" in kwargs:
self.resize_type = 2
self.resize_long = kwargs.get('resize_long', 960)
self.resize_long = kwargs.get("resize_long", 960)
else:
self.limit_side_len = kwargs.get('limit_side_len', 736)
self.limit_type = kwargs.get('limit_type', 'min')
self.limit_side_len = kwargs.get("limit_side_len", 736)
self.limit_type = kwargs.get("limit_type", "min")
def __call__(self, data):
img = data['image']
img = data["image"]
src_h, src_w = img.shape[:2]
if self.resize_type == 0:
@@ -128,8 +133,8 @@ class DetResizeForTest():
else:
# img, shape = self.resize_image_type1(img)
img, [ratio_h, ratio_w] = self.resize_image_type1(img)
data['image'] = img
data['shape'] = np.array([src_h, src_w, ratio_h, ratio_w])
data["image"] = img
data["shape"] = np.array([src_h, src_w, ratio_h, ratio_w])
return data
def resize_image_type1(self, img):
@@ -153,14 +158,14 @@ class DetResizeForTest():
h, w = img.shape[:2]
# limit the max side
if self.limit_type == 'max':
if self.limit_type == "max":
if max(h, w) > limit_side_len:
if h > w:
ratio = float(limit_side_len) / h
else:
ratio = float(limit_side_len) / w
else:
ratio = 1.
ratio = 1.0
else:
if min(h, w) < limit_side_len:
if h < w:
@@ -168,7 +173,7 @@ class DetResizeForTest():
else:
ratio = float(limit_side_len) / w
else:
ratio = 1.
ratio = 1.0
resize_h = int(h * ratio)
resize_w = int(w * ratio)
@@ -212,7 +217,7 @@ class DetResizeForTest():
def transform(data, ops=None):
""" transform """
"""transform"""
if ops is None:
ops = []
@@ -240,21 +245,22 @@ def draw_text_det_res(dt_boxes, img_path):
src_im = cv2.imread(img_path)
for box in dt_boxes:
box = np.array(box).astype(np.int32).reshape(-1, 2)
cv2.polylines(src_im, [box], True,
color=(255, 255, 0), thickness=2)
cv2.polylines(src_im, [box], True, color=(255, 255, 0), thickness=2)
return src_im
class DBPostProcess():
class DBPostProcess:
"""The post process for Differentiable Binarization (DB)."""
def __init__(self,
thresh=0.3,
box_thresh=0.7,
max_candidates=1000,
unclip_ratio=2.0,
score_mode="fast",
use_dilation=False):
def __init__(
self,
thresh=0.3,
box_thresh=0.7,
max_candidates=1000,
unclip_ratio=2.0,
score_mode="fast",
use_dilation=False,
):
self.thresh = thresh
self.box_thresh = box_thresh
self.max_candidates = max_candidates
@@ -268,16 +274,17 @@ class DBPostProcess():
self.dilation_kernel = None
def boxes_from_bitmap(self, pred, _bitmap, dest_width, dest_height):
'''
"""
_bitmap: single map with shape (1, H, W),
whose values are binarized as {0, 1}
'''
"""
bitmap = _bitmap
height, width = bitmap.shape
outs = cv2.findContours((bitmap * 255).astype(np.uint8), cv2.RETR_LIST,
cv2.CHAIN_APPROX_SIMPLE)
outs = cv2.findContours(
(bitmap * 255).astype(np.uint8), cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE
)
if len(outs) == 3:
img, contours, _ = outs[0], outs[1], outs[2]
elif len(outs) == 2:
@@ -306,10 +313,10 @@ class DBPostProcess():
continue
box = np.array(box)
box[:, 0] = np.clip(
np.round(box[:, 0] / width * dest_width), 0, dest_width)
box[:, 0] = np.clip(np.round(box[:, 0] / width * dest_width), 0, dest_width)
box[:, 1] = np.clip(
np.round(box[:, 1] / height * dest_height), 0, dest_height)
np.round(box[:, 1] / height * dest_height), 0, dest_height
)
boxes.append(box.astype(np.int16))
scores.append(score)
return np.array(boxes, dtype=np.int16), scores
@@ -341,9 +348,7 @@ class DBPostProcess():
index_2 = 3
index_3 = 2
box = [
points[index_1], points[index_2], points[index_3], points[index_4]
]
box = [points[index_1], points[index_2], points[index_3], points[index_4]]
return box, min(bounding_box[1])
def box_score_fast(self, bitmap, _box):
@@ -358,12 +363,12 @@ class DBPostProcess():
box[:, 0] = box[:, 0] - xmin
box[:, 1] = box[:, 1] - ymin
cv2.fillPoly(mask, box.reshape(1, -1, 2).astype(np.int32), 1)
return cv2.mean(bitmap[ymin:ymax + 1, xmin:xmax + 1], mask)[0]
return cv2.mean(bitmap[ymin : ymax + 1, xmin : xmax + 1], mask)[0]
def box_score_slow(self, bitmap, contour):
'''
"""
box_score_slow: use polyon mean score as the mean score
'''
"""
h, w = bitmap.shape[:2]
contour = contour.copy()
contour = np.reshape(contour, (-1, 2))
@@ -379,7 +384,7 @@ class DBPostProcess():
contour[:, 1] = contour[:, 1] - ymin
cv2.fillPoly(mask, contour.reshape(1, -1, 2).astype(np.int32), 1)
return cv2.mean(bitmap[ymin:ymax + 1, xmin:xmax + 1], mask)[0]
return cv2.mean(bitmap[ymin : ymax + 1, xmin : xmax + 1], mask)[0]
def __call__(self, pred, shape_list):
pred = pred[:, 0, :, :]
@@ -391,11 +396,13 @@ class DBPostProcess():
if self.dilation_kernel is not None:
mask = cv2.dilate(
np.array(segmentation[batch_index]).astype(np.uint8),
self.dilation_kernel)
self.dilation_kernel,
)
else:
mask = segmentation[batch_index]
boxes, scores = self.boxes_from_bitmap(pred[batch_index], mask,
src_w, src_h)
boxes, scores = self.boxes_from_bitmap(
pred[batch_index], mask, src_w, src_h
)
boxes_batch.append({'points': boxes})
boxes_batch.append({"points": boxes})
return boxes_batch
python/rapidocr_onnxruntime/ch_ppocr_v3_rec/text_recognize.py
+11 -13
View File
@@ -24,18 +24,18 @@ from rapidocr_onnxruntime.utils import OrtInferSession, read_yaml
from .utils import CTCLabelDecode
class TextRecognizer():
class TextRecognizer:
def __init__(self, config):
self.session = OrtInferSession(config)
if self.session.have_key():
self.character_dict_path = self.session.get_character_list()
else:
self.character_dict_path = config.get('keys_path', None)
self.character_dict_path = config.get("keys_path", None)
self.postprocess_op = CTCLabelDecode(self.character_dict_path)
self.rec_batch_num = config['rec_batch_num']
self.rec_image_shape = config['rec_img_shape']
self.rec_batch_num = config["rec_batch_num"]
self.rec_image_shape = config["rec_img_shape"]
def __call__(self, img_list: List[np.ndarray]):
if isinstance(img_list, np.ndarray):
@@ -48,7 +48,7 @@ class TextRecognizer():
indices = np.argsort(np.array(width_list))
img_num = len(img_list)
rec_res = [['', 0.0]] * img_num
rec_res = [["", 0.0]] * img_num
batch_num = self.rec_batch_num
elapse = 0
@@ -62,8 +62,7 @@ class TextRecognizer():
norm_img_batch = []
for ino in range(beg_img_no, end_img_no):
norm_img = self.resize_norm_img(img_list[indices[ino]],
max_wh_ratio)
norm_img = self.resize_norm_img(img_list[indices[ino]], max_wh_ratio)
norm_img_batch.append(norm_img[np.newaxis, :])
norm_img_batch = np.concatenate(norm_img_batch).astype(np.float32)
@@ -90,21 +89,20 @@ class TextRecognizer():
resized_w = int(math.ceil(img_height * ratio))
resized_image = cv2.resize(img, (resized_w, img_height))
resized_image = resized_image.astype('float32')
resized_image = resized_image.astype("float32")
resized_image = resized_image.transpose((2, 0, 1)) / 255
resized_image -= 0.5
resized_image /= 0.5
padding_im = np.zeros((img_channel, img_height, img_width),
dtype=np.float32)
padding_im = np.zeros((img_channel, img_height, img_width), dtype=np.float32)
padding_im[:, :, 0:resized_w] = resized_image
return padding_im
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--image_path', type=str, help='image_dir|image_path')
parser.add_argument('--config_path', type=str, default='config.yaml')
parser.add_argument("--image_path", type=str, help="image_dir|image_path")
parser.add_argument("--config_path", type=str, default="config.yaml")
args = parser.parse_args()
config = read_yaml(args.config_path)
@@ -112,4 +110,4 @@ if __name__ == "__main__":
img = cv2.imread(args.image_path)
rec_res, predict_time = text_recognizer(img)
print(f'rec result: {rec_res}\t cost: {predict_time}s')
print(f"rec result: {rec_res}\t cost: {predict_time}s")
python/rapidocr_onnxruntime/ch_ppocr_v3_rec/utils.py
+13 -13
View File
@@ -4,8 +4,8 @@
import numpy as np
class CTCLabelDecode():
""" Convert between text-label and text-index """
class CTCLabelDecode:
"""Convert between text-label and text-index"""
def __init__(self, character_dict_path):
super(CTCLabelDecode, self).__init__()
@@ -17,11 +17,11 @@ class CTCLabelDecode():
with open(character_dict_path, "rb") as fin:
lines = fin.readlines()
for line in lines:
line = line.decode('utf-8').strip("\n").strip("\r\n")
line = line.decode("utf-8").strip("\n").strip("\r\n")
self.character_str.append(line)
else:
self.character_str = character_dict_path
self.character_str.append(' ')
self.character_str.append(" ")
dict_character = self.add_special_char(self.character_str)
self.character = dict_character
@@ -33,22 +33,21 @@ class CTCLabelDecode():
def __call__(self, preds, label=None):
preds_idx = preds.argmax(axis=2)
preds_prob = preds.max(axis=2)
text = self.decode(preds_idx, preds_prob,
is_remove_duplicate=True)
text = self.decode(preds_idx, preds_prob, is_remove_duplicate=True)
if label is None:
return text
label = self.decode(label)
return text, label
def add_special_char(self, dict_character):
dict_character = ['blank'] + dict_character
dict_character = ["blank"] + dict_character
return dict_character
def get_ignored_tokens(self):
return [0] # for ctc blank
def decode(self, text_index, text_prob=None, is_remove_duplicate=False):
""" convert text-index into text-label. """
"""convert text-index into text-label."""
result_list = []
ignored_tokens = self.get_ignored_tokens()
@@ -61,15 +60,16 @@ class CTCLabelDecode():
continue
if is_remove_duplicate:
# only for predict
if idx > 0 and text_index[batch_idx][idx - 1] == text_index[
batch_idx][idx]:
if (
idx > 0
and text_index[batch_idx][idx - 1] == text_index[batch_idx][idx]
):
continue
char_list.append(self.character[int(text_index[batch_idx][
idx])])
char_list.append(self.character[int(text_index[batch_idx][idx])])
if text_prob is not None:
conf_list.append(text_prob[batch_idx][idx])
else:
conf_list.append(1)
text = ''.join(char_list)
text = "".join(char_list)
result_list.append((text, np.mean(conf_list + [1e-50])))
return result_list
python/rapidocr_onnxruntime/rapid_ocr_api.py
+50 -40
View File
@@ -12,17 +12,16 @@ import numpy as np
from .ch_ppocr_v2_cls import TextClassifier
from .ch_ppocr_v3_det import TextDetector
from .ch_ppocr_v3_rec import TextRecognizer
from .utils import (LoadImage, UpdateParameters, concat_model_path, init_args,
read_yaml)
from .utils import LoadImage, UpdateParameters, concat_model_path, init_args, read_yaml
root_dir = Path(__file__).resolve().parent
class RapidOCR():
class RapidOCR:
def __init__(self, **kwargs):
config_path = str(root_dir / 'config.yaml')
config_path = str(root_dir / "config.yaml")
if not Path(config_path).exists():
raise FileExistsError(f'{config_path} does not exist!')
raise FileExistsError(f"{config_path} does not exist!")
config = read_yaml(config_path)
config = concat_model_path(config)
@@ -30,30 +29,29 @@ class RapidOCR():
updater = UpdateParameters()
config = updater(config, **kwargs)
global_config = config['Global']
self.print_verbose = global_config['print_verbose']
self.text_score = global_config['text_score']
self.min_height = global_config['min_height']
self.width_height_ratio = global_config['width_height_ratio']
global_config = config["Global"]
self.print_verbose = global_config["print_verbose"]
self.text_score = global_config["text_score"]
self.min_height = global_config["min_height"]
self.width_height_ratio = global_config["width_height_ratio"]
self.use_text_det = config['Global']['use_text_det']
self.use_text_det = config["Global"]["use_text_det"]
if self.use_text_det:
self.text_detector = TextDetector(config['Det'])
self.text_detector = TextDetector(config["Det"])
self.text_recognizer = TextRecognizer(config['Rec'])
self.text_recognizer = TextRecognizer(config["Rec"])
self.use_angle_cls = config['Global']['use_angle_cls']
self.use_angle_cls = config["Global"]["use_angle_cls"]
if self.use_angle_cls:
self.text_cls = TextClassifier(config['Cls'])
self.text_cls = TextClassifier(config["Cls"])
self.load_img = LoadImage()
def __call__(self,
img_content: Union[str, np.ndarray, bytes, Path], **kwargs):
def __call__(self, img_content: Union[str, np.ndarray, bytes, Path], **kwargs):
if kwargs:
box_thresh = kwargs.get('box_thresh', 0.5)
unclip_ratio = kwargs.get('unclip_ratio', 1.6)
text_score = kwargs.get('text_score', 0.5)
box_thresh = kwargs.get("box_thresh", 0.5)
unclip_ratio = kwargs.get("unclip_ratio", 1.6)
text_score = kwargs.get("text_score", 0.5)
self.text_detector.postprocess_op.box_thresh = box_thresh
self.text_detector.postprocess_op.unclip_ratio = unclip_ratio
@@ -66,9 +64,7 @@ class RapidOCR():
else:
use_limit_ratio = w / h > self.width_height_ratio
if not self.use_text_det \
or h <= self.min_height \
or use_limit_ratio:
if not self.use_text_det or h <= self.min_height or use_limit_ratio:
dt_boxes, img_crop_list = self.get_boxes_img_without_det(img, h, w)
det_elapse = 0.0
else:
@@ -77,7 +73,7 @@ class RapidOCR():
return None, None
if self.print_verbose:
print(f'dt_boxes num: {len(dt_boxes)}, elapse: {det_elapse}')
print(f"dt_boxes num: {len(dt_boxes)}, elapse: {det_elapse}")
dt_boxes = self.sorted_boxes(dt_boxes)
img_crop_list = self.get_crop_img_list(img, dt_boxes)
@@ -87,16 +83,17 @@ class RapidOCR():
img_crop_list, _, cls_elapse = self.text_cls(img_crop_list)
if self.print_verbose:
print(f'cls num: {len(img_crop_list)}, elapse: {cls_elapse}')
print(f"cls num: {len(img_crop_list)}, elapse: {cls_elapse}")
rec_res, rec_elapse = self.text_recognizer(img_crop_list)
if self.print_verbose:
print(f'rec_res num: {len(rec_res)}, elapse: {rec_elapse}')
print(f"rec_res num: {len(rec_res)}, elapse: {rec_elapse}")
filter_boxes, filter_rec_res = self.filter_boxes_rec_by_score(dt_boxes,
rec_res)
fina_result = [[dt.tolist(), rec[0], str(rec[1])]
for dt, rec in zip(filter_boxes, filter_rec_res)]
filter_boxes, filter_rec_res = self.filter_boxes_rec_by_score(dt_boxes, rec_res)
fina_result = [
[dt.tolist(), rec[0], str(rec[1])]
for dt, rec in zip(filter_boxes, filter_rec_res)
]
if fina_result:
return fina_result, [det_elapse, cls_elapse, rec_elapse]
return None, None
@@ -118,20 +115,31 @@ class RapidOCR():
img_crop_width = int(
max(
np.linalg.norm(points[0] - points[1]),
np.linalg.norm(points[2] - points[3])))
np.linalg.norm(points[2] - points[3]),
)
)
img_crop_height = int(
max(
np.linalg.norm(points[0] - points[3]),
np.linalg.norm(points[1] - points[2])))
pts_std = np.float32([[0, 0], [img_crop_width, 0],
[img_crop_width, img_crop_height],
[0, img_crop_height]])
np.linalg.norm(points[1] - points[2]),
)
)
pts_std = np.float32(
[
[0, 0],
[img_crop_width, 0],
[img_crop_width, img_crop_height],
[0, img_crop_height],
]
)
M = cv2.getPerspectiveTransform(points, pts_std)
dst_img = cv2.warpPerspective(
img,
M, (img_crop_width, img_crop_height),
M,
(img_crop_width, img_crop_height),
borderMode=cv2.BORDER_REPLICATE,
flags=cv2.INTER_CUBIC)
flags=cv2.INTER_CUBIC,
)
dst_img_height, dst_img_width = dst_img.shape[0:2]
if dst_img_height * 1.0 / dst_img_width >= 1.5:
dst_img = np.rot90(dst_img)
@@ -159,8 +167,10 @@ class RapidOCR():
for i in range(num_boxes - 1):
for j in range(i, -1, -1):
if abs(_boxes[j + 1][0][1] - _boxes[j][0][1]) < 10 \
and _boxes[j + 1][0][0] < _boxes[j][0][0]:
if (
abs(_boxes[j + 1][0][1] - _boxes[j][0][1]) < 10
and _boxes[j + 1][0][0] < _boxes[j][0][0]
):
tmp = _boxes[j]
_boxes[j] = _boxes[j + 1]
_boxes[j + 1] = tmp
@@ -188,5 +198,5 @@ def main():
print(elapse_list)
if __name__ == '__main__':
if __name__ == "__main__":
main()
python/rapidocr_onnxruntime/utils.py
+109 -100
View File
@@ -10,69 +10,83 @@ from typing import Union
import cv2
import numpy as np
import yaml
from onnxruntime import (GraphOptimizationLevel, InferenceSession,
SessionOptions, get_available_providers, get_device)
from onnxruntime import (
GraphOptimizationLevel,
InferenceSession,
SessionOptions,
get_available_providers,
get_device,
)
from PIL import Image, UnidentifiedImageError
root_dir = Path(__file__).resolve().parent
InputType = Union[str, np.ndarray, bytes, Path]
class OrtInferSession():
class OrtInferSession:
def __init__(self, config):
sess_opt = SessionOptions()
sess_opt.log_severity_level = 4
sess_opt.enable_cpu_mem_arena = False
sess_opt.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_ALL
cpu_ep = 'CPUExecutionProvider'
cpu_ep = "CPUExecutionProvider"
cpu_provider_options = {
'arena_extend_strategy': 'kSameAsRequested',
"arena_extend_strategy": "kSameAsRequested",
}
cuda_ep = 'CUDAExecutionProvider'
cuda_ep = "CUDAExecutionProvider"
cuda_provider_options = {
'device_id': 0,
'arena_extend_strategy': 'kNextPowerOfTwo',
'cudnn_conv_algo_search': 'EXHAUSTIVE',
'do_copy_in_default_stream': True
"device_id": 0,
"arena_extend_strategy": "kNextPowerOfTwo",
"cudnn_conv_algo_search": "EXHAUSTIVE",
"do_copy_in_default_stream": True,
}
EP_list = []
if config['use_cuda'] and get_device() == 'GPU' \
and cuda_ep in get_available_providers():
if (
config["use_cuda"]
and get_device() == "GPU"
and cuda_ep in get_available_providers()
):
EP_list = [(cuda_ep, cuda_provider_options)]
EP_list.append((cpu_ep, cpu_provider_options))
self._verify_model(config['model_path'])
self.session = InferenceSession(config['model_path'],
sess_options=sess_opt,
providers=EP_list)
self._verify_model(config["model_path"])
self.session = InferenceSession(
config["model_path"], sess_options=sess_opt, providers=EP_list
)
if config['use_cuda'] and cuda_ep not in self.session.get_providers():
warnings.warn(f'{cuda_ep} is not avaiable for current env, the inference part is automatically shifted to be executed under {cpu_ep}.\n'
'Please ensure the installed onnxruntime-gpu version matches your cuda and cudnn version, '
'you can check their relations from the offical web site: '
'https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html',
RuntimeWarning)
if config["use_cuda"] and cuda_ep not in self.session.get_providers():
warnings.warn(
f"{cuda_ep} is not avaiable for current env, the inference part is automatically shifted to be executed under {cpu_ep}.\n"
"Please ensure the installed onnxruntime-gpu version matches your cuda and cudnn version, "
"you can check their relations from the offical web site: "
"https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html",
RuntimeWarning,
)
def __call__(self, input_content: np.ndarray) -> np.ndarray:
input_dict = dict(zip(self.get_input_names(), [input_content]))
try:
return self.session.run(self.get_output_names(), input_dict)
except Exception as e:
raise ONNXRuntimeError('ONNXRuntime inference failed.') from e
raise ONNXRuntimeError("ONNXRuntime inference failed.") from e
def get_input_names(self, ):
def get_input_names(
self,
):
return [v.name for v in self.session.get_inputs()]
def get_output_names(self,):
def get_output_names(
self,
):
return [v.name for v in self.session.get_outputs()]
def get_character_list(self, key: str = 'character'):
def get_character_list(self, key: str = "character"):
return self.meta_dict[key].splitlines()
def have_key(self, key: str = 'character') -> bool:
def have_key(self, key: str = "character") -> bool:
self.meta_dict = self.session.get_modelmeta().custom_metadata_map
if key in self.meta_dict.keys():
return True
@@ -82,23 +96,26 @@ class OrtInferSession():
def _verify_model(model_path):
model_path = Path(model_path)
if not model_path.exists():
raise FileNotFoundError(f'{model_path} does not exists.')
raise FileNotFoundError(f"{model_path} does not exists.")
if not model_path.is_file():
raise FileExistsError(f'{model_path} is not a file.')
raise FileExistsError(f"{model_path} is not a file.")
class ONNXRuntimeError(Exception):
pass
class LoadImage():
def __init__(self, ):
class LoadImage:
def __init__(
self,
):
pass
def __call__(self, img: InputType) -> np.ndarray:
if not isinstance(img, InputType.__args__):
raise LoadImageError(
f'The img type {type(img)} does not in {InputType.__args__}')
f"The img type {type(img)} does not in {InputType.__args__}"
)
img = self.load_img(img)
@@ -117,8 +134,7 @@ class LoadImage():
img = np.array(Image.open(img))
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
except UnidentifiedImageError as e:
raise LoadImageError(
f'cannot identify image file {img}') from e
raise LoadImageError(f"cannot identify image file {img}") from e
return img
if isinstance(img, bytes):
@@ -129,12 +145,11 @@ class LoadImage():
if isinstance(img, np.ndarray):
return img
raise LoadImageError(f'{type(img)} is not supported!')
raise LoadImageError(f"{type(img)} is not supported!")
@staticmethod
def cvt_four_to_three(img: np.ndarray) -> np.ndarray:
'''RGBA → RGB
'''
"""RGBA → RGB"""
r, g, b, a = cv2.split(img)
new_img = cv2.merge((b, g, r))
@@ -148,7 +163,7 @@ class LoadImage():
@staticmethod
def verify_exist(file_path: Union[str, Path]):
if not Path(file_path).exists():
raise LoadImageError(f'{file_path} does not exist.')
raise LoadImageError(f"{file_path} does not exist.")
class LoadImageError(Exception):
@@ -156,77 +171,74 @@ class LoadImageError(Exception):
def read_yaml(yaml_path):
with open(yaml_path, 'rb') as f:
with open(yaml_path, "rb") as f:
data = yaml.load(f, Loader=yaml.Loader)
return data
def concat_model_path(config):
key = 'model_path'
config['Det'][key] = str(root_dir / config['Det'][key])
config['Rec'][key] = str(root_dir / config['Rec'][key])
config['Cls'][key] = str(root_dir / config['Cls'][key])
key = "model_path"
config["Det"][key] = str(root_dir / config["Det"][key])
config["Rec"][key] = str(root_dir / config["Rec"][key])
config["Cls"][key] = str(root_dir / config["Cls"][key])
return config
def init_args():
parser = argparse.ArgumentParser()
parser.add_argument('-img', '--img_path', type=str, default=None,
required=True)
parser.add_argument('-p', '--print_cost',
action='store_true', default=False)
parser.add_argument("-img", "--img_path", type=str, default=None, required=True)
parser.add_argument("-p", "--print_cost", action="store_true", default=False)
global_group = parser.add_argument_group(title='Global')
global_group.add_argument('--text_score', type=float, default=0.5)
global_group.add_argument('--use_angle_cls', type=bool, default=True)
global_group.add_argument('--use_text_det', type=bool, default=True)
global_group.add_argument('--print_verbose', type=bool, default=False)
global_group.add_argument('--min_height', type=int, default=30)
global_group.add_argument('--width_height_ratio', type=int, default=8)
global_group = parser.add_argument_group(title="Global")
global_group.add_argument("--text_score", type=float, default=0.5)
global_group.add_argument("--use_angle_cls", type=bool, default=True)
global_group.add_argument("--use_text_det", type=bool, default=True)
global_group.add_argument("--print_verbose", type=bool, default=False)
global_group.add_argument("--min_height", type=int, default=30)
global_group.add_argument("--width_height_ratio", type=int, default=8)
det_group = parser.add_argument_group(title='Det')
det_group.add_argument('--det_model_path', type=str, default=None)
det_group.add_argument('--det_limit_side_len', type=float, default=736)
det_group.add_argument('--det_limit_type', type=str, default='min',
choices=['max', 'min'])
det_group.add_argument('--det_thresh', type=float, default=0.3)
det_group.add_argument('--det_box_thresh', type=float, default=0.5)
det_group.add_argument('--det_unclip_ratio', type=float, default=1.6)
det_group.add_argument('--det_use_dilation', type=bool, default=True)
det_group.add_argument('--det_score_mode', type=str, default='fast',
choices=['slow', 'fast'])
det_group = parser.add_argument_group(title="Det")
det_group.add_argument("--det_model_path", type=str, default=None)
det_group.add_argument("--det_limit_side_len", type=float, default=736)
det_group.add_argument(
"--det_limit_type", type=str, default="min", choices=["max", "min"]
)
det_group.add_argument("--det_thresh", type=float, default=0.3)
det_group.add_argument("--det_box_thresh", type=float, default=0.5)
det_group.add_argument("--det_unclip_ratio", type=float, default=1.6)
det_group.add_argument("--det_use_dilation", type=bool, default=True)
det_group.add_argument(
"--det_score_mode", type=str, default="fast", choices=["slow", "fast"]
)
cls_group = parser.add_argument_group(title='Cls')
cls_group.add_argument('--cls_model_path', type=str, default=None)
cls_group.add_argument('--cls_image_shape', type=list,
default=[3, 48, 192])
cls_group.add_argument('--cls_label_list', type=list,
default=['0', '180'])
cls_group.add_argument('--cls_batch_num', type=int, default=6)
cls_group.add_argument('--cls_thresh', type=float, default=0.9)
cls_group = parser.add_argument_group(title="Cls")
cls_group.add_argument("--cls_model_path", type=str, default=None)
cls_group.add_argument("--cls_image_shape", type=list, default=[3, 48, 192])
cls_group.add_argument("--cls_label_list", type=list, default=["0", "180"])
cls_group.add_argument("--cls_batch_num", type=int, default=6)
cls_group.add_argument("--cls_thresh", type=float, default=0.9)
rec_group = parser.add_argument_group(title='Rec')
rec_group.add_argument('--rec_model_path', type=str, default=None)
rec_group.add_argument('--rec_img_shape', type=list,
default=[3, 48, 320])
rec_group.add_argument('--rec_batch_num', type=int, default=6)
rec_group = parser.add_argument_group(title="Rec")
rec_group.add_argument("--rec_model_path", type=str, default=None)
rec_group.add_argument("--rec_img_shape", type=list, default=[3, 48, 320])
rec_group.add_argument("--rec_batch_num", type=int, default=6)
args = parser.parse_args()
return args
class UpdateParameters():
class UpdateParameters:
def __init__(self) -> None:
pass
def parse_kwargs(self, **kwargs):
global_dict, det_dict, cls_dict, rec_dict = {}, {}, {}, {}
for k, v in kwargs.items():
if k.startswith('det'):
if k.startswith("det"):
det_dict[k] = v
elif k.startswith('cls'):
elif k.startswith("cls"):
cls_dict[k] = v
elif k.startswith('rec'):
elif k.startswith("rec"):
rec_dict[k] = v
else:
global_dict[k] = v
@@ -235,11 +247,10 @@ class UpdateParameters():
def __call__(self, config, **kwargs):
global_dict, det_dict, cls_dict, rec_dict = self.parse_kwargs(**kwargs)
new_config = {
'Global': self.update_global_params(config['Global'],
global_dict),
'Det': self.update_det_params(config['Det'], det_dict),
'Cls': self.update_cls_params(config['Cls'], cls_dict),
'Rec': self.update_rec_params(config['Rec'], rec_dict)
"Global": self.update_global_params(config["Global"], global_dict),
"Det": self.update_det_params(config["Det"], det_dict),
"Cls": self.update_cls_params(config["Cls"], cls_dict),
"Rec": self.update_rec_params(config["Rec"], rec_dict),
}
return new_config
@@ -250,38 +261,36 @@ class UpdateParameters():
def update_det_params(self, config, det_dict):
if det_dict:
det_dict = {k.split('det_')[1]: v for k, v in det_dict.items()}
if not det_dict['model_path']:
det_dict['model_path'] = str(root_dir / config['model_path'])
det_dict = {k.split("det_")[1]: v for k, v in det_dict.items()}
if not det_dict["model_path"]:
det_dict["model_path"] = str(root_dir / config["model_path"])
config.update(det_dict)
return config
def update_cls_params(self, config, cls_dict):
if cls_dict:
need_remove_prefix = ['cls_label_list', 'cls_model_path']
need_remove_prefix = ["cls_label_list", "cls_model_path"]
new_cls_dict = {}
for k, v in cls_dict.items():
if k in need_remove_prefix:
k = k.split('cls_')[1]
k = k.split("cls_")[1]
new_cls_dict[k] = v
if not new_cls_dict['model_path']:
new_cls_dict['model_path'] = str(
root_dir / config['model_path'])
if not new_cls_dict["model_path"]:
new_cls_dict["model_path"] = str(root_dir / config["model_path"])
config.update(new_cls_dict)
return config
def update_rec_params(self, config, rec_dict):
if rec_dict:
need_remove_prefix = ['rec_model_path']
need_remove_prefix = ["rec_model_path"]
new_rec_dict = {}
for k, v in rec_dict.items():
if k in need_remove_prefix:
k = k.split('rec_')[1]
k = k.split("rec_")[1]
new_rec_dict[k] = v
if not new_rec_dict['model_path']:
new_rec_dict['model_path'] = str(
root_dir / config['model_path'])
if not new_rec_dict["model_path"]:
new_rec_dict["model_path"] = str(root_dir / config["model_path"])
config.update(new_rec_dict)
return config
python/rapidocr_openvino/ch_ppocr_v2_cls/text_cls.py
+12 -11
View File
@@ -25,12 +25,12 @@ from rapidocr_openvino.utils import OpenVINOInferSession, read_yaml
from .utils import ClsPostProcess
class TextClassifier():
class TextClassifier:
def __init__(self, config):
self.cls_image_shape = config['cls_image_shape']
self.cls_batch_num = config['cls_batch_num']
self.cls_thresh = config['cls_thresh']
self.postprocess_op = ClsPostProcess(config['label_list'])
self.cls_image_shape = config["cls_image_shape"]
self.cls_batch_num = config["cls_batch_num"]
self.cls_thresh = config["cls_thresh"]
self.postprocess_op = ClsPostProcess(config["label_list"])
self.infer = OpenVINOInferSession(config)
@@ -47,7 +47,7 @@ class TextClassifier():
indices = np.argsort(np.array(width_list))
img_num = len(img_list)
cls_res = [['', 0.0]] * img_num
cls_res = [["", 0.0]] * img_num
batch_num = self.cls_batch_num
elapse = 0
for beg_img_no in range(0, img_num, batch_num):
@@ -68,9 +68,10 @@ class TextClassifier():
for rno in range(len(cls_result)):
label, score = cls_result[rno]
cls_res[indices[beg_img_no + rno]] = [label, score]
if '180' in label and score > self.cls_thresh:
if "180" in label and score > self.cls_thresh:
img_list[indices[beg_img_no + rno]] = cv2.rotate(
img_list[indices[beg_img_no + rno]], 1)
img_list[indices[beg_img_no + rno]], 1
)
return img_list, cls_res, elapse
def resize_norm_img(self, img):
@@ -83,7 +84,7 @@ class TextClassifier():
resized_w = int(math.ceil(img_h * ratio))
resized_image = cv2.resize(img, (resized_w, img_h))
resized_image = resized_image.astype('float32')
resized_image = resized_image.astype("float32")
if img_c == 1:
resized_image = resized_image / 255
resized_image = resized_image[np.newaxis, :]
@@ -99,8 +100,8 @@ class TextClassifier():
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--image_path', type=str, help='image_dir|image_path')
parser.add_argument('--config_path', type=str, default='config.yaml')
parser.add_argument("--image_path", type=str, help="image_dir|image_path")
parser.add_argument("--config_path", type=str, default="config.yaml")
args = parser.parse_args()
config = read_yaml(args.config_path)
python/rapidocr_openvino/ch_ppocr_v2_cls/utils.py
+6 -5
View File
@@ -11,8 +11,8 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
class ClsPostProcess():
""" Convert between text-label and text-index """
class ClsPostProcess:
"""Convert between text-label and text-index"""
def __init__(self, label_list):
super(ClsPostProcess, self).__init__()
@@ -20,10 +20,11 @@ class ClsPostProcess():
def __call__(self, preds, label=None):
pred_idxs = preds.argmax(axis=1)
decode_out = [(self.label_list[idx], preds[i, idx])
for i, idx in enumerate(pred_idxs)]
decode_out = [
(self.label_list[idx], preds[i, idx]) for i, idx in enumerate(pred_idxs)
]
if label is None:
return decode_out
label = [(self.label_list[idx], 1.0) for idx in label]
return decode_out, label
return decode_out, label
python/rapidocr_openvino/ch_ppocr_v3_det/text_detect.py
+10 -9
View File
@@ -23,16 +23,16 @@ from rapidocr_openvino.utils import OpenVINOInferSession, read_yaml
from .utils import DBPostProcess, create_operators, transform
class TextDetector():
class TextDetector:
def __init__(self, config):
self.preprocess_op = create_operators(config['pre_process'])
self.postprocess_op = DBPostProcess(**config['post_process'])
self.preprocess_op = create_operators(config["pre_process"])
self.postprocess_op = DBPostProcess(**config["post_process"])
self.infer = OpenVINOInferSession(config)
def __call__(self, img):
ori_im = img.copy()
data = {'image': img}
data = {"image": img}
data = transform(data, self.preprocess_op)
img, shape_list = data
if img is None:
@@ -44,7 +44,7 @@ class TextDetector():
starttime = time.time()
preds = self.infer(img)
post_result = self.postprocess_op(preds, shape_list)
dt_boxes = post_result[0]['points']
dt_boxes = post_result[0]["points"]
dt_boxes = self.filter_tag_det_res(dt_boxes, ori_im.shape)
elapse = time.time() - starttime
return dt_boxes, elapse
@@ -96,8 +96,8 @@ class TextDetector():
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--config_path', type=str, default='config.yaml')
parser.add_argument('--image_path', type=str, default=None)
parser.add_argument("--config_path", type=str, default="config.yaml")
parser.add_argument("--image_path", type=str, default=None)
args = parser.parse_args()
config = read_yaml(args.config_path)
@@ -108,6 +108,7 @@ if __name__ == "__main__":
dt_boxes, elapse = text_detector(img)
from utils import draw_text_det_res
src_im = draw_text_det_res(dt_boxes, args.image_path)
cv2.imwrite('det_results.jpg', src_im)
print('The det_results.jpg has been saved in the current directory.')
cv2.imwrite("det_results.jpg", src_im)
print("The det_results.jpg has been saved in the current directory.")
python/rapidocr_openvino/ch_ppocr_v3_det/utils.py
+76 -69
View File
@@ -25,69 +25,74 @@ import six
from shapely.geometry import Polygon
class DecodeImage():
""" decode image """
class DecodeImage:
"""decode image"""
def __init__(self, img_mode='RGB', channel_first=False):
def __init__(self, img_mode="RGB", channel_first=False):
self.img_mode = img_mode
self.channel_first = channel_first
def __call__(self, data):
img = data['image']
img = data["image"]
if six.PY2:
assert type(img) is str and len(img) > 0, "invalid input 'img' in DecodeImage"
assert (
type(img) is str and len(img) > 0
), "invalid input 'img' in DecodeImage"
else:
assert type(img) is bytes and len(img) > 0, "invalid input 'img' in DecodeImage"
assert (
type(img) is bytes and len(img) > 0
), "invalid input 'img' in DecodeImage"
img = np.frombuffer(img, dtype='uint8')
img = np.frombuffer(img, dtype="uint8")
img = cv2.imdecode(img, 1)
if img is None:
return None
if self.img_mode == 'GRAY':
if self.img_mode == "GRAY":
img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
elif self.img_mode == 'RGB':
assert img.shape[2] == 3, f'invalid shape of image[{img.shape}]'
elif self.img_mode == "RGB":
assert img.shape[2] == 3, f"invalid shape of image[{img.shape}]"
img = img[:, :, ::-1]
if self.channel_first:
img = img.transpose((2, 0, 1))
data['image'] = img
data["image"] = img
return data
class NormalizeImage():
""" normalize image such as substract mean, divide std"""
class NormalizeImage:
"""normalize image such as substract mean, divide std"""
def __init__(self, scale=None, mean=None, std=None, order='chw'):
def __init__(self, scale=None, mean=None, std=None, order="chw"):
if isinstance(scale, str):
scale = eval(scale)
self.scale = np.float32(scale if scale is not None else 1.0 / 255.0)
mean = mean if mean is not None else [0.485, 0.456, 0.406]
std = std if std is not None else [0.229, 0.224, 0.225]
shape = (3, 1, 1) if order == 'chw' else (1, 1, 3)
self.mean = np.array(mean).reshape(shape).astype('float32')
self.std = np.array(std).reshape(shape).astype('float32')
shape = (3, 1, 1) if order == "chw" else (1, 1, 3)
self.mean = np.array(mean).reshape(shape).astype("float32")
self.std = np.array(std).reshape(shape).astype("float32")
def __call__(self, data):
img = np.array(data['image']).astype(np.float32)
data['image'] = (img * self.scale - self.mean) / self.std
img = np.array(data["image"]).astype(np.float32)
data["image"] = (img * self.scale - self.mean) / self.std
return data
class ToCHWImage():
""" convert hwc image to chw image"""
class ToCHWImage:
"""convert hwc image to chw image"""
def __init__(self):
pass
def __call__(self, data):
img = np.array(data['image'])
data['image'] = img.transpose((2, 0, 1))
img = np.array(data["image"])
data["image"] = img.transpose((2, 0, 1))
return data
class KeepKeys():
class KeepKeys:
def __init__(self, keep_keys):
self.keep_keys = keep_keys
@@ -98,26 +103,26 @@ class KeepKeys():
return data_list
class DetResizeForTest():
class DetResizeForTest:
def __init__(self, **kwargs):
super(DetResizeForTest, self).__init__()
self.resize_type = 0
if 'image_shape' in kwargs:
self.image_shape = kwargs['image_shape']
if "image_shape" in kwargs:
self.image_shape = kwargs["image_shape"]
self.resize_type = 1
elif 'limit_side_len' in kwargs:
self.limit_side_len = kwargs.get('limit_side_len', 736)
self.limit_type = kwargs.get('limit_type', 'min')
elif "limit_side_len" in kwargs:
self.limit_side_len = kwargs.get("limit_side_len", 736)
self.limit_type = kwargs.get("limit_type", "min")
if 'resize_long' in kwargs:
if "resize_long" in kwargs:
self.resize_type = 2
self.resize_long = kwargs.get('resize_long', 960)
self.resize_long = kwargs.get("resize_long", 960)
else:
self.limit_side_len = kwargs.get('limit_side_len', 736)
self.limit_type = kwargs.get('limit_type', 'min')
self.limit_side_len = kwargs.get("limit_side_len", 736)
self.limit_type = kwargs.get("limit_type", "min")
def __call__(self, data):
img = data['image']
img = data["image"]
src_h, src_w = img.shape[:2]
if self.resize_type == 0:
@@ -128,8 +133,8 @@ class DetResizeForTest():
else:
# img, shape = self.resize_image_type1(img)
img, [ratio_h, ratio_w] = self.resize_image_type1(img)
data['image'] = img
data['shape'] = np.array([src_h, src_w, ratio_h, ratio_w])
data["image"] = img
data["shape"] = np.array([src_h, src_w, ratio_h, ratio_w])
return data
def resize_image_type1(self, img):
@@ -153,14 +158,14 @@ class DetResizeForTest():
h, w = img.shape[:2]
# limit the max side
if self.limit_type == 'max':
if self.limit_type == "max":
if max(h, w) > limit_side_len:
if h > w:
ratio = float(limit_side_len) / h
else:
ratio = float(limit_side_len) / w
else:
ratio = 1.
ratio = 1.0
else:
if min(h, w) < limit_side_len:
if h < w:
@@ -168,7 +173,7 @@ class DetResizeForTest():
else:
ratio = float(limit_side_len) / w
else:
ratio = 1.
ratio = 1.0
resize_h = int(h * ratio)
resize_w = int(w * ratio)
@@ -212,7 +217,7 @@ class DetResizeForTest():
def transform(data, ops=None):
""" transform """
"""transform"""
if ops is None:
ops = []
for op in ops:
@@ -239,21 +244,22 @@ def draw_text_det_res(dt_boxes, img_path):
src_im = cv2.imread(img_path)
for box in dt_boxes:
box = np.array(box).astype(np.int32).reshape(-1, 2)
cv2.polylines(src_im, [box], True,
color=(255, 255, 0), thickness=2)
cv2.polylines(src_im, [box], True, color=(255, 255, 0), thickness=2)
return src_im
class DBPostProcess():
class DBPostProcess:
"""The post process for Differentiable Binarization (DB)."""
def __init__(self,
thresh=0.3,
box_thresh=0.7,
max_candidates=1000,
unclip_ratio=2.0,
score_mode="fast",
use_dilation=False):
def __init__(
self,
thresh=0.3,
box_thresh=0.7,
max_candidates=1000,
unclip_ratio=2.0,
score_mode="fast",
use_dilation=False,
):
self.thresh = thresh
self.box_thresh = box_thresh
self.max_candidates = max_candidates
@@ -267,16 +273,17 @@ class DBPostProcess():
self.dilation_kernel = None
def boxes_from_bitmap(self, pred, _bitmap, dest_width, dest_height):
'''
"""
_bitmap: single map with shape (1, H, W),
whose values are binarized as {0, 1}
'''
"""
bitmap = _bitmap
height, width = bitmap.shape
outs = cv2.findContours((bitmap * 255).astype(np.uint8), cv2.RETR_LIST,
cv2.CHAIN_APPROX_SIMPLE)
outs = cv2.findContours(
(bitmap * 255).astype(np.uint8), cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE
)
if len(outs) == 3:
img, contours, _ = outs[0], outs[1], outs[2]
elif len(outs) == 2:
@@ -305,10 +312,10 @@ class DBPostProcess():
continue
box = np.array(box)
box[:, 0] = np.clip(
np.round(box[:, 0] / width * dest_width), 0, dest_width)
box[:, 0] = np.clip(np.round(box[:, 0] / width * dest_width), 0, dest_width)
box[:, 1] = np.clip(
np.round(box[:, 1] / height * dest_height), 0, dest_height)
np.round(box[:, 1] / height * dest_height), 0, dest_height
)
boxes.append(box.astype(np.int16))
scores.append(score)
return np.array(boxes, dtype=np.int16), scores
@@ -340,9 +347,7 @@ class DBPostProcess():
index_2 = 3
index_3 = 2
box = [
points[index_1], points[index_2], points[index_3], points[index_4]
]
box = [points[index_1], points[index_2], points[index_3], points[index_4]]
return box, min(bounding_box[1])
def box_score_fast(self, bitmap, _box):
@@ -357,12 +362,12 @@ class DBPostProcess():
box[:, 0] = box[:, 0] - xmin
box[:, 1] = box[:, 1] - ymin
cv2.fillPoly(mask, box.reshape(1, -1, 2).astype(np.int32), 1)
return cv2.mean(bitmap[ymin:ymax + 1, xmin:xmax + 1], mask)[0]
return cv2.mean(bitmap[ymin : ymax + 1, xmin : xmax + 1], mask)[0]
def box_score_slow(self, bitmap, contour):
'''
"""
box_score_slow: use polyon mean score as the mean score
'''
"""
h, w = bitmap.shape[:2]
contour = contour.copy()
contour = np.reshape(contour, (-1, 2))
@@ -378,7 +383,7 @@ class DBPostProcess():
contour[:, 1] = contour[:, 1] - ymin
cv2.fillPoly(mask, contour.reshape(1, -1, 2).astype(np.int32), 1)
return cv2.mean(bitmap[ymin:ymax + 1, xmin:xmax + 1], mask)[0]
return cv2.mean(bitmap[ymin : ymax + 1, xmin : xmax + 1], mask)[0]
def __call__(self, pred, shape_list):
pred = pred[:, 0, :, :]
@@ -390,11 +395,13 @@ class DBPostProcess():
if self.dilation_kernel is not None:
mask = cv2.dilate(
np.array(segmentation[batch_index]).astype(np.uint8),
self.dilation_kernel)
self.dilation_kernel,
)
else:
mask = segmentation[batch_index]
boxes, scores = self.boxes_from_bitmap(pred[batch_index], mask,
src_w, src_h)
boxes, scores = self.boxes_from_bitmap(
pred[batch_index], mask, src_w, src_h
)
boxes_batch.append({'points': boxes})
boxes_batch.append({"points": boxes})
return boxes_batch
python/rapidocr_openvino/ch_ppocr_v3_rec/text_recognize.py
+12 -14
View File
@@ -28,13 +28,13 @@ from rapidocr_openvino.utils import OpenVINOInferSession, read_yaml
from .utils import CTCLabelDecode
class TextRecognizer():
class TextRecognizer:
def __init__(self, config):
self.rec_image_shape = config['rec_img_shape']
self.rec_batch_num = config['rec_batch_num']
self.rec_image_shape = config["rec_img_shape"]
self.rec_batch_num = config["rec_batch_num"]
dict_path = str(Path(__file__).parent / 'ppocr_keys_v1.txt')
self.character_dict_path = config.get('keys_path', dict_path)
dict_path = str(Path(__file__).parent / "ppocr_keys_v1.txt")
self.character_dict_path = config.get("keys_path", dict_path)
self.postprocess_op = CTCLabelDecode(self.character_dict_path)
self.infer = OpenVINOInferSession(config)
@@ -50,7 +50,7 @@ class TextRecognizer():
indices = np.argsort(np.array(width_list))
img_num = len(img_list)
rec_res = [['', 0.0]] * img_num
rec_res = [["", 0.0]] * img_num
batch_num = self.rec_batch_num
elapse = 0
@@ -64,8 +64,7 @@ class TextRecognizer():
norm_img_batch = []
for ino in range(beg_img_no, end_img_no):
norm_img = self.resize_norm_img(img_list[indices[ino]],
max_wh_ratio)
norm_img = self.resize_norm_img(img_list[indices[ino]], max_wh_ratio)
norm_img_batch.append(norm_img[np.newaxis, :])
norm_img_batch = np.concatenate(norm_img_batch).astype(np.float32)
@@ -91,21 +90,20 @@ class TextRecognizer():
resized_w = int(math.ceil(img_height * ratio))
resized_image = cv2.resize(img, (resized_w, img_height))
resized_image = resized_image.astype('float32')
resized_image = resized_image.astype("float32")
resized_image = resized_image.transpose((2, 0, 1)) / 255
resized_image -= 0.5
resized_image /= 0.5
padding_im = np.zeros((img_channel, img_height, img_width),
dtype=np.float32)
padding_im = np.zeros((img_channel, img_height, img_width), dtype=np.float32)
padding_im[:, :, 0:resized_w] = resized_image
return padding_im
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--image_path', type=str, help='image_dir|image_path')
parser.add_argument('--config_path', type=str, default='config.yaml')
parser.add_argument("--image_path", type=str, help="image_dir|image_path")
parser.add_argument("--config_path", type=str, default="config.yaml")
args = parser.parse_args()
config = read_yaml(args.config_path)
@@ -113,4 +111,4 @@ if __name__ == "__main__":
img = cv2.imread(args.image_path)
rec_res, predict_time = text_recognizer(img)
print(f'rec result: {rec_res}\t cost: {predict_time}s')
print(f"rec result: {rec_res}\t cost: {predict_time}s")
python/rapidocr_openvino/ch_ppocr_v3_rec/utils.py
+13 -13
View File
@@ -4,8 +4,8 @@
import numpy as np
class CTCLabelDecode():
""" Convert between text-label and text-index """
class CTCLabelDecode:
"""Convert between text-label and text-index"""
def __init__(self, character_dict_path):
super(CTCLabelDecode, self).__init__()
@@ -15,9 +15,9 @@ class CTCLabelDecode():
with open(character_dict_path, "rb") as fin:
lines = fin.readlines()
for line in lines:
line = line.decode('utf-8').strip("\n").strip("\r\n")
line = line.decode("utf-8").strip("\n").strip("\r\n")
self.character_str.append(line)
self.character_str.append(' ')
self.character_str.append(" ")
dict_character = self.add_special_char(self.character_str)
self.character = dict_character
@@ -29,22 +29,21 @@ class CTCLabelDecode():
def __call__(self, preds, label=None):
preds_idx = preds.argmax(axis=2)
preds_prob = preds.max(axis=2)
text = self.decode(preds_idx, preds_prob,
is_remove_duplicate=True)
text = self.decode(preds_idx, preds_prob, is_remove_duplicate=True)
if label is None:
return text
label = self.decode(label)
return text, label
def add_special_char(self, dict_character):
dict_character = ['blank'] + dict_character
dict_character = ["blank"] + dict_character
return dict_character
def get_ignored_tokens(self):
return [0] # for ctc blank
def decode(self, text_index, text_prob=None, is_remove_duplicate=False):
""" convert text-index into text-label. """
"""convert text-index into text-label."""
result_list = []
ignored_tokens = self.get_ignored_tokens()
@@ -57,15 +56,16 @@ class CTCLabelDecode():
continue
if is_remove_duplicate:
# only for predict
if idx > 0 and text_index[batch_idx][idx - 1] == text_index[
batch_idx][idx]:
if (
idx > 0
and text_index[batch_idx][idx - 1] == text_index[batch_idx][idx]
):
continue
char_list.append(self.character[int(text_index[batch_idx][
idx])])
char_list.append(self.character[int(text_index[batch_idx][idx])])
if text_prob is not None:
conf_list.append(text_prob[batch_idx][idx])
else:
conf_list.append(1)
text = ''.join(char_list)
text = "".join(char_list)
result_list.append((text, np.mean(conf_list + [1e-50])))
return result_list
python/rapidocr_openvino/rapid_ocr_api.py
+50 -40
View File
@@ -12,17 +12,16 @@ import numpy as np
from .ch_ppocr_v2_cls import TextClassifier
from .ch_ppocr_v3_det import TextDetector
from .ch_ppocr_v3_rec import TextRecognizer
from .utils import (LoadImage, UpdateParameters, concat_model_path, init_args,
read_yaml)
from .utils import LoadImage, UpdateParameters, concat_model_path, init_args, read_yaml
root_dir = Path(__file__).resolve().parent
class RapidOCR():
class RapidOCR:
def __init__(self, **kwargs):
config_path = str(root_dir / 'config.yaml')
config_path = str(root_dir / "config.yaml")
if not Path(config_path).exists():
raise FileExistsError(f'{config_path} does not exist!')
raise FileExistsError(f"{config_path} does not exist!")
config = read_yaml(config_path)
config = concat_model_path(config)
@@ -30,30 +29,29 @@ class RapidOCR():
updater = UpdateParameters()
config = updater(config, **kwargs)
global_config = config['Global']
self.print_verbose = global_config['print_verbose']
self.text_score = global_config['text_score']
self.min_height = global_config['min_height']
self.width_height_ratio = global_config['width_height_ratio']
global_config = config["Global"]
self.print_verbose = global_config["print_verbose"]
self.text_score = global_config["text_score"]
self.min_height = global_config["min_height"]
self.width_height_ratio = global_config["width_height_ratio"]
self.use_text_det = config['Global']['use_text_det']
self.use_text_det = config["Global"]["use_text_det"]
if self.use_text_det:
self.text_detector = TextDetector(config['Det'])
self.text_detector = TextDetector(config["Det"])
self.text_recognizer = TextRecognizer(config['Rec'])
self.text_recognizer = TextRecognizer(config["Rec"])
self.use_angle_cls = config['Global']['use_angle_cls']
self.use_angle_cls = config["Global"]["use_angle_cls"]
if self.use_angle_cls:
self.text_cls = TextClassifier(config['Cls'])
self.text_cls = TextClassifier(config["Cls"])
self.load_img = LoadImage()
def __call__(self,
img_content: Union[str, np.ndarray, bytes, Path], **kwargs):
def __call__(self, img_content: Union[str, np.ndarray, bytes, Path], **kwargs):
if kwargs:
box_thresh = kwargs.get('box_thresh', 0.5)
unclip_ratio = kwargs.get('unclip_ratio', 1.6)
text_score = kwargs.get('text_score', 0.5)
box_thresh = kwargs.get("box_thresh", 0.5)
unclip_ratio = kwargs.get("unclip_ratio", 1.6)
text_score = kwargs.get("text_score", 0.5)
self.text_detector.postprocess_op.box_thresh = box_thresh
self.text_detector.postprocess_op.unclip_ratio = unclip_ratio
@@ -66,9 +64,7 @@ class RapidOCR():
else:
use_limit_ratio = w / h > self.width_height_ratio
if not self.use_text_det \
or h <= self.min_height \
or use_limit_ratio:
if not self.use_text_det or h <= self.min_height or use_limit_ratio:
dt_boxes, img_crop_list = self.get_boxes_img_without_det(img, h, w)
det_elapse = 0.0
else:
@@ -77,7 +73,7 @@ class RapidOCR():
return None, None
if self.print_verbose:
print(f'dt_boxes num: {len(dt_boxes)}, elapse: {det_elapse}')
print(f"dt_boxes num: {len(dt_boxes)}, elapse: {det_elapse}")
dt_boxes = self.sorted_boxes(dt_boxes)
img_crop_list = self.get_crop_img_list(img, dt_boxes)
@@ -87,16 +83,17 @@ class RapidOCR():
img_crop_list, _, cls_elapse = self.text_cls(img_crop_list)
if self.print_verbose:
print(f'cls num: {len(img_crop_list)}, elapse: {cls_elapse}')
print(f"cls num: {len(img_crop_list)}, elapse: {cls_elapse}")
rec_res, rec_elapse = self.text_recognizer(img_crop_list)
if self.print_verbose:
print(f'rec_res num: {len(rec_res)}, elapse: {rec_elapse}')
print(f"rec_res num: {len(rec_res)}, elapse: {rec_elapse}")
filter_boxes, filter_rec_res = self.filter_boxes_rec_by_score(dt_boxes,
rec_res)
fina_result = [[dt.tolist(), rec[0], str(rec[1])]
for dt, rec in zip(filter_boxes, filter_rec_res)]
filter_boxes, filter_rec_res = self.filter_boxes_rec_by_score(dt_boxes, rec_res)
fina_result = [
[dt.tolist(), rec[0], str(rec[1])]
for dt, rec in zip(filter_boxes, filter_rec_res)
]
if fina_result:
return fina_result, [det_elapse, cls_elapse, rec_elapse]
return None, None
@@ -118,20 +115,31 @@ class RapidOCR():
img_crop_width = int(
max(
np.linalg.norm(points[0] - points[1]),
np.linalg.norm(points[2] - points[3])))
np.linalg.norm(points[2] - points[3]),
)
)
img_crop_height = int(
max(
np.linalg.norm(points[0] - points[3]),
np.linalg.norm(points[1] - points[2])))
pts_std = np.float32([[0, 0], [img_crop_width, 0],
[img_crop_width, img_crop_height],
[0, img_crop_height]])
np.linalg.norm(points[1] - points[2]),
)
)
pts_std = np.float32(
[
[0, 0],
[img_crop_width, 0],
[img_crop_width, img_crop_height],
[0, img_crop_height],
]
)
M = cv2.getPerspectiveTransform(points, pts_std)
dst_img = cv2.warpPerspective(
img,
M, (img_crop_width, img_crop_height),
M,
(img_crop_width, img_crop_height),
borderMode=cv2.BORDER_REPLICATE,
flags=cv2.INTER_CUBIC)
flags=cv2.INTER_CUBIC,
)
dst_img_height, dst_img_width = dst_img.shape[0:2]
if dst_img_height * 1.0 / dst_img_width >= 1.5:
dst_img = np.rot90(dst_img)
@@ -159,8 +167,10 @@ class RapidOCR():
for i in range(num_boxes - 1):
for j in range(i, -1, -1):
if abs(_boxes[j + 1][0][1] - _boxes[j][0][1]) < 10 \
and _boxes[j + 1][0][0] < _boxes[j][0][0]:
if (
abs(_boxes[j + 1][0][1] - _boxes[j][0][1]) < 10
and _boxes[j + 1][0][0] < _boxes[j][0][0]
):
tmp = _boxes[j]
_boxes[j] = _boxes[j + 1]
_boxes[j + 1] = tmp
@@ -188,5 +198,5 @@ def main():
print(elapse_list)
if __name__ == '__main__':
if __name__ == "__main__":
main()
python/rapidocr_openvino/utils.py
+73 -78
View File
@@ -16,14 +16,14 @@ root_dir = Path(__file__).resolve().parent
InputType = Union[str, np.ndarray, bytes, Path]
class OpenVINOInferSession():
class OpenVINOInferSession:
def __init__(self, config):
ie = Core()
config['model_path'] = str(root_dir / config['model_path'])
self._verify_model(config['model_path'])
model_onnx = ie.read_model(config['model_path'])
compile_model = ie.compile_model(model=model_onnx, device_name='CPU')
config["model_path"] = str(root_dir / config["model_path"])
self._verify_model(config["model_path"])
model_onnx = ie.read_model(config["model_path"])
compile_model = ie.compile_model(model=model_onnx, device_name="CPU")
self.session = compile_model.create_infer_request()
def __call__(self, input_content: np.ndarray) -> np.ndarray:
@@ -34,19 +34,22 @@ class OpenVINOInferSession():
def _verify_model(model_path):
model_path = Path(model_path)
if not model_path.exists():
raise FileNotFoundError(f'{model_path} does not exists.')
raise FileNotFoundError(f"{model_path} does not exists.")
if not model_path.is_file():
raise FileExistsError(f'{model_path} is not a file.')
raise FileExistsError(f"{model_path} is not a file.")
class LoadImage():
def __init__(self, ):
class LoadImage:
def __init__(
self,
):
pass
def __call__(self, img: InputType) -> np.ndarray:
if not isinstance(img, InputType.__args__):
raise LoadImageError(
f'The img type {type(img)} does not in {InputType.__args__}')
f"The img type {type(img)} does not in {InputType.__args__}"
)
img = self.load_img(img)
@@ -65,8 +68,7 @@ class LoadImage():
img = np.array(Image.open(img))
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
except UnidentifiedImageError as e:
raise LoadImageError(
f'cannot identify image file {img}') from e
raise LoadImageError(f"cannot identify image file {img}") from e
return img
if isinstance(img, bytes):
@@ -77,12 +79,11 @@ class LoadImage():
if isinstance(img, np.ndarray):
return img
raise LoadImageError(f'{type(img)} is not supported!')
raise LoadImageError(f"{type(img)} is not supported!")
@staticmethod
def cvt_four_to_three(img: np.ndarray) -> np.ndarray:
'''RGBA → RGB
'''
"""RGBA → RGB"""
r, g, b, a = cv2.split(img)
new_img = cv2.merge((b, g, r))
@@ -96,7 +97,7 @@ class LoadImage():
@staticmethod
def verify_exist(file_path: Union[str, Path]):
if not Path(file_path).exists():
raise LoadImageError(f'{file_path} does not exist.')
raise LoadImageError(f"{file_path} does not exist.")
class LoadImageError(Exception):
@@ -104,77 +105,74 @@ class LoadImageError(Exception):
def read_yaml(yaml_path):
with open(yaml_path, 'rb') as f:
with open(yaml_path, "rb") as f:
data = yaml.load(f, Loader=yaml.Loader)
return data
def concat_model_path(config):
key = 'model_path'
config['Det'][key] = str(root_dir / config['Det'][key])
config['Rec'][key] = str(root_dir / config['Rec'][key])
config['Cls'][key] = str(root_dir / config['Cls'][key])
key = "model_path"
config["Det"][key] = str(root_dir / config["Det"][key])
config["Rec"][key] = str(root_dir / config["Rec"][key])
config["Cls"][key] = str(root_dir / config["Cls"][key])
return config
def init_args():
parser = argparse.ArgumentParser()
parser.add_argument('-img', '--img_path', type=str, default=None,
required=True)
parser.add_argument('-p', '--print_cost',
action='store_true', default=False)
parser.add_argument("-img", "--img_path", type=str, default=None, required=True)
parser.add_argument("-p", "--print_cost", action="store_true", default=False)
global_group = parser.add_argument_group(title='Global')
global_group.add_argument('--text_score', type=float, default=0.5)
global_group.add_argument('--use_angle_cls', type=bool, default=True)
global_group.add_argument('--use_text_det', type=bool, default=True)
global_group.add_argument('--print_verbose', type=bool, default=False)
global_group.add_argument('--min_height', type=int, default=30)
global_group.add_argument('--width_height_ratio', type=int, default=8)
global_group = parser.add_argument_group(title="Global")
global_group.add_argument("--text_score", type=float, default=0.5)
global_group.add_argument("--use_angle_cls", type=bool, default=True)
global_group.add_argument("--use_text_det", type=bool, default=True)
global_group.add_argument("--print_verbose", type=bool, default=False)
global_group.add_argument("--min_height", type=int, default=30)
global_group.add_argument("--width_height_ratio", type=int, default=8)
det_group = parser.add_argument_group(title='Det')
det_group.add_argument('--det_model_path', type=str, default=None)
det_group.add_argument('--det_limit_side_len', type=float, default=736)
det_group.add_argument('--det_limit_type', type=str, default='min',
choices=['max', 'min'])
det_group.add_argument('--det_thresh', type=float, default=0.3)
det_group.add_argument('--det_box_thresh', type=float, default=0.5)
det_group.add_argument('--det_unclip_ratio', type=float, default=1.6)
det_group.add_argument('--det_use_dilation', type=bool, default=True)
det_group.add_argument('--det_score_mode', type=str, default='fast',
choices=['slow', 'fast'])
det_group = parser.add_argument_group(title="Det")
det_group.add_argument("--det_model_path", type=str, default=None)
det_group.add_argument("--det_limit_side_len", type=float, default=736)
det_group.add_argument(
"--det_limit_type", type=str, default="min", choices=["max", "min"]
)
det_group.add_argument("--det_thresh", type=float, default=0.3)
det_group.add_argument("--det_box_thresh", type=float, default=0.5)
det_group.add_argument("--det_unclip_ratio", type=float, default=1.6)
det_group.add_argument("--det_use_dilation", type=bool, default=True)
det_group.add_argument(
"--det_score_mode", type=str, default="fast", choices=["slow", "fast"]
)
cls_group = parser.add_argument_group(title='Cls')
cls_group.add_argument('--cls_model_path', type=str, default=None)
cls_group.add_argument('--cls_image_shape', type=list,
default=[3, 48, 192])
cls_group.add_argument('--cls_label_list', type=list,
default=['0', '180'])
cls_group.add_argument('--cls_batch_num', type=int, default=6)
cls_group.add_argument('--cls_thresh', type=float, default=0.9)
cls_group = parser.add_argument_group(title="Cls")
cls_group.add_argument("--cls_model_path", type=str, default=None)
cls_group.add_argument("--cls_image_shape", type=list, default=[3, 48, 192])
cls_group.add_argument("--cls_label_list", type=list, default=["0", "180"])
cls_group.add_argument("--cls_batch_num", type=int, default=6)
cls_group.add_argument("--cls_thresh", type=float, default=0.9)
rec_group = parser.add_argument_group(title='Rec')
rec_group.add_argument('--rec_model_path', type=str, default=None)
rec_group.add_argument('--rec_image_shape', type=list,
default=[3, 48, 320])
rec_group.add_argument('--rec_batch_num', type=int, default=6)
rec_group = parser.add_argument_group(title="Rec")
rec_group.add_argument("--rec_model_path", type=str, default=None)
rec_group.add_argument("--rec_image_shape", type=list, default=[3, 48, 320])
rec_group.add_argument("--rec_batch_num", type=int, default=6)
args = parser.parse_args()
return args
class UpdateParameters():
class UpdateParameters:
def __init__(self) -> None:
pass
def parse_kwargs(self, **kwargs):
global_dict, det_dict, cls_dict, rec_dict = {}, {}, {}, {}
for k, v in kwargs.items():
if k.startswith('det'):
if k.startswith("det"):
det_dict[k] = v
elif k.startswith('cls'):
elif k.startswith("cls"):
cls_dict[k] = v
elif k.startswith('rec'):
elif k.startswith("rec"):
rec_dict[k] = v
else:
global_dict[k] = v
@@ -183,11 +181,10 @@ class UpdateParameters():
def __call__(self, config, **kwargs):
global_dict, det_dict, cls_dict, rec_dict = self.parse_kwargs(**kwargs)
new_config = {
'Global': self.update_global_params(config['Global'],
global_dict),
'Det': self.update_det_params(config['Det'], det_dict),
'Cls': self.update_cls_params(config['Cls'], cls_dict),
'Rec': self.update_rec_params(config['Rec'], rec_dict)
"Global": self.update_global_params(config["Global"], global_dict),
"Det": self.update_det_params(config["Det"], det_dict),
"Cls": self.update_cls_params(config["Cls"], cls_dict),
"Rec": self.update_rec_params(config["Rec"], rec_dict),
}
return new_config
@@ -198,38 +195,36 @@ class UpdateParameters():
def update_det_params(self, config, det_dict):
if det_dict:
det_dict = {k.split('det_')[1]: v for k, v in det_dict.items()}
if not det_dict['model_path']:
det_dict['model_path'] = str(root_dir / config['model_path'])
det_dict = {k.split("det_")[1]: v for k, v in det_dict.items()}
if not det_dict["model_path"]:
det_dict["model_path"] = str(root_dir / config["model_path"])
config.update(det_dict)
return config
def update_cls_params(self, config, cls_dict):
if cls_dict:
need_remove_prefix = ['cls_label_list', 'cls_model_path']
need_remove_prefix = ["cls_label_list", "cls_model_path"]
new_cls_dict = {}
for k, v in cls_dict.items():
if k in need_remove_prefix:
k = k.split('cls_')[1]
k = k.split("cls_")[1]
new_cls_dict[k] = v
if not new_cls_dict['model_path']:
new_cls_dict['model_path'] = str(
root_dir / config['model_path'])
if not new_cls_dict["model_path"]:
new_cls_dict["model_path"] = str(root_dir / config["model_path"])
config.update(new_cls_dict)
return config
def update_rec_params(self, config, rec_dict):
if rec_dict:
need_remove_prefix = ['rec_model_path']
need_remove_prefix = ["rec_model_path"]
new_rec_dict = {}
for k, v in rec_dict.items():
if k in need_remove_prefix:
k = k.split('rec_')[1]
k = k.split("rec_")[1]
new_rec_dict[k] = v
if not new_rec_dict['model_path']:
new_rec_dict['model_path'] = str(
root_dir / config['model_path'])
if not new_rec_dict["model_path"]:
new_rec_dict["model_path"] = str(root_dir / config["model_path"])
config.update(new_rec_dict)
return config
python/setup_onnxruntime.py
+27 -20
View File
@@ -10,14 +10,14 @@ from get_pypi_latest_version import GetPyPiLatestVersion
def get_readme():
root_dir = Path(__file__).resolve().parent.parent
readme_path = str(root_dir / 'docs' / 'doc_whl_rapidocr_ort.md')
readme_path = str(root_dir / "docs" / "doc_whl_rapidocr_ort.md")
print(readme_path)
with open(readme_path, 'r', encoding='utf-8') as f:
with open(readme_path, "r", encoding="utf-8") as f:
readme = f.read()
return readme
MODULE_NAME = 'rapidocr_onnxruntime'
MODULE_NAME = "rapidocr_onnxruntime"
obtainer = GetPyPiLatestVersion()
latest_version = obtainer(MODULE_NAME)
@@ -25,7 +25,7 @@ VERSION_NUM = obtainer.version_add_one(latest_version)
# 优先提取commit message中的语义化版本号，如无，则自动加1
if len(sys.argv) > 2:
match_str = ' '.join(sys.argv[2:])
match_str = " ".join(sys.argv[2:])
matched_versions = obtainer.extract_version(match_str)
if matched_versions:
VERSION_NUM = matched_versions
@@ -37,31 +37,38 @@ setuptools.setup(
platforms="Any",
description="A cross platform OCR Library based on OnnxRuntime.",
long_description=get_readme(),
long_description_content_type='text/markdown',
long_description_content_type="text/markdown",
author="SWHL",
author_email="liekkaskono@163.com",
url="https://github.com/RapidAI/RapidOCR",
license='Apache-2.0',
license="Apache-2.0",
include_package_data=True,
install_requires=[
"pyclipper>=1.2.1", "onnxruntime>=1.7.0", "opencv_python>=4.5.1.48",
"numpy>=1.19.3", "six>=1.15.0", "Shapely>=1.7.1", 'PyYAML', 'Pillow'
"pyclipper>=1.2.1",
"onnxruntime>=1.7.0",
"opencv_python>=4.5.1.48",
"numpy>=1.19.3",
"six>=1.15.0",
"Shapely>=1.7.1",
"PyYAML",
"Pillow",
],
package_dir={'': MODULE_NAME},
package_dir={"": MODULE_NAME},
packages=setuptools.find_namespace_packages(where=MODULE_NAME),
package_data={'': ['*.onnx', '*.yaml']},
package_data={"": ["*.onnx", "*.yaml"]},
keywords=[
'ocr,text_detection,text_recognition,db,onnxruntime,paddleocr,openvino,rapidocr'
"ocr,text_detection,text_recognition,db,onnxruntime,paddleocr,openvino,rapidocr"
],
classifiers=[
'Programming Language :: Python :: 3.6',
'Programming Language :: Python :: 3.7',
'Programming Language :: Python :: 3.8',
'Programming Language :: Python :: 3.9',
'Programming Language :: Python :: 3.10',
'Programming Language :: Python :: 3.11',
"Programming Language :: Python :: 3.6",
"Programming Language :: Python :: 3.7",
"Programming Language :: Python :: 3.8",
"Programming Language :: Python :: 3.9",
"Programming Language :: Python :: 3.10",
"Programming Language :: Python :: 3.11",
],
python_requires='>=3.6,<3.12',
python_requires=">=3.6,<3.12",
entry_points={
'console_scripts': [f'{MODULE_NAME}={MODULE_NAME}.rapid_ocr_api:main'],
})
"console_scripts": [f"{MODULE_NAME}={MODULE_NAME}.rapid_ocr_api:main"],
},
)
python/setup_openvino.py
+28 -21
View File
@@ -10,14 +10,14 @@ from get_pypi_latest_version import GetPyPiLatestVersion
def get_readme():
root_dir = Path(__file__).resolve().parent.parent
readme_path = str(root_dir / 'docs' / 'doc_whl_rapidocr_vino.md')
readme_path = str(root_dir / "docs" / "doc_whl_rapidocr_vino.md")
print(readme_path)
with open(readme_path, 'r', encoding='utf-8') as f:
with open(readme_path, "r", encoding="utf-8") as f:
readme = f.read()
return readme
MODULE_NAME = 'rapidocr_openvino'
MODULE_NAME = "rapidocr_openvino"
obtainer = GetPyPiLatestVersion()
latest_version = obtainer(MODULE_NAME)
@@ -25,7 +25,7 @@ VERSION_NUM = obtainer.version_add_one(latest_version)
# 优先提取commit message中的语义化版本号，如无，则自动加1
if len(sys.argv) > 2:
match_str = ' '.join(sys.argv[2:])
match_str = " ".join(sys.argv[2:])
matched_versions = obtainer.extract_version(match_str)
if matched_versions:
VERSION_NUM = matched_versions
@@ -37,31 +37,38 @@ setuptools.setup(
platforms="Any",
description="A cross platform OCR Library based on OpenVINO.",
long_description=get_readme(),
long_description_content_type='text/markdown',
long_description_content_type="text/markdown",
author="SWHL",
author_email="liekkaskono@163.com",
url="https://github.com/RapidAI/RapidOCR",
license='Apache-2.0',
license="Apache-2.0",
include_package_data=True,
install_requires=["pyclipper>=1.2.1", "openvino>=2022.2.0",
"opencv_python>=4.5.1.48", "numpy>=1.19.3",
"six>=1.15.0", "Shapely>=1.7.1", 'PyYAML', 'Pillow'],
package_dir={'': MODULE_NAME},
install_requires=[
"pyclipper>=1.2.1",
"openvino>=2022.2.0",
"opencv_python>=4.5.1.48",
"numpy>=1.19.3",
"six>=1.15.0",
"Shapely>=1.7.1",
"PyYAML",
"Pillow",
],
package_dir={"": MODULE_NAME},
packages=setuptools.find_namespace_packages(where=MODULE_NAME),
package_data={'': ['*.onnx', '*.yaml', '*.txt']},
package_data={"": ["*.onnx", "*.yaml", "*.txt"]},
keywords=[
'ocr,text_detection,text_recognition,db,onnxruntime,paddleocr,openvino,rapidocr'
"ocr,text_detection,text_recognition,db,onnxruntime,paddleocr,openvino,rapidocr"
],
classifiers=[
'Programming Language :: Python :: 3.6',
'Programming Language :: Python :: 3.7',
'Programming Language :: Python :: 3.8',
'Programming Language :: Python :: 3.9',
'Programming Language :: Python :: 3.10',
'Programming Language :: Python :: 3.11',
"Programming Language :: Python :: 3.6",
"Programming Language :: Python :: 3.7",
"Programming Language :: Python :: 3.8",
"Programming Language :: Python :: 3.9",
"Programming Language :: Python :: 3.10",
"Programming Language :: Python :: 3.11",
],
python_requires='>=3.6,<3.12',
python_requires=">=3.6,<3.12",
entry_points={
'console_scripts': [f'{MODULE_NAME}={MODULE_NAME}.rapid_ocr_api:main'],
}
"console_scripts": [f"{MODULE_NAME}={MODULE_NAME}.rapid_ocr_api:main"],
},
)
python/tests/base_module.py
+6 -6
View File
@@ -8,25 +8,25 @@ from pathlib import Path
import yaml
class BaseModule():
def __init__(self, package_name: str = 'rapidocr_onnxruntime'):
class BaseModule:
def __init__(self, package_name: str = "rapidocr_onnxruntime"):
self.package_name = package_name
self.root_dir = Path(__file__).resolve().parent.parent
self.package_dir = self.root_dir / self.package_name
self.tests_dir = self.root_dir / 'tests'
self.tests_dir = self.root_dir / "tests"
sys.path.append(str(self.root_dir))
sys.path.append(str(self.package_dir))
def init_module(self, module_name: str, class_name: str = None):
if class_name is None:
module_part = importlib.import_module(f'{self.package_name}')
module_part = importlib.import_module(f"{self.package_name}")
return module_part
module_part = importlib.import_module(f'{self.package_name}.{module_name}')
module_part = importlib.import_module(f"{self.package_name}.{module_name}")
return getattr(module_part, class_name)
@staticmethod
def read_yaml(yaml_path: str):
with open(yaml_path, 'rb') as f:
with open(yaml_path, "rb") as f:
data = yaml.load(f, Loader=yaml.Loader)
return data
python/tests/benchmark/benchmark.py
+10 -8
View File
@@ -14,22 +14,22 @@ import cv2
from rapidocr_onnxruntime import RapidOCR
if __name__ == '__main__':
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--yaml_path', type=str, default='config.yaml')
parser.add_argument("--yaml_path", type=str, default="config.yaml")
args = parser.parse_args()
yaml_path = cur_dir / args.yaml_path
rapid_ocr = RapidOCR(yaml_path)
image_dir = cur_dir / 'test_images_benchmark'
image_dir = cur_dir / "test_images_benchmark"
if not image_dir.exists():
raise FileNotFoundError(f'{image_dir} does not exits!!')
raise FileNotFoundError(f"{image_dir} does not exits!!")
image_list = list(image_dir.iterdir())
cost_time_list = []
for image_path in tqdm(image_list, desc='Test'):
for image_path in tqdm(image_list, desc="Test"):
img = cv2.imread(str(image_path))
start_time = time.time()
@@ -40,6 +40,8 @@ if __name__ == '__main__':
total_time = sum(cost_time_list)
avg_time = total_time / len(cost_time_list)
print(f'Total Files: {len(image_list)}, '
f'Total Time: {total_time:.5f}, '
f'Average Time: {avg_time:.5f}')
print(
f"Total Files: {len(image_list)}, "
f"Total Time: {total_time:.5f}, "
f"Average Time: {avg_time:.5f}"
)
python/tests/test_all_ort.py
+17 -17
View File
@@ -16,14 +16,14 @@ sys.path.append(str(root_dir))
from rapidocr_onnxruntime import RapidOCR, LoadImageError
rapid_ocr = RapidOCR()
tests_dir = root_dir / 'tests' / 'test_files'
tests_dir = root_dir / "tests" / "test_files"
def test_normal():
image_path = tests_dir / 'ch_en_num.jpg'
image_path = tests_dir / "ch_en_num.jpg"
img = cv2.imread(str(image_path))
result, _ = rapid_ocr(img)
assert result[0][1] == '正品促销'
assert result[0][1] == "正品促销"
assert len(result) == 17
@@ -42,29 +42,29 @@ def test_zeros():
def test_input_str():
image_path = tests_dir / 'ch_en_num.jpg'
image_path = tests_dir / "ch_en_num.jpg"
result, _ = rapid_ocr(str(image_path))
assert result[0][1] == '正品促销'
assert result[0][1] == "正品促销"
assert len(result) == 17
def test_input_bytes():
image_path = tests_dir / 'ch_en_num.jpg'
with open(image_path, 'rb') as f:
image_path = tests_dir / "ch_en_num.jpg"
with open(image_path, "rb") as f:
result, _ = rapid_ocr(f.read())
assert result[0][1] == '正品促销'
assert result[0][1] == "正品促销"
assert len(result) == 17
def test_input_path():
image_path = tests_dir / 'ch_en_num.jpg'
image_path = tests_dir / "ch_en_num.jpg"
result, _ = rapid_ocr(image_path)
assert result[0][1] == '正品促销'
assert result[0][1] == "正品促销"
assert len(result) == 17
def test_input_parameters():
image_path = tests_dir / 'ch_en_num.jpg'
image_path = tests_dir / "ch_en_num.jpg"
rapid_ocr = RapidOCR(text_score=1)
result, _ = rapid_ocr(image_path)
@@ -72,27 +72,27 @@ def test_input_parameters():
def test_input_det_parameters():
image_path = tests_dir / 'ch_en_num.jpg'
image_path = tests_dir / "ch_en_num.jpg"
with pytest.raises(FileNotFoundError) as exc_info:
rapid_ocr = RapidOCR(det_model_path='1.onnx')
rapid_ocr = RapidOCR(det_model_path="1.onnx")
result, _ = rapid_ocr(image_path)
raise FileNotFoundError()
assert exc_info.type is FileNotFoundError
def test_input_cls_parameters():
image_path = tests_dir / 'ch_en_num.jpg'
image_path = tests_dir / "ch_en_num.jpg"
with pytest.raises(FileNotFoundError) as exc_info:
rapid_ocr = RapidOCR(cls_model_path='1.onnx')
rapid_ocr = RapidOCR(cls_model_path="1.onnx")
result, _ = rapid_ocr(image_path)
raise FileNotFoundError()
assert exc_info.type is FileNotFoundError
def test_input_rec_parameters():
image_path = tests_dir / 'ch_en_num.jpg'
image_path = tests_dir / "ch_en_num.jpg"
with pytest.raises(FileNotFoundError) as exc_info:
rapid_ocr = RapidOCR(rec_model_path='1.onnx')
rapid_ocr = RapidOCR(rec_model_path="1.onnx")
result, _ = rapid_ocr(image_path)
raise FileNotFoundError()
assert exc_info.type is FileNotFoundError
python/tests/test_all_vino.py
+17 -17
View File
@@ -16,14 +16,14 @@ from rapidocr_openvino import RapidOCR, LoadImageError
rapid_ocr = RapidOCR()
tests_dir = root_dir / 'tests' / 'test_files'
tests_dir = root_dir / "tests" / "test_files"
def test_normal():
image_path = tests_dir / 'ch_en_num.jpg'
image_path = tests_dir / "ch_en_num.jpg"
img = cv2.imread(str(image_path))
result, _ = rapid_ocr(img)
assert result[0][1] == '正品促销'
assert result[0][1] == "正品促销"
assert len(result) == 17
@@ -42,29 +42,29 @@ def test_zeros():
def test_input_str():
image_path = tests_dir / 'ch_en_num.jpg'
image_path = tests_dir / "ch_en_num.jpg"
result, _ = rapid_ocr(str(image_path))
assert result[0][1] == '正品促销'
assert result[0][1] == "正品促销"
assert len(result) == 17
def test_input_bytes():
image_path = tests_dir / 'ch_en_num.jpg'
with open(image_path, 'rb') as f:
image_path = tests_dir / "ch_en_num.jpg"
with open(image_path, "rb") as f:
result, _ = rapid_ocr(f.read())
assert result[0][1] == '正品促销'
assert result[0][1] == "正品促销"
assert len(result) == 17
def test_input_path():
image_path = tests_dir / 'ch_en_num.jpg'
image_path = tests_dir / "ch_en_num.jpg"
result, _ = rapid_ocr(image_path)
assert result[0][1] == '正品促销'
assert result[0][1] == "正品促销"
assert len(result) == 17
def test_input_parameters():
image_path = tests_dir / 'ch_en_num.jpg'
image_path = tests_dir / "ch_en_num.jpg"
rapid_ocr = RapidOCR(text_score=1)
result, _ = rapid_ocr(image_path)
@@ -72,27 +72,27 @@ def test_input_parameters():
def test_input_det_parameters():
image_path = tests_dir / 'ch_en_num.jpg'
image_path = tests_dir / "ch_en_num.jpg"
with pytest.raises(FileNotFoundError) as exc_info:
rapid_ocr = RapidOCR(det_model_path='1.onnx')
rapid_ocr = RapidOCR(det_model_path="1.onnx")
result, _ = rapid_ocr(image_path)
raise FileNotFoundError()
assert exc_info.type is FileNotFoundError
def test_input_cls_parameters():
image_path = tests_dir / 'ch_en_num.jpg'
image_path = tests_dir / "ch_en_num.jpg"
with pytest.raises(FileNotFoundError) as exc_info:
rapid_ocr = RapidOCR(cls_model_path='1.onnx')
rapid_ocr = RapidOCR(cls_model_path="1.onnx")
result, _ = rapid_ocr(image_path)
raise FileNotFoundError()
assert exc_info.type is FileNotFoundError
def test_input_rec_parameters():
image_path = tests_dir / 'ch_en_num.jpg'
image_path = tests_dir / "ch_en_num.jpg"
with pytest.raises(FileNotFoundError) as exc_info:
rapid_ocr = RapidOCR(rec_model_path='1.onnx')
rapid_ocr = RapidOCR(rec_model_path="1.onnx")
result, _ = rapid_ocr(image_path)
raise FileNotFoundError()
assert exc_info.type is FileNotFoundError
python/tests/test_cls.py
+7 -9
View File
@@ -7,24 +7,22 @@ from base_module import BaseModule
@pytest.mark.parametrize(
'package_name',
[('rapidocr_onnxruntime'),
('rapidocr_openvino')]
"package_name", [("rapidocr_onnxruntime"), ("rapidocr_openvino")]
)
def test_cls(package_name: str):
module_name = 'ch_ppocr_v2_cls'
class_name = 'TextClassifier'
module_name = "ch_ppocr_v2_cls"
class_name = "TextClassifier"
base = BaseModule(package_name=package_name)
TextClassifier = base.init_module(module_name, class_name)
yaml_path = base.package_dir / module_name / 'config.yaml'
yaml_path = base.package_dir / module_name / "config.yaml"
config = base.read_yaml(str(yaml_path))
config['model_path'] = str(base.package_dir / config['model_path'])
config["model_path"] = str(base.package_dir / config["model_path"])
text_cls = TextClassifier(config)
img_path = base.tests_dir / 'test_files' / 'text_cls.jpg'
img_path = base.tests_dir / "test_files" / "text_cls.jpg"
img = cv2.imread(str(img_path))
result = text_cls([img])
assert result[1][0][0] == '180'
assert result[1][0][0] == "180"
python/tests/test_det.py
+6 -9
View File
@@ -6,23 +6,20 @@ import pytest
from base_module import BaseModule
@pytest.mark.parametrize(
'package_name',
['rapidocr_onnxruntime', 'rapidocr_openvino']
)
@pytest.mark.parametrize("package_name", ["rapidocr_onnxruntime", "rapidocr_openvino"])
def test_det(package_name):
module_name = 'ch_ppocr_v3_det'
class_name = 'TextDetector'
module_name = "ch_ppocr_v3_det"
class_name = "TextDetector"
base = BaseModule(package_name)
TextDetector = base.init_module(module_name, class_name)
yaml_path = base.package_dir / module_name / 'config.yaml'
yaml_path = base.package_dir / module_name / "config.yaml"
config = base.read_yaml(str(yaml_path))
config['model_path'] = str(base.package_dir / config['model_path'])
config["model_path"] = str(base.package_dir / config["model_path"])
text_det = TextDetector(config)
img_path = base.tests_dir / 'test_files' / 'text_det.jpg'
img_path = base.tests_dir / "test_files" / "text_det.jpg"
img = cv2.imread(str(img_path))
dt_boxes, elapse = text_det(img)
assert dt_boxes.shape == (18, 4, 2)
python/tests/test_rec.py
+7 -10
View File
@@ -6,24 +6,21 @@ import pytest
from base_module import BaseModule
@pytest.mark.parametrize(
'package_name',
['rapidocr_onnxruntime', 'rapidocr_openvino']
)
@pytest.mark.parametrize("package_name", ["rapidocr_onnxruntime", "rapidocr_openvino"])
def test_det(package_name):
module_name = 'ch_ppocr_v3_rec'
class_name = 'TextRecognizer'
module_name = "ch_ppocr_v3_rec"
class_name = "TextRecognizer"
base = BaseModule(package_name)
TextRecognizer = base.init_module(module_name, class_name)
yaml_path = base.package_dir / module_name / 'config.yaml'
yaml_path = base.package_dir / module_name / "config.yaml"
config = base.read_yaml(str(yaml_path))
config['model_path'] = str(base.package_dir / config['model_path'])
config["model_path"] = str(base.package_dir / config["model_path"])
text_rec = TextRecognizer(config)
img_path = base.tests_dir / 'test_files' / 'text_rec.jpg'
img_path = base.tests_dir / "test_files" / "text_rec.jpg"
img = cv2.imread(str(img_path))
rec_res, elapse = text_rec([img])
assert rec_res[0][0] == '韩国小馆'
assert rec_res[0][0] == "韩国小馆"