CV学习笔记 — 数据集预处理常用脚本总结

笔者在学习计算机视觉时，需要经常借助脚本对数据集进行预处理，现将常用的脚本总结如下：

1. 批量修改文件后缀名

# 批量修改
import os
import sys
# 需要修改后缀的文件目录
os.chdir(r'H:\葡萄\datasets\JPEGImages')

# 列出当前目录下所有的文件
files = os.listdir('./')
print('files',files)

for fileName in files:
    portion = os.path.splitext(fileName)
    newName = portion[0] + ".jpg" # 修改为目标后缀
    os.rename(fileName, newName)

2. 对数据集图片进行裁剪

import cv2
import os
import sys
import time

def get_img(input_dir):
    img_paths = []
    for (path,dirname,filenames) in os.walk(input_dir):
        for filename in filenames:
            img_paths.append(path+'/'+filename)
    print("img_paths:",img_paths)
    return img_paths

def cut_img(img_paths,output_dir):
    scale = len(img_paths)
    for i,img_path in enumerate(img_paths):
        a = "#"* int(i/1000)
        b = "."*(int(scale/1000)-int(i/1000))
        c = (i/scale)*100
        time.sleep(0.2)
        print('正在处理图片: %s' % img_path.split('/')[-1])
        img = cv2.imread(img_path)

        cropImg = img[0:200, 0:200]    # 裁剪【y1,y2:x1,x2】

        cv2.imwrite(output_dir + '/' + img_path.split('/')[-1], cropImg)
        print('{:^3.3f}%[{}>>{}]'.format(c,a,b))

if __name__ == '__main__':
    output_dir = "C:/Users/XY/Desktop/222"    # 处理后图片保存目录
    input_dir = "C:/Users/XY/Desktop/111"     # 处理前图片保存目录
    img_paths = get_img(input_dir)
    print('图片读取完成~')
    cut_img(img_paths,output_dir)

3. VOC格式转COCO格式

import xml.etree.ElementTree as ET
import os
import json
 
coco = dict()
coco['images'] = []
coco['type'] = 'instances'
coco['annotations'] = []
coco['categories'] = []
 
category_set = dict()
image_set = set()
 
category_item_id = -1
image_id = 20180000000
annotation_id = 0
 
def addCatItem(name):
    global category_item_id
    category_item = dict()
    category_item['supercategory'] = 'none'
    category_item_id += 1
    category_item['id'] = category_item_id
    category_item['name'] = name
    coco['categories'].append(category_item)
    category_set[name] = category_item_id
    return category_item_id
 
def addImgItem(file_name, size):
    global image_id
    if file_name is None:
        raise Exception('Could not find filename tag in xml file.')
    if size['width'] is None:
        raise Exception('Could not find width tag in xml file.')
    if size['height'] is None:
        raise Exception('Could not find height tag in xml file.')
    image_id += 1
    image_item = dict()
    image_item['id'] = image_id
    image_item['file_name'] = file_name
    image_item['width'] = size['width']
    image_item['height'] = size['height']
    coco['images'].append(image_item)
    image_set.add(file_name)
    return image_id
 
def addAnnoItem(object_name, image_id, category_id, bbox):
    global annotation_id
    annotation_item = dict()
    annotation_item['segmentation'] = []
    seg = []
    # bbox[] is x,y,w,h
    # left_top
    seg.append(bbox[0])
    seg.append(bbox[1])
    # left_bottom
    seg.append(bbox[0])
    seg.append(bbox[1] + bbox[3])
    # right_bottom
    seg.append(bbox[0] + bbox[2])
    seg.append(bbox[1] + bbox[3])
    # right_top
    seg.append(bbox[0] + bbox[2])
    seg.append(bbox[1])
 
    annotation_item['segmentation'].append(seg)
 
    annotation_item['area'] = bbox[2] * bbox[3]
    annotation_item['iscrowd'] = 0
    annotation_item['ignore'] = 0
    annotation_item['image_id'] = image_id
    annotation_item['bbox'] = bbox
    annotation_item['category_id'] = category_id
    annotation_id += 1
    annotation_item['id'] = annotation_id
    coco['annotations'].append(annotation_item)
 
