数据集处理def __getitem__(self,index)
从前面缓存中进行读取(这里先不说mosaic数据增强)。通过索引index加载图像,获取图像的尺寸,通过letterbox对图像进行reshape,使其适合网络输入大小。将labels中的box信息从xywh转为xyxy(中心点,w,h--->左上角和右下角)【可以用来做数据增强处理】。
# Load image img, (h0, w0), (h, w) = load_image(self, index) # Letterbox shape = self.batch_shapes[self.batch[index]] if self.rect else self.img_size # final letterboxed shape img, ratio, pad = letterbox(img, shape, auto=False, scaleup=self.augment) shapes = (h0, w0), ((h / h0, w / w0), pad) # for COCO mAP rescaling labels = self.labels[index].copy() if labels.size: # normalized xywh to pixel xyxy format labels[:, 1:] = xywhn2xyxy(labels[:, 1:], ratio[0] * w, ratio[1] * h, padw=pad[0], padh=pad[1])
数据增强后再转为xywh(不增强也要转,我这里是以不增强为例)。
nL = len(labels) # number of labels if nL: labels[:, 1:5] = xyxy2xywh(labels[:, 1:5]) # convert xyxy to xywh labels[:, [2, 4]] /= img.shape[0] # normalized height 0-1 labels[:, [1, 3]] /= img.shape[1] # normalized width 0-1
创建一个labels_out 用来存储labels
labels_out = torch.zeros((nL, 6))
这里再说一下,前面我们得到的labels是shape为(1,5),这里的labels_out是(1,6)在最前面多了一列。
图像shape的转换:
img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416
img = np.ascontiguousarray(img)
最终返回的是图像img(转张量),labels_out,img_files[index]【对应img的路径】,shapes
return torch.from_numpy(img), labels_out, self.img_files[index], shapes
获得的dataset形式
通过LoadImagesAndLabels可以获得我们上述处理的数据集。可以看下现在的dataset都包含哪些内容。【里面的segments是用来做图像分割的,我这里没有,所以是空】
获得batch_size和number workers 。DistributedSampler是分布式采样器。可以将数据集放在多卡上进行训练,但我这里用的是单卡。每个GPU上数据的划分为:一个epoch数据集数量/num_gpu=每个GPU得到的数量,如果shuffle=True,那么每个GPU得到的数据是随机的,否则是按顺序划分。
batch_size = min(batch_size, len(dataset)) nw = min([os.cpu_count() // world_size, batch_size if batch_size > 1 else 0, workers]) # number of workers sampler = torch.utils.data.distributed.DistributedSampler(dataset) if rank != -1 else None loader = torch.utils.data.DataLoader if image_weights else InfiniteDataLoader # Use torch.utils.data.DataLoader() if dataset.properties will update during training else InfiniteDataLoader() dataloader = loader(dataset, batch_size=batch_size, num_workers=nw, sampler=sampler, pin_memory=True, collate_fn=LoadImagesAndLabels.collate_fn4 if quad else LoadImagesAndLabels.collate_fn) return dataloader, dataset
dataloader
继承DataLoader可以写自己的loader函数,或者是直接调用torch中的dataloader也可以。主要放入:dataset,batch_size,num_workers,pin_memory,以及自己实现的collate_fn函数。
dataloader = loader(dataset, batch_size=batch_size, num_workers=nw, sampler=sampler, pin_memory=True, collate_fn=LoadImagesAndLabels.collate_fn4 if quad else LoadImagesAndLabels.collate_fn)
dataloader内容:
以上就是yolov7中数据集处理以及加载过程,还有一些细节后面再补充。