笔者在学习计算机视觉时，需要经常借助脚本对数据集进行预处理，现将常用的脚本总结如下：
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)
文中部分代码引用自其他博主博客，仅用作学习用途，在此表示感谢，如有侵权，可联系我删除。
版权声明：本文为qq_38233258原创文章，遵循 CC 4.0 BY-SA 版权协议，转载请附上原文出处链接和本声明。