# -*- coding: utf-8 -*-

"""

@author:XuMing(xuming624@qq.com)

@description: Cli for image ocr.

"""

import argparse

import csv

import os

import sys

from glob import glob

from loguru import logger

from tqdm import tqdm

sys.path.append('..')

from imgocr.ppocr_onnx import ImgOcr, draw_ocr_boxes

def save_partial_results(rows, output_file, is_first_chunk):

"""Save rows (list of [image, ocr_result]) to CSV file."""

mode = 'w' if is_first_chunk else 'a'

file_dir = os.path.dirname(output_file)

if file_dir and not os.path.exists(file_dir):

os.makedirs(file_dir)

with open(output_file, mode, encoding='utf-8', newline='') as f:

writer = csv.writer(f)

if is_first_chunk:

writer.writerow(['images', 'ocr_results'])

writer.writerows(rows)

def parse_args():

parser = argparse.ArgumentParser(description='imgocr cli')

parser.add_argument('--image_dir', type=str, help='input image dir path, required', required=True)

parser.add_argument('--output_dir', type=str, default='outputs', help='output ocr result dir path, default outputs')

parser.add_argument("--save_box_img", action='store_true', help='save ocr box img, default False')

parser.add_argument("--no_efficiency_mode", action='store_true', help='disable efficiency mode, default False')

parser.add_argument('--chunk_size', type=int, default=50, help='chunk size')

parser.add_argument('--use_gpu', action='store_true', help='use gpu, default False')

args = parser.parse_args()

logger.debug(args)

return args

def cli(args):

m = ImgOcr(use_gpu=args.use_gpu, is_efficiency_mode=not args.no_efficiency_mode)

if not os.path.exists(args.output_dir):

os.makedirs(args.output_dir)

images = glob(args.image_dir + '/*.[jJpP][pPnN][gG]')

logger.info(f"Found {len(images)} images in {args.image_dir}")

chunk_size = args.chunk_size

for i in range(0, len(images), chunk_size):

chunk_imgs = images[i:i + chunk_size]

ocr_results = []

for path in tqdm(chunk_imgs):

try:

res = m.ocr(path)

res_list = [i['text'] for i in res if i]

result = "\n".join(res_list)

ocr_results.append(result)

if args.save_box_img:

# Save ocr box img

saved_img_path = os.path.join(args.output_dir, os.path.basename(path))

draw_ocr_boxes(path, res, saved_img_path)

except Exception as e:

logger.error(f'error: {e}, img: {path}')

ocr_results.append("")

# Save partial results to CSV

rows = list(zip(chunk_imgs, ocr_results))

output_file = os.path.join(args.output_dir, 'ocr_results.csv')

save_partial_results(rows, output_file, i == 0)

logger.debug(f'saved partial results. size: {len(chunk_imgs)}, ocr_results size: {len(ocr_results)}')

logger.info(f"Input image_dir {args.image_dir}, saved to {args.output_dir} success.")

def main():

args = parse_args()

cli(args)

if __name__ == "__main__":

main()