def _read_image_ids(image_sets_file):
    ids = []
    with open(image_sets_file) as f:
        for line in f:
            ids.append(line.rstrip())
    return ids
 
"""通过txt文件生成"""
#split ='train' 'va' 'trainval' 'test'
def parseXmlFiles_by_txt(data_dir,json_save_path,split='train'):
    print("hello")
    labelfile=split+".txt"
    image_sets_file = data_dir + "/ImageSets/Main/"+labelfile
    ids=_read_image_ids(image_sets_file)
 
    for _id in ids:
        xml_file=data_dir + f"/Annotations/{_id}.xml"
 
        bndbox = dict()
        size = dict()
        current_image_id = None
        current_category_id = None
        file_name = None
        size['width'] = None
        size['height'] = None
        size['depth'] = None
 
        tree = ET.parse(xml_file)
        root = tree.getroot()
        if root.tag != 'annotation':
            raise Exception('pascal voc xml root element should be annotation, rather than {}'.format(root.tag))
 
        # elem is <folder>, <filename>, <size>, <object>
        for elem in root:
            current_parent = elem.tag
            current_sub = None
            object_name = None
 
            if elem.tag == 'folder':
                continue
 
            if elem.tag == 'filename':
                file_name = elem.text
                if file_name in category_set:
                    raise Exception('file_name duplicated')
 
            # add img item only after parse <size> tag
            elif current_image_id is None and file_name is not None and size['width'] is not None:
                if file_name not in image_set:
                    current_image_id = addImgItem(file_name, size)
                    print('add image with {} and {}'.format(file_name, size))
                else:
                    raise Exception('duplicated image: {}'.format(file_name))
                    # subelem is <width>, <height>, <depth>, <name>, <bndbox>
            for subelem in elem:
                bndbox['xmin'] = None
                bndbox['xmax'] = None
                bndbox['ymin'] = None
                bndbox['ymax'] = None
 
                current_sub = subelem.tag
                if current_parent == 'object' and subelem.tag == 'name':
                    object_name = subelem.text
                    if object_name not in category_set:
                        current_category_id = addCatItem(object_name)
                    else:
                        current_category_id = category_set[object_name]
 
                elif current_parent == 'size':
                    if size[subelem.tag] is not None:
                        raise Exception('xml structure broken at size tag.')
                    size[subelem.tag] = int(subelem.text)
 
                # option is <xmin>, <ymin>, <xmax>, <ymax>, when subelem is <bndbox>
                for option in subelem:
                    if current_sub == 'bndbox':
                        if bndbox[option.tag] is not None:
                            raise Exception('xml structure corrupted at bndbox tag.')
                        bndbox[option.tag] = int(option.text)
 
                # only after parse the <object> tag
                if bndbox['xmin'] is not None:
                    if object_name is None:
                        raise Exception('xml structure broken at bndbox tag')
                    if current_image_id is None:
                        raise Exception('xml structure broken at bndbox tag')
                    if current_category_id is None:
                        raise Exception('xml structure broken at bndbox tag')
                    bbox = []
                    # x
                    bbox.append(bndbox['xmin'])
                    # y
                    bbox.append(bndbox['ymin'])
                    # w
                    bbox.append(bndbox['xmax'] - bndbox['xmin'])
                    # h
                    bbox.append(bndbox['ymax'] - bndbox['ymin'])
                    print('add annotation with {},{},{},{}'.format(object_name, current_image_id, current_category_id,
                                                                   bbox))
                    addAnnoItem(object_name, current_image_id, current_category_id, bbox)
    json.dump(coco, open(json_save_path, 'w'))
 
