基于Python的离线OCR图片文字识别（三）——支持PDF文件

前面第一个版本实现了基本的ocr功能，可以对某图像文件进行处理，将ocr结果以同名txt文件的方式保存在图像文件同路径下；

然后在第二个版本中又实现了对文件夹参数的支持，也即可以对某个包含大量图像文件的文件夹进行处理；同时还支持参数配置文件，以json文件的形式支持关键参数的配置，例如：设置txt文件的保存结果（当然为空时就还是以前的保存在图像文件同目录下）、设置排除字符（离线ocr过程中容易出现无意义的乱码的字符）等；

但是，系统开发的小伙伴们又提出了新需求了，说实际应用中发现虽然大部分人员档案都是扫描后保存为图像格式，但也有一些人的档案由于历史原因是以pdf图像的方式保存为一个pdf文件的，能不能也支持对这类文件的离线ocr？当然可以了，本质上都是图像格式嘛！只需要找一个可以将pdf转换为图像的工具即可（当然，对于本身就是文字类型的pdf文件，可以利用其它工具直接提取其中的文字部分，不需要脱裤子放屁多此ocr识别一举）！于是，对img2txt.py的第三次升级改造来了——

#!/home/super/miniconda3/bin/python
#encoding=utf-8
#author: superchao1982, [email protected]

#帮助信息
strhelp='''
img2txt is one program to get ocr texts from image files!

default threshold is 0.1;
default langpath is '/home/langdata' for linux and 'C:\ocr\langdata' for win;
default remove char is '| _^~`&';
default path storing the ocr texts are the same directory with images;
default settings above can be changed in the file 'config.json' which stored in langpath;

contents in config.json like:
{
    "threshold": 0.1,
    "batchsize": 2,
    "workernum": 4,
    "maximgsize": 2000,
    "allowlist": "",
    "langpath": "/home/langdata",
    "removechar": " _^~`&"
    "txtpath": ""
}
------------------------------------
e.g.
./img2txt.py img1.jpg jmg2.jpg #follow by one or more image files
./img2txt.py ./img1 ./img home/usr/Document/img #follow by one or more directory contain image files
./img2txt.py --help #output the help info
./img2txt.py --config #generate the default config.json file in the langpath
------------------------------------
'''
import sys
import json
import os
import pdf2image
import numpy as np

#------------------默认参数设置----------------------
threshold=0.1        #(default = 0.1)阈值
batchsize=2            # (default = 1) - batch_size>1 will make EasyOCR faster but use more memory
workernum=4            # (default = 0) - Number thread used in of dataloader
maximgsize=2000        #(default = 1000) - Max image width & height when using pdf
allowlist=''        # (string) - Force EasyOCR to recognize only subset of characters
removechar='| _^~`&'#待删除无效字符
txtpath=''            #ocr识别后同名txt文件存放的位置:空表示同一目录，点表示相对目录，其他表示绝对目录
#根据系统设置默认的语言包路径
if sys.platform.lower().startswith('linux'):
    langpath='/home/langdata'
elif sys.platform.lower().startswith('win'):
    langpath='C:\ocr\langdata'
else:
    print('Error: Unknow System!')
    sys.exit()
#配置参数字典
config={
    "threshold": threshold,
    "batchsize": batchsize,
    "workernum": workernum,
    "maximgsize": maximgsize,
    "allowlist": allowlist,
    "langpath": langpath,
    "removechar": removechar,
    "txtpath": txtpath
}

#------------------命令行参数处理----------------------
#首先对输入的命令行参数进行处理，在加载ocr包之前排查的好处是避免临处理时出错白白浪费时间
for i in range(1,len(sys.argv)):#获取命令行参数：argv[0]表示可执行文件本身
    if sys.argv[i] in ['-h', '--help']:
        print(strhelp)
        sys.exit()
    elif sys.argv[i] in ['-c', '--config']:
        #保存字典到文件
        try:
            with open(os.path.join(langpath,'config.json'), 'w') as jsonfile:
                json.dump(config, jsonfile, ensure_ascii=False,indent=4)
            print('Genrerating config.json success! ---> ', os.path.join(langpath,'config.json'))
        except(Exception) as e:
            print('\tSaving config file config.json Error: ', e)#输出异常错误
        sys.exit()
    else:
        #check the image file or directory is valid-提前校验，免得浪费时间加载easyocr模型
        if not os.path.exists(sys.argv[i]):
            print(sys.argv[i], ' is invalid, please input the correct file or directory path!')
            sys.exit()

#检查语言包路径是否正确check the langpath is valid
if not os.path.exists(langpath):
    print('Error: Invalid langpath! Checking the path again!')
    sys.exit()

#判断是否存在配置文件config.json,存在就使用,格式如下：
configfile=os.path.join(langpath,'config.json')
if os.path.exists(configfile):
    try:
        with open(configfile, 'r') as jsonfile:
            configdict=json.load(jsonfile)
        threshold=configdict['threshold']
        batchsize=configdict['batchsize']
        workernum=configdict['workernum']
        maximgsize=configdict['maximgsize']
        langpath=configdict['langpath']
        allowlist=configdict['allowlist']
        removechar=configdict['removechar']
        txtpath=configdict['txtpath']
        print('using the config in ', configfile)
    except(Exception) as e:
        print('\tReading config file ', configfile ,' Error: ', e)#输出异常错误
        print('\tCheck the json file, or remove the config.json file to use defaulting configs!')
        sys.exit()
