dataset.py
torch.utils.data分为三部分：Dataset、DataLoader、sampler.Sampler
一、Dataset部分
1、先看准备部分
① def get_hash(files): # 输入[list of files_path]，返回所有文件的大小总和
return sum(os.path.getsize(f) for f in files if os.path.isfile(f))
② def cache_labels： 输入[list of txts_path] [list of imgs_path]输出labels的cache.
values = {'./crop1024/images/1.jpg':[array( [对应1.txt的labels信息] ，dtype=float32), (weights, heights))]，……，'hash':int_value}
def cache_labels(self, path='labels.cache'): values = {} for (img, label) in zip(self.img_files, self.label_files): ls = [] image = Image.open(img) image.verify() # PIL verify shape = img.size assert (shape[0] > 9) & (shape[1] > 9), 'image size <10 pixels' if os.path.isfile(label): with open(label, 'r') as f: ls = np.array([x.split() for x in f.read().splitlines()], dtype=np.float32) # labels if len(ls) == 0: # 当labels文件中内容为空时也要确保shape一致 ls = np.zeros((0, 6), dtype=np.float32) values[img] = [ls, shape] values['hash'] = get_hash(self.label_files + self.img_files) torch.save(values, path) # save for next time return values
③ def load_image(self, index):
input:self.img_files[index]图片路径； output：img，hw_original, hw_resized; img是cv2直接读出的.shape=(height, weight, 3), 关于resize：
cv2.INTER_AREA和cv2.INTER_LINEAR是两种无关的参数，resize后仍然为(w0:h0)的比例，但长边为self.img_size。
def load_image(self, index): path = self.img_files[index] img = cv2.imread(path) assert img is not None, 'Image Not Found ' + path h0, w0 = img.shape[:2] # orig hw r = self.img_size / max(h0, w0) # resize image to img_size if r != 1: # always resize down, only resize up if training with augmentation interp = cv2.INTER_AREA if r < 1 and not self.augment else cv2.INTER_LINEAR img = cv2.resize(img, (int(w0 * r), int(h0 * r)), interpolation=interp) return img, (h0, w0), img.shape[:2] # img, hw_original, hw_resized
④ def load_mosaic(self, index): mosaic数据增强, 这里只是把index和其他随机的三张图片拼接在一起(同时修改label)，数据增强部分由random_perspective实现。
input：self.labels, self.img_size; output: img4.size=[resized_height,resized_ width, 3], labels4.shape=(img4中的目标GT数量, [classid ,LT_x,LT_y,RB_x,RB_y,Θ])
注意这里img4和label4中都是用的上一步load_image中的resized_height和resized_width, label中的xywh格式转换为xyxy格式。
def load_mosaic(self, index): labels4 = [] s = self.img_size yc, xc = [int(random.uniform(-x, 2 * s + x)) for x in self.mosaic_border] # mosaic center x, y indices = [index] + [random.randint(0, len(self.labels) - 1) for _ in range(3)] # 3 additional image indices for i, index in enumerate(indices): img, _, (h, w) = load_image(self, index) # img.size = [resized_height,resized_ width, 3] # place img in img4 if i == 0: # top left img4 = np.full((s * 2, s * 2, img.shape[2]), 114, dtype=np.uint8) # base image with 4 tiles # xmin, ymin, xmax, ymax (large image) 用于确定原图片在img4左上角的坐标（左上右下） x1a, y1a, x2a, y2a = max(xc - w, 0), max(yc - h, 0), xc, yc # xmin, ymin, xmax, ymax (small image) 用于确定原图片剪裁进img4中的图像内容范围 x1b, y1b, x2b, y2b = w - (x2a - x1a), h - (y2a - y1a), w, h elif i == 1: # top right x1a, y1a, x2a, y2a = xc, max(yc - h, 0), min(xc + w, s * 2), yc x1b, y1b, x2b, y2b = 0, h - (y2a - y1a), min(w, x2a - x1a), h elif i == 2: # bottom left x1a, y1a, x2a, y2a = max(xc - w, 0), yc, xc, min(s * 2, yc + h) x1b, y1b, x2b, y2b = w - (x2a - x1a), 0, w, min(y2a - y1a, h) elif i == 3: # bottom right x1a, y1a, x2a, y2a = xc, yc, min(xc + w, s * 2), min(s * 2, yc + h) x1b, y1b, x2b, y2b = 0, 0, min(w, x2a - x1a), min(y2a - y1a, h) # img4.size = [resized_height,resized_ width, 3] img4[y1a:y2a, x1a:x2a] = img[y1b:y2b, x1b:x2b] # img4[ymin:ymax, xmin:xmax] padw = x1a - x1b # 原图片未剪裁进img4中的宽度 padh = y1a - y1b # 原图片未剪裁进img4中的高度 x = self.labels[index] labels = x.copy() if x.size > 0: # Normalized xywh to pixel xyxy format labels[:, 1] = w * (x[:, 1] - x[:, 3] / 2) + padw # Left_top_x labels[:, 2] = h * (x[:, 2] - x[:, 4] / 2) + padh # Left_top_y labels[:, 3] = w * (x[:, 1] + x[:, 3] / 2) + padw # right_bottom_x labels[:, 4] = h * (x[:, 2] + x[:, 4] / 2) + padh # right_bottom_y labels4.append(labels) # Concat/clip labels if len(labels4): labels4 = np.concatenate(labels4, 0) # 将第一个维度shape=4取消 np.clip(labels4[:, 1:5], 0, 2 * s, out=labels4[:, 1:5]) # 限定labels4[:, 1:5]中最小值只能为0，最大值只能为2*self.size # Augment img4, labels4 = random_perspective(img4, labels4, degrees=self.hyp['degrees'], translate=self.hyp['translate'], scale=self.hyp['scale'], shear=self.hyp['shear'], perspective=self.hyp['perspective'], border=self.mosaic_border) # border to remove return img4, labels4
2、class LoadImagesAndLabels(Dataset)：
先看函数调用情况：
def create_dataloader(path, imgsz, batch_size, stride, opt, hyp=None, augment=False, cache=False, pad=0.0, rect=False, rank=-1, world_size=1, workers=8) dataset = LoadImagesAndLabels(path, imgsz, batch_size, augment=augment, # augment images hyp=hyp, # augmentation hyperparameters rect=rect, # rectangular training cache_images=cache, single_cls=opt.single_cls, stride=int(stride), pad=pad, rank=rank) # ---- train.py ----- dataloader, dataset = create_dataloader(train_path, imgsz, batch_size, gs, opt, hyp=hyp, augment=True, cache=opt.cache_images, rect=opt.rect,
rank=rank, world_size=opt.world_size, workers=opt.workers)
parser.add_argument('--rect', action='store_true', help='rectangular training')
self.rect = False if image_weights else rect # False
self.mosaic = self.augment and not self.rect # load 4 images at a time into a mosaic (only during training - True)
self.mosaic_border = [-img_size // 2, -img_size // 2]
①def __init__(self, path, img_size=640, batch_size=16, augment=False, hyp=None, rect=False, image_weights=False,
cache_images=False, single_cls=False, stride=32, pad=0.0, rank=-1):
f = [] for p in path if isinstance(path, list) else [path]: if os.path.isdir(p): # folder f += glob.iglob(p + os.sep + '*.*') self.img_files = sorted([x.replace('/', os.sep) for x in f if os.path.splitext(x)[-1].lower() in img_formats]) sa, sb = os.sep + 'images' + os.sep, os.sep + 'labels' + os.sep # sa=/images/, sb=/labels/ self.label_files = [x.replace(sa, sb, 1).replace(os.path.splitext(x)[-1], '.txt') for x in self.img_files] self.batch = np.floor(np.arange(n) / batch_size).astype(np.int) cache_path = str(Path(self.label_files[0]).parent) + '.cache' # cached labels if os.path.isfile(cache_path): cache = torch.load(cache_path) # load if cache['hash'] != get_hash(self.label_files + self.img_files): # dataset changed cache = self.cache_labels(cache_path) # re-cache else: cache = self.cache_labels(cache_path) labels, shapes = zip(*[cache[x] for x in self.img_files]) self.shapes = np.array(shapes, dtype=np.float64) self.labels = list(labels)
path:图片所在文件夹路径，可以是一个str，也可以多个路径str list
self.img_files = ['/**/images/**/1.jpg', '/**/images/**/2.png',…] self.label_files = ['/**/labels/**/1.txt', '/**/images/**/2.txt'，…]
self.batch记录每一张图片所在的batch_id，ndarray.size=(num_img)
self.shapes记录每一张图片的(width,height), ndarray.size=(num_img)
self.labels记录每一张图片的.txt-label信息，type is list(num_img);其中某张图片的信息self.labels[i].shape=(num_box, 6) [class, xywh_center(归一化),Θ])
nb = self.batch[-1] + 1
if self.rect:
# Sort by aspect ratio 按纵横比的数值从小到大重新进行排序，矩形训练通常以成批处理 s = self.shapes # wh ar = s[:, 1] / s[:, 0] # aspect ratio, h/w irect = ar.argsort() self.img_files = [self.img_files[i] for i in irect] self.label_files = [self.label_files[i] for i in irect] self.labels = [self.labels[i] for i in irect] self.shapes = s[irect] # wh ar = ar[irect] # Set training image shapes shapes = [[1, 1]] * nb for i in range(nb): ari = ar[bi == i] mini, maxi = ari.min(), ari.max() if maxi < 1: shapes[i] = [maxi, 1] elif mini > 1: shapes[i] = [1, 1 / mini] self.batch_shapes = np.ceil(np.array(shapes) * img_size / stride + pad).astype(np.int) * stride
这步为Rectangular Training作准备，默认是不执行的。
self.batch_shapes.shape=(num_batch,2),这一步是为batch中的图片计算出一个统一的shape，要能代表batch内图片的总体情况。
② def __getitem__(self, index)
整个过程是对self.img_files[index]和self.labels[index]的变换，其中也用到hyp等数据增强的参数！
def __getitem__(self, index): if self.image_weights: # False, 取参数初始化默认值 index = self.indices[index] hyp = self.hyp mosaic = self.mosaic and random.random() < hyp['mosaic'] # 1 ,True and True if mosaic: # True img, labels = load_mosaic(self, index) shapes = None else: img, (h0, w0), (h, w) = load_image(self, index) # Letterbox, 如果进行矩形训练，则获取每个batch的输入图片的shape 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) labels = [] x = self.labels[index] if x.size > 0: # Normalized format xywh to xyxy labels = x.copy() labels[:, 1] = ratio[0] * w * (x[:, 1] - x[:, 3] / 2) + pad[0] # pad width labels[:, 2] = ratio[1] * h * (x[:, 2] - x[:, 4] / 2) + pad[1] # pad height labels[:, 3] = ratio[0] * w * (x[:, 1] + x[:, 3] / 2) + pad[0] labels[:, 4] = ratio[1] * h * (x[:, 2] + x[:, 4] / 2) + pad[1] if self.augment: if not mosaic: # Augment imagespace # 随机对图片进行旋转，平移，缩放，裁剪 -- train中不执行 img, labels = random_perspective(img, labels, degrees=hyp['degrees'], translate=hyp['translate'], scale=hyp['scale'], shear=hyp['shear'], perspective=hyp['perspective']) # Augment colorspace 随机改变图片的色调（H），饱和度（S），亮度（V） augment_hsv(img, hgain=hyp['hsv_h'], sgain=hyp['hsv_s'], vgain=hyp['hsv_v']) # labels.size = (目标数量, [class, xyxy, Θ]), convert xyxy to xywh , 归一化标签 [0.0~1.0] labels[:, 1:5] = xyxy2xywh(labels[:, 1:5]) labels[:, [2, 4]] /= img.shape[0] labels[:, [1, 3]] /= img.shape[1] if self.augment: # flip up-down 上下翻转 沿x轴翻转 （y变x不变）, θ根据左右偏转也进行改变, 原θ=0时，情况特殊不做改变 if random.random() < hyp['flipud']: # 配置文件默认值为0.5 img = np.flipud(img) labels[:, 2] = 1 - labels[:, 2] labels[:, -1] = 180 - labels[:, -1] labels[labels[:, -1] == 180, -1] = 0 # flip left-right 左右翻转 沿y轴翻转（ x变y不变）,θ根据左右偏转也进行改变, 原θ=0时，情况特殊不做改变 if random.random() < hyp['fliplr']: # 配置文件默认值为0.5 img = np.fliplr(img) labels[:, 1] = 1 - labels[:, 1] labels[:, -1] = 180 - labels[:, -1] labels[labels[:, -1] == 180, -1] = 0 # img.size=[resized_height,resized_width,3] -> [3, resized_height, resized_width] img = img[:, :, ::-1].transpose(2, 0, 1) img = np.ascontiguousarray(img) labels_out = torch.zeros((len(labels), 7)) labels_out[:, 1:] = torch.from_numpy(labels) return torch.from_numpy(img), labels_out, self.img_files[index], shapes # train -> mosaic==True -> shape=None
load_mosaic：经过load_image取出cv2的图片并resize成(resized_height, resized_weight), 这在load_mosaic中没有改变 并且把label由[0~1]转成[r_h,r_w]。
进行了4张图片和label的拼接，xywh格式改成了xyxy格式，然后是random_perspective做数据增强。
augment：augment_hsv(img, 数据增强调整色调、亮度、饱和度； 转格式xyxy -> xywh, 并label归一化(resized_height, resized_weight)； np.flipud(img)， np.fliplr(img)
两种翻转按一定概率执行！
最有由cv2的(height,width,3) -> (3,height,width)。
③ def __len__(self):
return len(self.img_files)
@staticmethod def collate_fn(batch): img, label, path, shapes = zip(*batch) for i, l in enumerate(label): l[:, 0] = i return torch.stack(img, 0), torch.cat(label, 0), path, shapes
④ def conllate_fn(batch):
@staticmethod def collate_fn(batch): img, label, path, shapes = zip(*batch) for i, l in enumerate(label): l[:, 0] = i return torch.stack(img, 0), torch.cat(label, 0), path, shapes
返回img.size = (batch_size, 3 , resized_height, resized_width) 没有归一化;
labels.size=(batch中的目标数量, [图片在当前batch中的索引,classid,归一化后的xywh, Θ])
path.size=batch_size, 该batch中所有image的路径
浙公网安备 33010602011771号