上章节讲到了如何将要数据模型转换成网络训练的时候方便读取的数据格式，本文讲解下pytorch模块如何使用这个转换后的数据代码块如下：
#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File : datasets.py
@Time : 2021/08/19 16:37:10
@Author : XIA Yan
@Contact : 2440212215@qq.com
@Version : 0.1
@License : Apache License Version 2.0, January 2004
@Language: python3.8
@Desc : torch的训练数据生成
'''
import glob
import random
import numpy as np
from PIL import Image
import torch
import json
import torch.nn.functional as F
from torch.utils.data import Dataset
import torchvision.transforms as transforms
def pad_to_square(img, pad_value = 0):
c, h, w = img.shape
dim_diff = np.abs(h - w)
# (upper / left) padding and (lower / right) padding
pad1, pad2 = dim_diff // 2, dim_diff - dim_diff // 2
# Determine padding
pad = (0, 0, pad1, pad2) if h <= w else (pad1, pad2, 0, 0)
# Add padding
img = F.pad(img, pad, "constant", value=pad_value)
return img, pad
def resize(image, size):
image = F.interpolate(image.unsqueeze(0), size=size, mode="nearest").squeeze(0)
return image
class ImageFolder(Dataset):
def __init__(self, folder_path, img_size=416):
self.files = sorted(glob.glob("%s/*.*" % folder_path))
self.img_size = img_size
def __getitem__(self, index):
img_path = self.files[index % len(self.files)]
# 图片解析成pytorch格式的tensor
img = transforms.ToTensor()(Image.open(img_path).convert("RGB"))
# Pad to square resolution
img, _ = pad_to_square(img, 0)
# Resize
img = resize(img, self.img_size)
return img_path, img
def __len__(self):
return len(self.files)
class ListDataset(Dataset):
'''
@description:
加载图像和标签在train阶段或者test阶段
@Args:
json_path :(string) 数据的json格式，由createdata/create_targets.py产生
img_size :(int) 图像尺寸默认416
augment :(bool) 是否采用数据增强方式
@Return:
img_path, img, targets
'''
def __init__(self, json_path, img_size = 416, augment = None):
with open(json_path, "r") as fp:
self.json_dict = json.load(fp) #读取json中字典
self.img_size = img_size
self.image_path = list(self.json_dict.keys()) #提取json文件中的image路径
self.augment = augment
def __getitem__(self, index):
# --------
# Image
# --------
img_path = self.image_path[index % len(self.image_path)] #获取当前index的图像路径
labels = np.array(self.json_dict[img_path]) #获取当前index图像的targets
img = transforms.ToTensor()(Image.open(str(img_path)).convert('RGB'))
_, h, w = img.shape
img, pad = pad_to_square(img, 0)
_, padded_h, padded_w = img.shape
# ---------
# Label
# ---------
# 解析坐标(原标记坐标是非normal的[x1,y1,x2,y2]坐标)
boxes = torch.from_numpy(labels.reshape(-1, 5))
x1 = boxes[:,1]
y1 = boxes[:,2]
x2 = boxes[:,3]
y2 = boxes[:,4]
# Adjust for added padding
x1 += pad[0]
y1 += pad[2]
x2 += pad[1]
y2 += pad[3]
# Returns (x, y, w, h)
cx = (x1 + x2) / 2
cy = (y1 + y2) / 2
w = x2 - x1
h = y2 - y1
# normal coords
boxes[:, 1] = cx / padded_w
boxes[:, 2] = cy / padded_h
boxes[:, 3] = w / padded_w
boxes[:, 4] = h / padded_h
targets = torch.zeros((len(boxes), 6))
targets[:, 1:] = boxes
# TODO 数据增强
if self.augment:
pass
return img_path, img, targets
def collate_fn(self, batch):
'''
@description:
重写collate_fn方法，用于处理每个batch信息
@Args:
batch:
@Return:
返回经过batch处理后的，图片路径，图片，targets
'''
paths, imgs, targets = list(zip(*batch))
for i, boxes in enumerate(targets):
boxes[:, 0] = i
imgs = torch.stack([resize(img, self.img_size) for img in imgs])
return paths, imgs, torch.cat(targets, 0)
def __len__(self):
return len(self.image_path)
if __name__ == '__main__':
dataset = ListDataset('data/data.json')
print(dataset[0])
pytorch中给网络模型产生数据的主要模块是torch.utils.data中的Dataset
from torch.utils.data import Dataset
首先使用如下代码讲之前的转换的json数据进行解析
def __init__(self, json_path, img_size = 416, augment = None):
with open(json_path, "r") as fp:
self.json_dict = json.load(fp) #读取json中字典
self.img_size = img_size
self.image_path = list(self.json_dict.keys()) #提取json文件中的image路径
self.augment = augment
self.image_path是获取了所有的图像路径
_, h, w = img.shape # 获取图像的通道数、高和宽 img, pad = pad_to_square(img, 0) # 将图像填充为正方形 _, padded_h, padded_w = img.shape # 获取填充后的数据格式
下面给出本模块的示意图：
图中将本模块的大致功能讲解了一下，通过读取图像数据和对应的图像标签数据，需要将数据转换成tensor格式，因为图像数据标注的时候有长宽比例不一致的情况，因此在这里需要将图像规范为正方形，那这边就牵涉到图像的正方形填充。填充完毕后需要记录下填充后的多少个内容，这样在将图像映射回原图的时候能够正确绘出矩形框。
这里顺便提示下，由于我之前参加一个项目是利用红外摄像头读取数据进行检测，然后读取检测后的图像温度，红外一起的尺寸为640*512的分辨率。采集数据和实际使用数据都是这个分辨率，这个项目中我并没有使用图像的正方形化，训练的时候就是用得640*512尺寸。这样的好处是省去了图像生成和推理的时间。因此我们在知道为什么需要填充成正方形的原因才能知道什么情况下不使用。算法和技巧并不是一成不变的，需要灵活应用。
