轻松学Pytorch-自定义数据集制作与使用

  磐创AI分享   

来源 | OpenCV学堂

作者 | gloomyfish

大家好,这是轻松学Pytorch系列的第六篇分享,本篇你将学会如何从头开始制作自己的数据集,并通过DataLoader实现加载。本文以人脸Landmard五点的数据集标定与之制作为例来说明pytorch中如何实现自定义数据集读取与加载。

数据来源

首先要实现人脸landmark五点的数据标定,就得找到人脸数据,我使用的人脸数据是celebA数据集,大概有20W张多点,我从中选择了1000张,然后通过OpenCV写了个程序对人脸进行了简单的裁剪。然后还选择了一个音乐MV(上次就被人打call的宇少)通过opencv实现采集了一些人脸数据,这个数据的好处是有不同的光照,各种角度,丰富了数据的多样性。这些数据加起来1500张左右。图示如下:

Landmark标定

我这里选择的对得到1500张图像做数据标注,刚开始的选择标定工具都让我头疼,这个是我第一次标定一系列的点,经过一番尝试之后,终于发现一个很好用的工具,同时支持人脸检测与五点标定。贴上地址:

https://github.com/Mukosame/Face-Annotation-Tool

废话也不多说了,只说一句话,简单靠谱,然后我就对这个1500张图像进行五点标定,本来我可以不这么干的,我可以用其它的模型来直接找这些图像的landmark五点然后生成文件即可,但是我还是决定手动标注一番。结果让我眼睛疼了两天之后,终于给标注好拉,发誓以后再也不干这种活了,我太难了。截图如下:

现在自定义数据已经准备完毕,下面就应该是pytorch登场了。

自定义数据集实现

基于Pytorch中的torch.utils.data.Dataset类实现自定义的FaceLandmarksDataset类,主要是重写了getitem这个方法。完整的代码实现如下:

 1class FaceLandmarksDataset(Dataset):
2 def __init__(self, txt_file):
3 self.transform = transforms.Compose([transforms.ToTensor()])
4 lines = []
5 with open(txt_file) as read_file:
6 for line in read_file:
7 line = line.replace('\n', '')
8 lines.append(line)
9 self.landmarks_frame = lines
10
11 def __len__(self):
12 return len(self.landmarks_frame)
13
14 def num_of_samples(self):
15 return len(self.landmarks_frame)
16
17 def __getitem__(self, idx):
18 if torch.is_tensor(idx):
19 idx = idx.tolist()
20 contents = self.landmarks_frame[idx].split('\t')
21 image_path = contents[0]
22 img = cv.imread(image_path) # BGR order
23 h, w, c = img.shape
24 # rescale
25 img = cv.resize(img, (64, 64))
26 img = (np.float32(img) /255.0 - 0.5) / 0.5
27 landmarks = np.zeros(10, dtype=np.float32)
28 for i in range(1, len(contents), 2):
29 landmarks[i - 1] = np.float32(contents[i]) / w
30 landmarks[i] = np.float32(contents[i + 1]) / h
31 landmarks = landmarks.astype('float32').reshape(-1, 2)
32 # H, W C to C, H, W
33 img = img.transpose((2, 0, 1))
34 sample = {'image': torch.from_numpy(img), 'landmarks': torch.from_numpy(landmarks)}
35 return sample

加载与显示

实现了自定义的Dataset类之后,就可以通过自定义的Dataset来构建一个DataLoader对象实现数据的加载跟批次处理,对自定义的dataset完成测试。代码如下:

 1ds = FaceLandmarksDataset("D:/facedb/Face-Annotation-Tool/landmark_output.txt")
2for i in range(len(ds)):
3 sample = ds[i]
4 print(i, sample['image'].size(), sample['landmarks'].size())
5 if i == 3:
6 break
7
8dataloader = DataLoader(ds, batch_size=4, shuffle=True, num_workers=4)
9# data loader
10for i_batch, sample_batched in enumerate(dataloader):
11 print(i_batch, sample_batched['image'].size(), sample_batched['landmarks'].size())

运行显示如下:

关注我们,后台输入关键字 landmark 获取本人辛苦标注的数据集。

记得点赞支持,这是本人继续写下去的动力!

更多资料:

PyTorch 中文文档: http://pytorch123.com/

GitHub : https://github.com/fendouai/PyTorchDocs

- End -

✄------------------------------------------------

看到这里,说明你喜欢这篇文章,请点击「 在看 」或顺手「 转发 」「 点赞 」。

欢迎微信搜索「 panchuangxx 」,添加小编 磐小小仙 微信,每日朋友圈更新一篇高质量推文(无广告),为您提供更多精彩内容。

▼       扫描二维码添加小编   ▼    ▼  

我来评几句
登录后评论

已发表评论数()

相关站点

+订阅
热门文章