else:
    print('\tusing the default config in ', langpath)
print(configdict)

#如果用户在config.json中指定的txt文件保存路径不存在就生成一个
if len(txtpath)>0 and not os.path.exists(txtpath):
    print('txtpath in config.json is not exists, generating ', txtpath, '!\n')
    os.system('mkdir '+txtpath)

#------------------开始OCR识别----------------------
import easyocr
ocrreader=easyocr.Reader(['ch_sim', 'en'], model_storage_directory=langpath)#Linux: r'/home/langdata', Windows: r'C:\ocr\langdata'
for ind in range(1,len(sys.argv)):#获取命令行参数：argv[0]表示可执行文件本身
    argpath=sys.argv[ind]
    #如果是文件...
    if os.path.isfile(argpath):
        paper=''
        #获取文件后缀名
        filext=os.path.splitext(argpath)[-1]
        if filext.upper() not in ['.JPG','.JPEG','.PNG','.BMP','.PDF']:#转换为大写后再比对
            print('\t', argpath, ' 不是有效图片格式(jpg/jpeg/png/bmp/pdf)!')
            continue
        if filext.upper() in['.PDF']:#如果是pdf文档
            images=pdf2image.convert_from_path(argpath)#将pdf文档转换为图像序列
            for i in range(len(images)):#如果图片尺寸过大，缩小到特定尺寸，避免内存崩溃
                ratio=max(images[i].width, images[i].height)/maximgsize
                if ratio>1:
                    images[i]=images[i].resize((round(images[i].width/ratio),round(images[i].height/ratio)))
                result = ocrreader.readtext(np.asarray(images[i]),batch_size=batchsize,workers=workernum)
                for w in result:
                    if w[2]>threshold:#设置一定的置信度阈值
                        paper = paper+w[1]
        else:
            result = ocrreader.readtext(argpath,batch_size=batchsize,workers=workernum)
            for w in result:
                if w[2]>threshold:#设置一定的置信度阈值
                    paper = paper+w[1]
        #print(paper)
        for item in removechar:
            paper=paper.replace(item, '')
        paper=paper.replace('\r', '')
        paper=paper.replace('\n', '')
        #记录当前文件的识别结果，保存为同名的txt文件
        if(len(txtpath)>0):#如果设置了txt文件目录
            basename=os.path.basename(argpath)+'.txt'#与原文件同名的txt文件（不含目录仅文件名）
            txtfilename=os.path.join(txtpath, basename)
        else:
            txtfilename=os.path.splitext(argpath)[0]+'.txt'#与原文件同名的txt文件（包括目录）
        print('saving file ---> ', txtfilename)#保存的文件名字
        try:
            with open(txtfilename, 'w') as txtfile:
                txtfile.write(paper)
        except(Exception) as e:
            print('\t', txtfilename, ' Saving txt File Error: ', e)#输出异常错误
            continue
    #如果是文件夹...
    if os.path.isdir(argpath):
        for root, _, filenames in os.walk(argpath):
            for imgfile in filenames:
                paper=''
                filext=os.path.splitext(imgfile)[-1]#文件后缀名
                if filext.upper() not in ['.JPG','.JPEG','.PNG','.BMP','.PDF']:
                    print('\t', imgfile, '的后缀名不是有效的图像格式，跳过该文件！')
                    continue
                imgfilepath=os.path.join(root, imgfile)#文件绝对路径
                if filext.upper() in['.PDF']:#如果是pdf文档
                    images=pdf2image.convert_from_path(imgfilepath)#将pdf文档转换为图像序列
                    for i in range(len(images)):#如果图片尺寸过大，缩小到特定尺寸，避免内存崩溃
                        ratio=max(images[i].width, images[i].height)/maximgsize
                        if ratio>1:
                            images[i]=images[i].resize((round(images[i].width/ratio),round(images[i].height/ratio)))
                        result = ocrreader.readtext(np.asarray(images[i]),batch_size=batchsize,workers=workernum)
                        for w in result:
                            if w[2]>threshold:#设置一定的置信度阈值
                                paper = paper+w[1]
                else:
                    result = ocrreader.readtext(imgfilepath,batch_size=batchsize,workers=workernum)
                    for w in result:
                        if w[2]>threshold:#设置一定的置信度阈值
                            paper = paper+w[1]
                #print(paper)
                for item in removechar:
                    paper=paper.replace(item, '')
                paper=paper.replace('\r', '')
                paper=paper.replace('\n', '')
                #记录当前文件的识别结果，保存为同名的txt文件
                basename=os.path.splitext(imgfile)[0]+'.txt'#与原文件同名的txt文件（不包括目录）
                if(len(txtpath)>0):#如果设置了txt文件目录
                    txtfilename=os.path.join(txtpath, basename)
                else:
                    txtfilename=os.path.join(root, basename)#需要加上绝对路径
                print('saving file ---> ', txtfilename)#保存的文件名字
                try:
                    with open(txtfilename, 'w') as txtfile:
                        txtfile.write(paper)
                except(Exception) as e:
                    print('\t', txtfilename, ' Saving txt File Error: ', e)#输出异常错误
                    continue