基于Python的离线OCR图片文字识别（二）——支持文件夹、支持配置文件

在前面基于EasyOCR包实现了基本的离线OCR功能，也即在CentOS系统下通过布设Python环境，即可利用命令行实现对指定的图像或图像序列进行OCR内容识别，例如：

python ./img2txt.py 001.jpg 002.jpg

或者

./img2txt.py 001.jpg 002.jpg

但是，目前这种命令行方式后面只能识别图像文件名（含路径），对于某个文件夹中包含很多扫描的图像文件，就需要一一列出文件名字，显得很麻烦，能不能让命令行除了识别图像文件、还识别文件夹参数呢？如果命令行后面跟的参数是一个文件夹（也即某个路径），那么就自动的对该文件夹进行遍历，对其中的图像类文件进行OCR？答案当然是可以的，这也是第二个版本的img2txt.py，主要升级的地方除了上面所述的命令行参数支持文件夹外，还首次支持了帮助信息输出、通过配置文件支持关键参数的设置等，例如可以通过配置文件指定ocr后的txt文件保存的位置（之前是ocr识别后保存在图像文件所在的目录下，文件名与图像文件相同），如果指定的文本保存文件夹不存在就自动创建等，如下所示：

#!/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,
    "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

#默认参数设置
threshold=0.1#阈值
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,
    "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']
            langpath=configdict['langpath']
            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)

import easyocr
ocrreader=easyocr.Reader(['ch_sim', 'en'], model_storage_directory=langpath)#Linux: r'/home/langdata', Windows: r'C:\ocr\langdata'
for i in range(1,len(sys.argv)):#获取命令行参数：argv[0]表示可执行文件本身
    argpath=sys.argv[i]
    #如果是文件...
    if os.path.isfile(argpath):
        #获取文件后缀名
        argpathxt=os.path.splitext(argpath)[-1]
        if argpathxt.upper() not in ['.JPG','.JPEG','.PNG','.BMP']:#转换为大写后再比对
            print('\t', argpath, ' 不是有效图片格式(jpg/jpeg/png/bmp)!')
            continue
        result = ocrreader.readtext(argpath)
        paper=''
        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:
                filext=os.path.splitext(imgfile)[-1]#文件后缀名
                if filext.upper() not in ['.JPG','.JPEG','.PNG','.BMP']:
                    print('\t', imgfile, '的后缀名不是图像格式，跳过该文件！')
                    continue
                imgfilepath=os.path.join(root, imgfile)#文件绝对路径
                result = ocrreader.readtext(imgfilepath)
                paper=''
                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

注意，并没有在一开始就import easyOCR包，这是因为导入ocr包比较耗费资源，如果前面路径部分出错就白白浪费了时间，所以将这种比较耗费资源的行为放在了确实需要的地方。