"""直接从xml文件夹中生成"""
def parseXmlFiles(xml_path,json_save_path):
    for f in os.listdir(xml_path):
        if not f.endswith('.xml'):
            continue
 
        bndbox = dict()
        size = dict()
        current_image_id = None
        current_category_id = None
        file_name = None
        size['width'] = None
        size['height'] = None
        size['depth'] = None
 
        xml_file = os.path.join(xml_path, f)
        print(xml_file)
 
        tree = ET.parse(xml_file)
        root = tree.getroot()
        if root.tag != 'annotation':
            raise Exception('pascal voc xml root element should be annotation, rather than {}'.format(root.tag))
 
        # elem is <folder>, <filename>, <size>, <object>
        for elem in root:
            current_parent = elem.tag
            current_sub = None
            object_name = None
 
            if elem.tag == 'folder':
                continue
 
            if elem.tag == 'filename':
                file_name = elem.text
                if file_name in category_set:
                    raise Exception('file_name duplicated')
 
            # add img item only after parse <size> tag
            elif current_image_id is None and file_name is not None and size['width'] is not None:
                if file_name not in image_set:
                    current_image_id = addImgItem(file_name, size)
                    print('add image with {} and {}'.format(file_name, size))
                else:
                    raise Exception('duplicated image: {}'.format(file_name))
                    # subelem is <width>, <height>, <depth>, <name>, <bndbox>
            for subelem in elem:
                bndbox['xmin'] = None
                bndbox['xmax'] = None
                bndbox['ymin'] = None
                bndbox['ymax'] = None
 
                current_sub = subelem.tag
                if current_parent == 'object' and subelem.tag == 'name':
                    object_name = subelem.text
                    if object_name not in category_set:
                        current_category_id = addCatItem(object_name)
                    else:
                        current_category_id = category_set[object_name]
 
                elif current_parent == 'size':
                    if size[subelem.tag] is not None:
                        raise Exception('xml structure broken at size tag.')
                    size[subelem.tag] = int(subelem.text)
 
                # option is <xmin>, <ymin>, <xmax>, <ymax>, when subelem is <bndbox>
                for option in subelem:
                    if current_sub == 'bndbox':
                        if bndbox[option.tag] is not None:
                            raise Exception('xml structure corrupted at bndbox tag.')
                        bndbox[option.tag] = int(option.text)
 
                # only after parse the <object> tag
                if bndbox['xmin'] is not None:
                    if object_name is None:
                        raise Exception('xml structure broken at bndbox tag')
                    if current_image_id is None:
                        raise Exception('xml structure broken at bndbox tag')
                    if current_category_id is None:
                        raise Exception('xml structure broken at bndbox tag')
                    bbox = []
                    # x
                    bbox.append(bndbox['xmin'])
                    # y
                    bbox.append(bndbox['ymin'])
                    # w
                    bbox.append(bndbox['xmax'] - bndbox['xmin'])
                    # h
                    bbox.append(bndbox['ymax'] - bndbox['ymin'])
                    print('add annotation with {},{},{},{}'.format(object_name, current_image_id, current_category_id,
                                                                   bbox))
                    addAnnoItem(object_name, current_image_id, current_category_id, bbox)
    json.dump(coco, open(json_save_path, 'w'))
 
if __name__ == '__main__':
 
    ann_path="E:/data/datasets/VOC/Annotations"            # VOC数据集标注存储路径
    json_save_path="E:/data/datasets/coco128/test.json"    # COCO数据集标注存储路径
    parseXmlFiles(ann_path,json_save_path)

4. 数据集图片批量png转为jpg

import os
from PIL import Image

dirname_read="C:/Users/xiey/Desktop/CityPerson/png/"
dirname_write="C:/Users/xiey/Desktop/CityPerson/jpg/"
names=os.listdir(dirname_read)
count=0
for name in names:
    img=Image.open(dirname_read+name)
    name=name.split(".")
    if name[-1] == "png":
        name[-1] = "jpg"
        name = str.join(".", name)
        to_save_path = dirname_write + name
        img.save(to_save_path)
        count+=1
        print(to_save_path, "------conut：",count)
    else:
        continue

5. 数据集批量txt转为xml

import os
import glob
from PIL import Image

voc_annotations = 'C:/Users/xiey/Desktop/CityPerson/A_jpg'  # 图片xml文件存储路径
yolo_txt = 'C:/Users/xiey/Desktop/CityPerson/A'       # 图片txt文件存储路径
img_path = 'C:/Users/xiey/Desktop/CityPerson/I_jpg'   # 图片路径
labels = ['person']  # label for datasets
# 图片存储位置
src_img_dir = img_path 
# 图片的txt文件存放位置
src_txt_dir = yolo_txt
# 图片的xml文件存放位置
src_xml_dir = voc_annotations

