目标检测的dataset
首先我们需要先将数据集做一个整理的txt,文件叫做ann.txt吧
首先是图片路径,xmin,ymin,xmax,ymax,类别
最简单的dataset就构造出来了
import numpy as np
from torch.utils.data import Dataset
from PIL import Image
class MyDataSet(Dataset):
# 初始化
def __init__(self,annotations_file):
# 定义一个列表，里面有图片名称，bbox，类别
self.annotations=annotations_file
self.length=len(annotations_file)
# 长度
def __len__(self):
return self.length
# 返回第index个数据
def __getitem__(self, index):
# index:序号
# 得到diindex数据的图片数据和y，y包括了(bbox,label)
img,y=self.get_data(self.annotations[index])
# 将img转成(c,h,w)
img=np.transpose(img,(2,0,1))
'''
由于拿到的y数据是字符串类型的，所以需要重新构造一个(len(y),5)的矩阵存储
'''
box_data = np.zeros((len(y), 5))
if len(y) > 0:
box_data[:len(y)] = y
box = box_data[:, :4]
label = box_data[:, -1]
return img,box,label
def get_data(self,annotations_file):
line=annotations_file.split()
img_data=Image.open(line[0])
# 输出图片大小要一致，此时只是简单的resize一下
img_data=img_data.resize((360,360))
bbox=np.array([box.split(',') for box in line[1:]])
img_data=np.array(img_data,np.float32)
return img_data,bbox
# DataLoader中collate_fn使用
# collate_fn：如何取样本的，我们可以定义自己的函数来准确地实现想要的功能,
# 将img一个batch合并，bbox一个batch合并，label一个batch合并
def frcnn_dataset_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
来运行一下看看效果
path='ann.txt'
train_line=[]
with open(path) as f:
train_line=f.readlines()
myd=MyDataSet(train_line)
from torch.utils.data import DataLoader
myloader=DataLoader(myd,batch_size=16,shuffle=False,num_workers=0,collate_fn=frcnn_dataset_collate)
for i,(img,bbox,label) in enumerate(myloader):
print(f'第{i}个batch的数据')
print(f"img的维度:{np.array(img).shape}")
print(f'16个batch的bbox{len(bbox)},每个batch的bbox维度：{bbox[0].shape}')
print(f'16个batch的label:{len(label)},每个batch的bbox维度{label[0].shape}')
break
输出:
第0个batch的数据
img的维度:(16, 3, 360, 360)
16个batch的bbox16,每个batch的bbox维度：(11, 4)
16个batch的label:16,每个batch的bbox维度(11,)
1
对于图片我们还是要进行一些处理，例如如果图片不是三通道怎么办，图片的resize失真怎么办，我们在get_data方法中进行优化
首先来解决通道问题
'''
1.如果图片不是三通道的，那么需要转成三通道
'''
if np.shape(img_data)[2]!=3:
img_data = img_data.convert('RGB')
然后解决resize失真,使用灰条填充，尽量填充较少
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image
path='data/coco128/images/train2017/000000000009.jpg'
img=Image.open(path)
# (480, 640, 3)
h,w=600,600
# PIL.Image打开的图片是(w,h,c)的格式
iw,ih=img.size
scale=min(h/ih,w/iw)
print(scale)
nh=int(ih*scale)
nw=int(iw*scale)
dx = (w-nw)//2
dy = (h-nh)//2
print(nh,nw,dx,dy)
# 缩放到(nw,nh)
img=img.resize((nw,nh),Image.BICUBIC)
img.show()
# new : 这个函数创建一幅给定模式（mode）和尺寸（size）的图片。如果省略 color 参数，
# 则创建的图片被黑色填充满，如果 color 参数是 None 值，则图片还没初始化。
# 创建一个全是灰色的图片
new_image = Image.new('RGB', (w, h), (128, 128, 128))
# python中PIL库中的paste函数的作用为将一张图片覆盖到另一张图片的指定位置去
new_image.paste(img, (dx, dy))
plt.imshow(new_image)
plt.savefig('img/img2')
# image_data = np.array(new_image, np.float32)
创建的都是灰色的图片new_image
将原图复制上去，就实现了添加灰条的操作
将上面代码补充进get_data函数中就可:
def get_data(self,annotations_file):
line=annotations_file.split()
img_data=Image.open(line[0])
'''
1.如果图片不是三通道的，那么需要转成三通道
'''
if np.shape(img_data)[2]!=3:
img_data = img_data.convert('RGB')
# 输出图片大小要一致，此时只是简单的resize一下
# img_data=img_data.resize((360,360))
'''
resize就会发生图片尺寸的变化，训练不会像原图那样好，bbox也会发生变化，该如何处理呢
图片缩放会导致失真，那么能不能不缩放也达到我们想要的尺寸呢，或者最小限度的失真
我们将图片的最大的长度保留下来，把短边用灰边填充
scale就是图片(ih,iw)变化成(h,w)的最小比例,(nh,nw)就是我们按照这个比例resize的图片大小
当然(nh,nw)和(h,w)有空白，空白我们就用灰条填充
'''
h,w=360,360#想要的尺寸
iw,ih=img_data.size
# 比列
scale = min(w / iw, h / ih)
nw = int(iw * scale)
nh = int(ih * scale)
dx = (w - nw) // 2
dy = (h - nh) // 2
img=img_data.resize((nw,nh),Image.BICUBIC)
new_image=Image.new('RGB', (w, h), (128, 128, 128))
new_image.paste(img,(dx,dy))
bbox=np.array([box.split(',') for box in line[1:]])
img_data=np.array(img_data,np.float32)
return img_data,bbox
2
既然图片大小调整了，那么bbox的位置也发生变化了，需要对bbox的坐标也进行调整
if len(box) > 0:
np.random.shuffle(box)
# 图片大小都是按照scale变化的，bbox也应该如此
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
加入到get_data函数中
def get_data(self,annotations_file):
line=annotations_file.split()
img_data=Image.open(line[0])
'''
1.如果图片不是三通道的，那么需要转成三通道
'''
if np.shape(img_data)[2]!=3:
img_data = img_data.convert('RGB')
# 输出图片大小要一致，此时只是简单的resize一下
# img_data=img_data.resize((360,360))
'''
resize就会发生图片尺寸的变化，训练不会像原图那样好，bbox也会发生变化，该如何处理呢
图片缩放会导致失真，那么能不能不缩放也达到我们想要的尺寸呢，或者最小限度的失真
我们将图片的最大的长度保留下来，把短边用灰边填充
scale就是图片(ih,iw)变化成(h,w)的最小比例,(nh,nw)就是我们按照这个比例resize的图片大小
当然(nh,nw)和(h,w)有空白，空白我们就用灰条填充
'''
h,w=360,360#想要的尺寸
iw,ih=img_data.size
# 比列
scale = min(w / iw, h / ih)
nw = int(iw * scale)
nh = int(ih * scale)
dx = (w - nw) // 2
dy = (h - nh) // 2
img=img_data.resize((nw,nh),Image.BICUBIC)
new_image=Image.new('RGB', (w, h), (128, 128, 128))
new_image.paste(img,(dx,dy))
bbox=np.array([box.split(',') for box in line[1:]])
if len(bbox) > 0:
np.random.shuffle(bbox)
# 图片大小都是按照scale变化的，bbox也应该如此
bbox[:, [0, 2]] = bbox[:, [0, 2]] * nw / iw + dx
bbox[:, [1, 3]] = bbox[:, [1, 3]] * nh / ih + dy
bbox[:, 0:2][bbox[:, 0:2] < 0] = 0
bbox[:, 2][bbox[:, 2] > w] = w
bbox[:, 3][bbox[:, 3] > h] = h
img_data=np.array(img_data,np.float32)
return img_data,bbox
3
对于训练时图片可以进行增强，进行旋转，明亮度变化等操作，使得网络认识不同状态的图片，扩充数据集。
#左右旋转
flip=np.random.random()<.5
if flip:
new_image=new_image.transpose(Image.FLIP_LEFT_RIGHT)
h,s,v=np.random.random(),np.random.random(),np.random.random()
#HSV调节
img=cv2.cvtColor(np.array(new_image),cv2.COLOR_RGB2HSV)
a=np.random.random()<.5
if a:
img[...,0]=img[...,0]*h
img[...,1]=img[...,1]*s
img[...,2]=img[...,2]*v
img=cv2.cvtColor(img,cv2.COLOR_HSV2RGB)*255
img=Image.fromarray(img)
img.show()
添加到get_data函数中
def get_data(self,annotations_file):
line=annotations_file.split()
img_data=Image.open(line[0])
'''
1.如果图片不是三通道的，那么需要转成三通道
'''
if np.shape(img_data)[2]!=3:
img_data = img_data.convert('RGB')
# 输出图片大小要一致，此时只是简单的resize一下
# img_data=img_data.resize((360,360))
'''
resize就会发生图片尺寸的变化，训练不会像原图那样好，bbox也会发生变化，该如何处理呢
图片缩放会导致失真，那么能不能不缩放也达到我们想要的尺寸呢，或者最小限度的失真
我们将图片的最大的长度保留下来，把短边用灰边填充
scale就是图片(ih,iw)变化成(h,w)的最小比例,(nh,nw)就是我们按照这个比例resize的图片大小
当然(nh,nw)和(h,w)有空白，空白我们就用灰条填充
'''
h,w=360,360#想要的尺寸
iw,ih=img_data.size
# 比列
scale = min(w / iw, h / ih)
nw = int(iw * scale)
nh = int(ih * scale)
dx = (w - nw) // 2
dy = (h - nh) // 2
img=img_data.resize((nw,nh),Image.BICUBIC)
new_image=Image.new('RGB', (w, h), (128, 128, 128))
new_image.paste(img,(dx,dy))
bbox=np.array([box.split(',') for box in line[1:]])
if len(bbox) > 0:
np.random.shuffle(bbox)
# 图片大小都是按照scale变化的，bbox也应该如此
bbox[:, [0, 2]] = bbox[:, [0, 2]] * nw / iw + dx
bbox[:, [1, 3]] = bbox[:, [1, 3]] * nh / ih + dy
bbox[:, 0:2][bbox[:, 0:2] < 0] = 0
bbox[:, 2][bbox[:, 2] > w] = w
bbox[:, 3][bbox[:, 3] > h] = h
img_data=np.array(img_data,np.float32)
# 只有训练时才会进行图片增强，测试和预测时不用
# slef.train可以自己定义传入，我就不写了
if not self.train:
return img_data,bbox
flip = np.random.random() < .5
if flip:
new_image = new_image.transpose(Image.FLIP_LEFT_RIGHT)
img_data = np.array(img_data, np.float32)
h, s, v = np.random.random(), np.random.random(), np.random.random()
img_data = cv2.cvtColor(np.array(new_image,np.float32), cv2.COLOR_RGB2HSV)
a = np.random.random() < .5
if a:
img[..., 0] = img[..., 0] * h
img[..., 1] = img[..., 1] * s
img[..., 2] = img[..., 2] * v
img_data = cv2.cvtColor(img, cv2.COLOR_HSV2RGB)
return img_data,bbox
到这儿基本的dataset已经完成了，当然还有其他的数据增强可以自己添加，后面可以写dataloader部分的代码了
1795
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
