本文基于Pytorch框架，制作自己的Dataset及DataLoader。完成代码，请移步Go-for-i。
开始步骤：
1.准备自己需要分类的数据。该文以二分类为例，其中，0_image和1_image中分别存放自己的数据集。
2. 生成train.txt文件和test.txt文件。其中，两个.txt文件中存放，图片的路径以及图片所对应的label。具体代码及生成的结果如下：
train_ratio = 0.9
test_ratio = 1-train_ratio
rootdata = r"dataset"
train_list, test_list = [], []
data_list = []
class_flag = -1
for a, b, c in os.walk(rootdata):
for i in range(0, int(len(c)*train_ratio)):
train_data = os.path.join(a, c[i])+'t'+str(class_flag)+'n'
train_list.append(train_data)
for i in range(int(len(c) * train_ratio), len(c)):
test_data = os.path.join(a, c[i]) + 't' + str(class_flag)+'n'
test_list.append(test_data)
class_flag += 1
print(train_list)
random.shuffle(train_list)
random.shuffle(test_list)
with open('train.txt', 'w', encoding='UTF-8') as f:
for train_img in train_list:
f.write(str(train_img))
with open('test.txt', 'w', encoding='UTF-8') as f:
for test_img in test_list:
f.write(test_img)
3.对图片进行处理，放入DataLoader中。具体代码如下：
# 图像标准化
transform_BZ = transforms.Normalize(
mean=[0.5, 0.5, 0.5], # 取决于数据集
std=[0.5, 0.5, 0.5]
)
class LoadData(Dataset):
def __init__(self, txt_path, train_flag=True):
self.imgs_info = self.get_images(txt_path)
self.train_flag = train_flag
self.train_tf = transforms.Compose([
transforms.Resize(224),
transforms.RandomHorizontalFlip(),
transforms.RandomVerticalFlip(),
transforms.ToTensor(),
transform_BZ
])
self.val_tf = transforms.Compose([
transforms.Resize(224),
transforms.ToTensor(),
transform_BZ
])
def get_images(self, txt_path):
with open(txt_path, 'r', encoding='utf-8') as f:
imgs_info = f.readlines()
imgs_info = list(map(lambda x: x.strip().split('t'), imgs_info))
return imgs_info
def padding_black(self, img):
w, h = img.size
scale = 224. / max(w, h)
img_fg = img.resize([int(x) for x in [w * scale, h * scale]])
size_fg = img_fg.size
size_bg = 224
img_bg = Image.new("RGB", (size_bg, size_bg))
img_bg.paste(img_fg, ((size_bg - size_fg[0]) // 2,
(size_bg - size_fg[1]) // 2))
img = img_bg
return img
def __getitem__(self, index):
img_path, label = self.imgs_info[index]
img = Image.open(img_path)
img = img.convert('RGB')
img = self.padding_black(img)
if self.train_flag:
img = self.train_tf(img)
else:
img = self.val_tf(img)
label = int(label)
return img, label
def __len__(self):
return len(self.imgs_info)
if __name__ == "__main__":
train_dataset = LoadData("train.txt", True)
print("训练数据个数：", len(train_dataset))
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
batch_size=10,
shuffle=True)
for image, label in train_loader:
print(image.shape)
print(image)
# img = transform_BZ(image)
# print(img)
print(label)
test_dataset = LoadData("test.txt", False)
print("测试数据个数：", len(test_dataset))
test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
batch_size=10,
shuffle=False)
for image, label in test_loader:
print(image.shape)
print(label)
欢迎交流、沟通、批评、指正。最后，别忘了点赞哦！
(本文仅为记录自己的学习过程，无其他用途)
