DocumentAIKit/ocr_func_test/get_ocr_boxes.py

import datetime
import os.path
import uuid

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


# 例如`ch`, `en`, `fr`, `german`, `korean`, `japan`
from tqdm import tqdm


def get_boxes(img_path, language):
    ocr = PaddleOCR(use_angle_cls=True, lang=language)  # need to run only once to download and load model into memory
    result = ocr.ocr(img_path, cls=True)
    # for line in idcard_imgs:
    #     print(line)
    #
    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('result_lisa_ch.jpg')
    return [line[0] for line in result]


def image_cut_save(img_path, boxes, save_dir, cnt):
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    img = cv2.imread(img_path)  # 打开图像
    for box in tqdm(boxes):
        upper = int(box[0][1])
        lower = int(box[2][1])
        left = int(box[0][0])
        right = int(box[2][0])
        cropped = img[upper:lower, left:right]
        time = datetime.datetime.now()
        t = str(time.year).zfill(4) + '-' + str(time.month).zfill(2) + '-' + str(time.day).zfill(2)
        save_path = save_dir + '/' + 'kdan' + '_' + t + '_' + str(cnt).zfill(4) + '_' + str(uuid.uuid1())[0:8] + '.jpg'
        cv2.imwrite(save_path, cropped)
        cnt += 1


def get_all_img(img_dir):
    result = []
    img_paths = os.listdir(img_dir)
    for it in img_paths:
        temp = os.path.join(img_dir, it)
        result.append(temp)
    return result


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--img_path', type=str, default='idcard_lisa.png')
    parser.add_argument('--img_dir', type=str, default='')
    parser.add_argument('--language', type=str, default='chinese_cht')
    parser.add_argument('--save_dir', type=str, default='./idcard_imgs')
    args = parser.parse_args()
    cnt = 1
    if args.img_dir == '':
        boxes = get_boxes(args.img_path, args.language)
        # image_cut_save(args.img_path, boxes, args.save_dir, cnt)
    else:
        img_paths = get_all_img(args.img_dir)
        for img_path in tqdm(img_paths):
            boxes = get_boxes(img_path, args.language)
            image_cut_save(img_path, boxes, args.save_dir, cnt)