基于Python的离线OCR图片文字识别（四）——支持txt文件指定路径保存

虽然在前面在第二次升级时就已经通过json配置文件支持将ocr识别结果txt保存到指定的文件夹里了，但由于指定待识别文件夹时文件夹里面可能包含多个不同的子文件夹、不同的子文件夹里面可能包含同名的图像文件，而原来的方式是直接把所有的txt全部放在json文件指定的一个文件夹中，当不同文件夹中存在同名的图像文件时，会存在txt文件覆盖的情况，虽然几率很小但是开发小伙伴们反映确实出现了这种情况，那么就需要改进，最好是改进为下面的方式：在指定的文件夹下面按照原图像文件的目录结构新建相同的文件夹结构并存放txt文件，也即在json指定的txt保存路径下重新按照待识别图像文件夹的结构，完全重新新建相同的文件夹结果，这样就可以完全避免由于大文件夹中的文件重名而带来的识别结果txt文件覆盖从而丢失的情况发生了，升级后的代码如下所示：

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

#帮助信息
strhelp='''
img2txt is one program to get ocr texts from image or pdf 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": 1000,
    "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=1000        #(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')
    try:
        os.system('mkdir '+txtpath)
    except(Exception) as e:
        print('\tMaking txt directory Error: ', e)#输出异常错误
        print('\tPlease input a legal txtpath in the config.json file and try again!\n')
        sys.exit()
        

#------------------开始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文件目录
                    #原来的方式是直接把所有的txt全部放在指定的一个文件夹中，当不同文件夹中存在同名的图像文件时，会存在txt文件覆盖的情况
                    #txtfilename=os.path.join(txtpath, basename)#拼接得到txt文件的绝对路径
                    #下面的方式是在指定的文件夹下面按照原图像文件的目录结构新建相同的文件夹结构并存放txt文件
                    relativeimgpath=imgfilepath.lstrip(argpath)#图片绝对路径左减去命令行指定的路径argpath得到图像文件的内部相对路径
                    newtxtpath=os.path.join(txtpath,relativeimgpath)#指定txt文件路径+图像内部相对路径（还带有图像文件名和后缀名）
                    basedir=os.path.dirname(newtxtpath)#去掉图像文件名和后缀名
                    if not os.path.exists(basedir):#上面的新文件路径不一定存在
                        try:
                            os.system('mkdir '+basedir)#新建文件夹
                        except(Exception) as e:
                            print('\tMaking txt directory Error: ', e)#输出异常错误
                            print('\tTxt file will be storded in the image file directory!')
                            txtfilename=os.path.join(root, basename)#路径+txt文件名
                    txtfilename=os.path.join(basedir, basename)#新路径+txt文件名
                else:
                    txtfilename=os.path.join(root, basename)#路径+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