import os
from paddleocr import PaddleOCR, draw_ocr
import argparse
from PIL import Image


def get_boxes(img_path, language, save_dir):
    ocr = PaddleOCR(use_angle_cls=True, rec_model_dir='model/rec/tw_idcard_rec_infer_0103/Student', lang=language)  # need to run only once to download and load model into memory
    result = ocr.ocr(img_path, cls=True)
    for line in result:
        # 左上
        line[0][0][0] -= 5
        line[0][0][1] -= 5
        # 右上
        line[0][1][0] += 5
        line[0][1][1] -= 5
        # 右下
        line[0][2][0] += 5
        line[0][2][1] += 5
        # 左下
        line[0][3][0] -= 5
        line[0][3][1] += 5
    # 显示结果
    image = Image.open(img_path).convert('RGB')
    boxes = [line[0] for line in result]
    texts = [line[1][0] for line in result]
    scores = [line[1][1] for line in result]
    im_show = draw_ocr(image, boxes, texts, scores, font_path='./fonts/STSONG.TTF')
    im_show = Image.fromarray(im_show)
    im_show.save(save_dir + "/" + language + "_" + str(img_path).split("/")[-1])


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--img_dir', type=str, default='imgs')
    parser.add_argument('--img_path', type=str, default='')
    parser.add_argument('--language', type=str, default='ch')
    parser.add_argument('--save_dir', type=str, default='./result')
    args = parser.parse_args()
    if args.img_path == '':
        path_list = os.listdir(args.img_dir)
        for path in path_list:
            get_boxes(args.img_dir + '/' + path, args.language, args.save_dir)
    else:
        get_boxes(args.img_path, args.language, args.save_dir)