img_Lists = glob.glob(src_img_dir + '/*.jpg')

img_basenames = []
for item in img_Lists:
    img_basenames.append(os.path.basename(item))

img_names = []
for item in img_basenames:
    temp1, temp2 = os.path.splitext(item)
    img_names.append(temp1)

for img in img_names:
    im = Image.open((src_img_dir + '/' + img + '.jpg'))
    width, height = im.size

    # 打开txt文件
    gt = open(src_txt_dir + '/' + img + '.txt').read().splitlines()
    print(gt)
    if gt:
        # 将主干部分写入xml文件中
        xml_file = open((src_xml_dir + '/' + img + '.xml'), 'w')
        xml_file.write('<annotation>\n')
        xml_file.write('    <folder>VOC2007</folder>\n')
        xml_file.write('    <filename>' + str(img) + '.jpg' + '</filename>\n')
        xml_file.write('    <size>\n')
        xml_file.write('        <width>' + str(width) + '</width>\n')
        xml_file.write('        <height>' + str(height) + '</height>\n')
        xml_file.write('        <depth>3</depth>\n')
        xml_file.write('    </size>\n')

        # write the region of image on xml file
        for img_each_label in gt:
            spt = img_each_label.split(' ')  # 这里如果txt里面是以逗号‘，’隔开的，那么就改为spt = img_each_label.split(',')。
            print(f'spt:{spt}')
            xml_file.write('    <object>\n')
            xml_file.write('        <name>' + str(labels[int(spt[0])]) + '</name>\n')
            xml_file.write('        <pose>Unspecified</pose>\n')
            xml_file.write('        <truncated>0</truncated>\n')
            xml_file.write('        <difficult>0</difficult>\n')
            xml_file.write('        <bndbox>\n')

            center_x = round(float(spt[2].strip()) * width)
            center_y = round(float(spt[3].strip()) * height)
            bbox_width = round(float(spt[4].strip()) * width)
            bbox_height = round(float(spt[5].strip()) * height)
            xmin = str(int(center_x - bbox_width / 2))
            ymin = str(int(center_y - bbox_height / 2))
            xmax = str(int(center_x + bbox_width / 2))
            ymax = str(int(center_y + bbox_height / 2))

            xml_file.write('            <xmin>' + xmin + '</xmin>\n')
            xml_file.write('            <ymin>' + ymin + '</ymin>\n')
            xml_file.write('            <xmax>' + xmax + '</xmax>\n')
            xml_file.write('            <ymax>' + ymax + '</ymax>\n')
            xml_file.write('        </bndbox>\n')
            xml_file.write('    </object>\n')

        xml_file.write('</annotation>')

6. 数据集批量后缀".JPG"转为".jpg"

import os

# 列出当前目录下所有的文件
files = os.listdir(".")

for filename in files:
    portion = os.path.splitext(filename)
    # 如果后缀是.JPG
    if portion[1] == ".JPG":
        # 重新组合文件名和后缀名
        newname = portion[0] + ".jpg"
        os.rename(filename,newname)

7. 混淆矩阵

#coding=utf-8
import matplotlib.pyplot as plt
import numpy as np
# 二进制网络
#confusion = np.array(([349,9,11,4, 10],
                      #[21,87,10,5,  2],
                      #[12,3,171,5,  0],
                      #[2, 1,  8,86, 0],
                      #[16,0,  0,0, 88]))
# Faster R-CNN
confusion = np.array(([37,18,16,17,10,13,14],
                      [19,37,17,3, 2,7,15],
                      [10,12,38,12, 18,10,11],
                      [5, 7,  9,38, 5,12,5],
                      [4, 10, 12,12,38,8,3],
                      [5, 12, 5,9, 7,39,14],
                      [20,8, 5, 9, 20,11,38]))
# VGG-16
#confusion = np.array(([35,18,16,17,10,13,14],
                      #[21,33,17,3, 2,7,15],
                      #[10,14,35,12, 18,10,11],
                      #[5, 5,  10,34, 5,12,4],
                      #[4, 10, 12,16,35,8,4],
                      #[5, 12, 5,9, 10,36,15],
                      #[20,8, 5, 9, 20,14,35]))
