玩玩paddle ocr
paddle ocr demo封装
前言
对于ocr一直比较感兴趣，随便看了下好像训练开源的比较少，百度还是比较良心的，提供了保姆级的教程给开发者，但是发现整个项目好像没有提供一个可以修改的脚本，处于好玩的目的，clone了下paddle ocr，简单的写了个paddle ocr demo，记录一下，
1.配置文件
百度的ocr有三个模型：检测、识别、分类，原始的仓库上对齐有较为详细的说明。这里简单写到yaml配置文件里面
image_dir: ./imgs/11.jpg
det: true
rec: true
use_angle_cls: false
drop_score: 0.5
# DB parmas
det_algorithm:
name: DB
use_gpu: true
model_dir: 2.0_model/det
pre_process:
DetResizeForTest:
limit_side_len: 960
limit_type: max
NormalizeImage:
std: [0.229, 0.224, 0.225]
mean: [0.485, 0.456, 0.406]
scale: 1./255.
order: hwc
ToCHWImage:
KeepKeys:
keep_keys: ['image', 'shape']
postprocess:
name: DBPostProcess
thresh: 0.4
box_thresh: 0.5
max_candidates: 1000
unclip_ratio: 2.0
use_dilation: True
lang:
name: ch
dict_path: ./ppocr/utils/ppocr_keys_v1.txt
rec_algorithm:
name: CRNN
use_gpu: true
model_dir: 2.0_model/rec/ch
rec_image_shape: "3, 32, 320"
rec_char_type: ch
rec_batch_num: 30
max_text_length: 25
postprocess:
name: CTCLabelDecode
character_dict_path: ./dict/ppocr_keys_v1.txt
use_space_char: True
text_classifier:
use_gpu: true
model_dir: 2.0_model/cls
cls_image_shape: 3, 48, 192
label_list: ['0', '180']
cls_batch_num: 30
cls_thresh: 0.9
postprocess:
name: ClsPostProcess
2.类提取和封装
demo里面所有类实现均参考原始的仓库
class PaddleOCR(TextSystem):
def __init__(self, args):
"""
paddleocr package
args:
**kwargs: other params show in paddleocr --help
"""
self.use_angle_cls = args['use_angle_cls']
super().__init__(args)
def ocr(self, img, det=True, rec=True, cls=False):
"""
ocr with paddleocr
args：
img: img for ocr, support ndarray, img_path and list or ndarray
det: use text detection or not, if false, only rec will be exec. default is True
rec: use text recognition or not, if false, only det will be exec. default is True
"""
if isinstance(img, np.ndarray) and len(img.shape) == 2:
img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
if det and rec:
dt_boxes, rec_res = self.__call__(img)
return [[box.tolist(), res] for box, res in zip(dt_boxes, rec_res)]
elif det and not rec:
dt_boxes, elapse = self.text_detector(img)
if dt_boxes is None:
return None
return [box.tolist() for box in dt_boxes]
else:
if not isinstance(img, list):
img = [img]
if self.use_angle_cls:
img, cls_res, elapse = self.text_classifier(img)
if not rec:
return cls_res
rec_res, elapse = self.text_recognizer(img)
return rec_res
3添加demo可视化
import argparse
import yaml
import cv2
import numpy as np
from PIL import ImageFont, ImageDraw, Image
from core.PaddleOCR import PaddleOCR
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--params", type=str, default='data/params.yaml')
return parser.parse_args()
def draw_result(img, result):
img_rec = np.ones_like(img, np.uint8)*255
img_pil = Image.fromarray(img_rec)
draw = ImageDraw.Draw(img_pil)
fontpath = "font/simsun.ttc"
font = ImageFont.truetype(fontpath, 16)
for info in result:
bbox, rec_info = info
pts=np.array(bbox, np.int32)
pts=pts.reshape((-1,1,2))
cv2.polylines(img,[pts],True,(0,255,0),2)
txt = rec_info[0] + str(rec_info[1])
draw.text(tuple(pts[0][0]), txt, font=font, fill =(0,255,0))
bk_img = np.array(img_pil)
draw_img = np.hstack([img,bk_img])
return draw_img
if __name__ == '__main__':
args = parse_args()
with open(args.params) as f:
data_dict = yaml.load(f, Loader=yaml.FullLoader) # data dict
ocr_engine = PaddleOCR(data_dict)
img = cv2.imread(data_dict['image_dir'])
result = ocr_engine.ocr(img,
det=data_dict['det'],
rec=data_dict['rec'],
cls=data_dict['use_angle_cls'])
draw_img = draw_result(img, result)
cv2.imwrite('result.jpg', draw_img)
cv2.imshow("img", draw_img)
cv2.waitKey(0)
最后结果输出
