python+opencv批量处理文件夹（包括子文件夹）下所有图片

导读

深度学习作为机器视觉研究的热点，每训练一个模型都要经过数据采集与处理、设计网络模型、训练模型、验证模型、测试模型等阶段。俗话说“万事开头难”，由于深度学习需要大量数据，而且构造不同的数据集相对复杂，这就需要有一定批量处理数据的能力。在这里以批量处理图片数据为例，用矩形框标出文件夹下所有图片中的缺陷。本文章记录批量处理数据的方法，以备后续学习。

环境

python3.7.7+opencv3.4.11+pycharm社区版+windows10 环境很好搭建 到官网下载安装包安装即可

文件目录结构

一个目录下有多个子目录，如图所示：

子目录下有包括图片等文件，如图所示：

详细步骤

记住几个关键的python函数：
os.listdir()： 以列表形式返回文件夹下所有文件包括子文件夹
file.open(): 读取文件，返回文件对象，类似于c的FILE。
str.split() 以空格拆分字符串，返回列表

完整代码

'''函数名：draw_rectangle'''
'''此函数画出原始图像上所有缺陷的位置'''

import cv2 as cv
import os
import argparse as arg

'''设置命令行解析器'''
parse=arg.ArgumentParser()
args=parse.add_argument('-i','--input',type=str,help='输入图像父级目录')
args=parse.add_argument('-o','--output',type=str,help='输出图像父级目录')
args=parse.parse_args() #解析命令
args=vars(args)#转变为字典形式

#args['input']='H:\datasets\Castings\Castings'
#args['output']='H:\datasets\Castings\Castings_defects'


paths=os.listdir(args['input'])#获取父级目录下所有子文件夹
child_paths=[]#存放子文件夹路径


print('[INFO]reading parent_dir......')
for path in paths:
    child_paths.append(os.path.join(args['input'],path))

print('[INFO]reading child_dir......')
for child_path in child_paths:
    box_path=child_path+'\\'+'ground_truth.txt'
    i=1 #图片序号
    imgpaths=[] #清空列表
    if os.path.exists(box_path): #如果有真实标签文件 就处理
        try:
            file=open(box_path)
            child_path_file = os.listdir(child_path)  # 获取子文件夹下所有文件
            for path in child_path_file:
                path = path.split('.')
                if path[1] == 'png':
                    imgpath = path[0] + '.png'
                    imgpaths.append(imgpath)  # 存放图片名称
            lines = file.readlines()  # 读取标签文件所有行
            for imgpath in imgpaths:
                fullpath = child_path + '\\' + imgpath
                print(fullpath)
                img = cv.imread(fullpath)  # 读入图像
                for line in lines:
                    line = line.strip().split()  # 去掉开头空格 并以空格分割
                    if float(line[0]) == i:  # 注意转换数据类型
                        img = cv.rectangle(img, (int(float(line[1])), int(float(line[3]))),
                                           (int(float(line[2])), int(float(line[4]))), color=(0, 0, 255), thickness=1)
                # 用矩形框住缺陷 注意先将str转换成float 再转成int型
                outdir=args['output']+'\\'+imgpath[0:5]
                if os.path.exists(outdir):
                    print('[INFO] saving picture {} of the {}'.format(i, imgpath[0:5]))
                    cv.imwrite(outdir + '\\' + imgpath, img)
                else:
                    print('[INFO] saving picture {} of the {}'.format(i, imgpath[0:5]))
                    os.mkdir(outdir)
                    cv.imwrite(outdir + '\\' + imgpath, img)
                i+=1
        except IOError:
            print('读取txt文件失败')

print('[INFO] finished')

如何运行程序
命令行终端输入 python draw_rectangle.py -i 输入图像目录 -o 输出图像目录

图片示例