# ResNet50
#confusion = np.array(([36,18,16,17,10,13,14],
                      #[20,35,17,3, 2,7,15],
                      #[10,12,37,12, 18,10,11],
                      #[5, 5,  8,38, 5,12,5],
                      #[4, 10, 12,12,37,8,3],
                      #[5, 12, 5,9, 8,38,16],
                      #[20,8, 5, 9, 20,12,36]))
# 热度图，后面是指定的颜色块，可设置其他的不同颜色
plt.imshow(confusion, cmap=plt.cm.Blues)
# ticks 坐标轴的坐标点
# label 坐标轴标签说明
indices = range(len(confusion))
# 第一个是迭代对象，表示坐标的显示顺序，第二个参数是坐标轴显示列表
#plt.xticks(indices, [0, 1, 2])
#plt.yticks(indices, [0, 1, 2])
plt.xticks(indices, ['深灰色泥岩', '黑色煤', '灰色细砂岩','浅灰色细砂岩','深灰色粉砂质泥岩','灰黑色泥岩','灰色泥质粉砂岩'], rotation='vertical')
plt.yticks(indices, ['深灰色泥岩', '黑色煤', '灰色细砂岩','浅灰色细砂岩','深灰色粉砂质泥岩','灰黑色泥岩','灰色泥质粉砂岩'])

plt.colorbar()
plt.xlabel('预测值')
plt.ylabel('真实值')
#plt.title('Binary Faster R-CNN')
#plt.title('Faster R-CNN')
#plt.title('VGG-16')
plt.title('S-ResNet50')

