将自己的数据集引入Pytorch是搭建属于自己的神经网络的重要一步,这里我设计了一个简单的实验,结合这个实验代码,我将逐句教会大家如何将数据引入DataLoader。
这里以目标检测为例,一个batch中包含图片文件、先验框的框体坐标、目标类型,相对而言更加全面。大家亦可根据自己的数据结构和需求进行修改。
一、数据文件分析
标准的Voc格式是无法直接注入模型的,而如果在训练程序中进行处理即拖慢了运算速度,又难以保证数据集分割的一致性。最好是使用一个独立程序完成数据集的分割、组织、暂存。这里参考了Bubbliiiing的做法,将数据集信息暂存为txt文件。其中一条具体数据的格式如下:
../VOC2007/JPEGImages/0.jpg 166,121,336,323,0 1052,372,1371,924,1
# 文件路径(绝对路径为佳)/ # 先验框框体信息1/ # 先验框框体信息1/
文件路径和框体信息之间采用空格分开;框体信息内部以逗号分开,前4个为坐标信息,最后一个为分类信息。
在随后的程序中,我们将循环读取这个文件中的数据来获取数据集信息。
二、载入数据
1.定义DataSet超参
在开始重写DataSet前,我们需要定义一些用来控制DataSet的参数。
#----自定义DataSet,继承自torch.utils.data.DataSet
class MyDataSet(Dataset):
#----参数定义,输入的参数分别为数据行、输入图像尺寸,类型数
def __init__(self,file_Lines,inp_Shape,num_Classes):
super(MyDataSet,self).__init__()
# 将局部形参变为类内的全局变量
self.length = len(file_Lines) #将文件数赋值给length
self.file_Lines = file_Lines
self.inp_Shape = inp_Shape
self.num_Classes = num_Classes
2.重写len函数
没什么技巧,因为我们刚刚将数据行(file_Lines)的长度赋值给了self.length,直接返回这个值就能拿到数据集的长度了。这也是设置超参数的意义所在。
def __len__(self):
return self.length
3.重写getitem函数
此函数每次会获取单个文件,DataLoader通过反复调用这个函数最终获取整个数据集,我们写入的,这个函数自带一个index用于控制获取的行数。
①分解数据集文件
按照我们上面解析的文件格式,我们用split函数分割index行的文本(空格分割),得到的file_item中第1个元素为文件的绝对路径,后续元素为目标先验框的信息。
def __getitem__(self, index):
index = index % self.length #计算batch长度
file_item = self.file_Lines[index].split()
#按空格拆分文件行,其中的元素分别为:文件路径、先验框坐标(n个)
file_item的值如下图所示:
②加载图片文件
同样没什么技巧,拿到文件路径后直接打开就好了。需要注意的是神经网络的输入需要为固定的形状(图片尺寸和颜色通道数),如果图片为灰阶图则需要将其颜色通道扩充为3个(RGB图)
img = Image.open(file_item[0]) #打开图片
img = cvt2RGB(img)
#若图像为灰度图需要先转换为RGB图(神经网络输入为3通道)
#----将图像转换为RGB----#
def cvt2RGB(img):
if len(np.shape(img)) == 3 and np.shape(img)[2]==3:
return img #为RGB不需要转换
else:
img = img.convert('RGB')
return img
③拆分框体信息
同样同样没什么好说的,遍历分割file_item从1开始的元素就好了
box_info = np.array([np.array(list(map(int,box.split(',')))) for box in file_item[1:]])
#从文件中加载先验框坐标和类型(从第1个元素开始)
④将图片变形
这一步也不是必须的,可以选择在开始训练之前对图片信息进行处理。但是在程序中处理需要注意一点,在改变图像的同时需要以同样的比例改变先验框的坐标。
img,new_box = self.resize_img_withBox(img,box_info,self.inp_Shape)
这里给出一个无损变换大小的函数,若不指定参数则不变化。
def resize_img_withBox(self,img, box, size=[0,0]): #输入参数分别为:原图、先验框列表、变形后的图片大小
iw,ih = img.size
w,h = size
new_box=[]
# 若没有指定大小则不需要变形,若指定了大小则进行变形
if size!=(0,0):
scale = min(w/iw,h/ih) #获取变形比例
nw = int(iw*scale) #计算变形后的长宽
nh = int(ih*scale)
dx = (w-nw)//2
dy = (h-nh)//2
# 图像变形
img = img.resize((nw,nh),Image.BICUBIC)
new_img = Image.new('RGB',size,(128,128,128)) #创建一张灰色背景
new_img.paste(img,((w-nw)//2,(h-nh)//2)) #将变形后的图片贴进背景中央
# 先验框变形
if len(box)>0:
np.random.shuffle(box)
box[:, [0,2]] = box[:, [0,2]]*nw/iw + dx
box[:, [1,3]] = box[:, [1,3]]*nh/ih + dy
box[:, 0:2][box[:, 0:2]<0] = 0
box[:, 2][box[:, 2]>w] = w
box[:, 3][box[:, 3]>h] = h
box_w = box[:, 2] - box[:, 0]
box_h = box[:, 3] - box[:, 1]
new_box = box[np.logical_and(box_w>1, box_h>1)] # discard invalid box
return new_img,new_box
⑤变换图片的通道
标准的图片通道为RGB,而在pytorch中图片的通道为BGR,所以我们需要对通道进行调整,同时为其附加batch通道。
img = np.transpose(preprocess_input(np.array(img, dtype=np.float32)), (2, 0, 1))
前面的函数是一个增强函数,会给RGB三个通道加上不同的权值,至于权值则是一个默认权值(我也不知道为什么用这个)
#----为图像加权,这是一般默认的参数----#
def preprocess_input(image):
image = np.array(image,dtype = np.float32)[:, :, ::-1]
mean = [0.40789655, 0.44719303, 0.47026116]
std = [0.2886383, 0.27408165, 0.27809834]
return (image / 255. - mean) / std
⑥拆分坐标信息和分类信息
如题,将坐标信息和分类信息从先验框信息组中进行分割
# 拆分先验框坐标和类型
box_data = np.zeros((len(new_box),5))
if(len(box_info)>0):
box_data[:len(box_info)] = box_info
box = box_data[:,:4]
label = box_data[:,-1]
完成上述步骤后,将得到的数据返回即可,完整的getitem函数如下:
def __getitem__(self, index):
index = index % self.length #计算batch长度
# 读取文件
file_item = self.file_Lines[index].split() #按空格拆分文件行,其中的元素分别为:文件路径、先验框坐标(n个)
img = Image.open(file_item[0]) #打开图片
img = cvt2RGB(img) #若图像为灰度图需要先转换为RGB图(神经网络输入为3通道)
box_info = np.array([np.array(list(map(int,box.split(',')))) for box in file_item[1:]]) #从文件中加载先验框坐标和类型(从第1个元素开始)
# 对图像进行变形(含先验框变形)
img,new_box = self.resize_img_withBox(img,box_info,self.inp_Shape)
# 将图像进行加权
#img = np.transpose(preprocess_input(np.array(img, dtype=np.float32)), (2, 0, 1))
img = np.transpose(np.array(img))
# 拆分先验框坐标和类型
box_data = np.zeros((len(new_box),5))
if(len(box_info)>0):
box_data[:len(box_info)] = box_info
box = box_data[:,:4]
label = box_data[:,-1]
return img,box,label
三、数据封包
我们在训练时肯定不能这样一个一个训练,一般情况我们训练时会将这些数据打包成一个个的patch送给迭代器,而collate_fn就是做这个的,需要注意collate_fn并不是DataSet类的成员。
这个函数使dataloader自动使用的,其中的images、bboxes、labels 将会在训练过程中用到,这里我们只要确保将数据装入对应的容器中即可。
# DataLoader中collate_fn使用
def my_collate(batch):
images = []
bboxes = []
labels = []
for img, box, label in batch:
images.append(img)
bboxes.append(box)
labels.append(label)
images = np.array(images)
return images, bboxes, labels
四、调用
①读取数据集文件(txt)
file_path = "数据集文件的路径"
with open(file_path) as f:
train_lines = f.readlines()
②实例化DataSet
train_dataset = MyDataSet(train_lines,input_shape,num_classes)
train_loader = DataLoader(train_dataset, shuffle = True, batch_size = 32, num_workers = 1, collate_fn=my_collate)
其中num_workers是线程数;batch_size是单个batch的大小;collate_fn指向我们刚刚重写的collate_fn;shuffle表示是否打乱数据集的顺序。
③训练
当然,这里不是真的训练,我们用一个展示函数来代替训练。
print("开始打印结果")
for item in train_dataset:
img, box, label = item
img = img.transpose(1,2,0)
print(img.shape)
im = Image.fromarray(np.uint8(img))
im.show()
input("按任意键继续")
这个昆虫数据集太恶心了就不给大家看了