import argparse
import os

from paddleocr import PaddleOCR, draw_ocr


def rec_result(img_dir, save_dir):
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    # 模型路径下必须含有model和params文件
    ocr = PaddleOCR(det_model_dir='model/det/tw_idcard_det_infer_1202', rec_model_dir='model/rec/ch_PP-OCRv3_rec_infer', rec_char_dict_path='ppocr_keys_v1.txt', cls_model_dir='model/cls/ch_ppocr_mobile_v2.0_cls_slim_infer', use_angle_cls=True)
    path_list = os.listdir(img_dir)
    for img_p in path_list:
        img_path = img_dir + '/' + img_p
        result = ocr.ocr(img_path, cls=True)
        # for line in result:
        #     print(line)

        # 显示结果
        from PIL import Image
        image = Image.open(img_path).convert('RGB')
        boxes = [line[0] for line in result]
        txts = [line[1][0] for line in result]
        scores = [line[1][1] for line in result]
        im_show = draw_ocr(image, boxes, txts, scores, font_path='font/simfang.ttf')
        im_show = Image.fromarray(im_show)
        im_show.save(save_dir + '/' + img_p)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--img_dir', type=str, default='./imgs')
    parser.add_argument('--save_dir', type=str, default='./result')
    args = parser.parse_args()
    rec_result(args.img_dir, args.save_dir)