# plt.rcParams两行是用于解决标签不能显示汉字的问题
plt.rcParams['font.sans-serif']=['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 显示数据
for first_index in range(len(confusion)):    #第几行
    for second_index in range(len(confusion[first_index])):    #第几列
        plt.text(first_index, second_index, confusion[first_index][second_index],va='center', ha='center')
# 在matlab里面可以对矩阵直接imagesc(confusion)
# 显示
plt.show()
plt.savefig('Data/Binary Faster R-CNN.png')

8. 图片亮度处理

import cv2
import numpy as np
import os
import time

def get_img(input_dir):
    img_paths = []
    for (path,dirname,filenames) in os.walk(input_dir):
        for filename in filenames:
            img_paths.append(path+'/'+filename)
    print("img_paths:",img_paths)
    return img_paths

# def contrast_brightness_demo(image, c, b):  # C 是对比度，b 是亮度
#     h, w, ch = image.shape
#     blank = np.zeros([h, w, ch], image.dtype)
#     dst = cv2.addWeighted(image, c, blank, 1-c, b)   # 改变像素的API
#     cv2.imshow("con-bri-demo", dst)

# src=cv2.imread('E:/imgyuchuli/1.jpg')  
# cv2.namedWindow("input image",cv2.WINDOW_AUTOSIZE)  
# print(src)
# cv2.imshow("input image",src)  # 显示图片
# contrast_brightness_demo(src, 1.2, 10)
# cv2.waitKey(0)  
# cv2.destroyAllWindows()  

def process_img(img_paths,output_dir):
    scale = len(img_paths)
    for i,img_path in enumerate(img_paths):
        a = "#"* int(i/1000)
        b = "."*(int(scale/1000)-int(i/1000))
        c = (i/scale)*100
        time.sleep(0.2)
        print('正在处理图片: %s' % img_path.split('/')[-1])
        img = cv2.imread(img_path)

        c=1.2
        b=25
        h, w, ch = img.shape
        blank = np.zeros([h, w, ch], img.dtype)
        dst = cv2.addWeighted(img, c, blank, 1 - c, b)  # 改变像素的API

        cv2.imwrite(output_dir + '/' + img_path.split('/')[-1], dst)
        print('{:^3.3f}%[{}>>{}]'.format(c,a,b))

if __name__ == '__main__':

    output_dir = "H:\dataset\ld"  # 保存图片目录
    input_dir = "H:\dataset\yt"   # 读取图片目录
    img_paths = get_img(input_dir)
    print('图片读取完成~')
    process_img(img_paths,output_dir)

9. 图片添加噪声

包括椒盐噪声、高斯噪声以及随机噪声

import os
import cv2
import numpy as np
import random

# def sp_noise(noise_img, proportion):
#     '''
#     添加椒盐噪声
#     proportion的值表示加入噪声的量，可根据需要自行调整
#     return: img_noise
#     '''
#     height, width = noise_img.shape[0], noise_img.shape[1]  # 获取高度宽度像素值
#     num = int(height * width * proportion)                 
#     for i in range(num):
#         w = random.randint(0, width - 1)
#         h = random.randint(0, height - 1)
#         if random.randint(0, 1) == 0:
#             noise_img[h, w] = 0
#         else:
#             noise_img[h, w] = 255
#     return noise_img

def gaussian_noise(img, mean, sigma):
    '''
    此函数将产生的高斯噪声加到图片上
    入参:
        img   :  原图
        mean  :  均值
        sigma :  标准差
    返回:
        gaussian_out : 噪声处理后的图片
    '''
    # 将图片灰度标准化
    img = img / 255
    # 产生高斯 noise
    noise = np.random.normal(mean, sigma, img.shape)
    # 将噪声和图片叠加
    gaussian_out = img + noise
    # 将超过 1 的置 1，低于 0 的置 0
    gaussian_out = np.clip(gaussian_out, 0, 1)
    # 将图片灰度范围的恢复为 0-255
    gaussian_out = np.uint8(gaussian_out*255)
    # 将噪声范围搞为 0-255
    # noise = np.uint8(noise*255)
    return gaussian_out# 这里也会返回噪声，注意返回值

# def random_noise(image,noise_num):
#     '''
#     添加随机噪点（实际上就是随机在图像上将像素点的灰度值变为255即白色）
#     param image: 需要加噪的图片
#     param noise_num: 添加的噪音点数目
#     return: img_noise
#     '''
#     # 参数image：，noise_num：
#     img_noise = image
#     # cv2.imshow("src", img)
#     rows, cols, chn = img_noise.shape
#     # 加噪声
#     for i in range(noise_num):
#         x = np.random.randint(0, rows)#随机生成指定范围的整数
#         y = np.random.randint(0, cols)
#         img_noise[x, y, :] = 255
#     return img_noise

def convert(input_dir, output_dir):
    for filename in os.listdir(input_dir):
        path = input_dir + "/" + filename # 获取文件路径
        print("doing... ", path)
        noise_img = cv2.imread(path)#读取图片
        img_noise = gaussian_noise(noise_img, 0, 0.12) # 高斯噪声
        #img_noise = sp_noise(noise_img,0.025)         # 椒盐噪声
        #img_noise  = random_noise(noise_img,500)      # 随机噪声
        cv2.imwrite(output_dir+'/'+filename,img_noise )

if __name__ == '__main__':
    input_dir = "H:/dataset/yt"    # 输入数据文件夹
    output_dir = "H:/dataset/gszs" # 输出数据文件夹
    convert(input_dir, output_dir)

10. 图片翻转处理

包括水平翻转、垂直翻转以及45°顺时针翻转

from PIL import Image
import os
import os.path

rootdir = r'H:\dataset'   # 读取文件夹位置

for parent, dirnames, filenames in os.walk(rootdir):
    for filename in filenames:
        print('parentis :' + parent)
        print('filenameis :' + filename)
        currentPath = os.path.join(parent, filename)
        print('thefulll name of the file is :' + currentPath)
        im = Image.open(currentPath)
        im = im.convert('RGB')
        #out = im.transpose(Image.FLIP_LEFT_RIGHT)      # 水平翻转
        out = im.transpose(Image.FLIP_TOP_BOTTOM)       # 垂直翻转
        #out = im.rotate(45)                            # 45°顺时针翻转

        newname = r"H:\a" + '\\' + filename

        out.save(newname)

文中部分代码引用自其他博主博客，仅用作学习用途，在此表示感谢，如有侵权，可联系我删除。