轻松学PyTorch：自定义数据集制作与使用（ImageFolder/DataLoader）

点击上方“小白学视觉”，选择加"星标"或“置顶”

重磅干货，第一时间送达

大家好，这是轻松学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 = lines1011    def __len__(self):12        return len(self.landmarks_frame)1314    def num_of_samples(self):15        return len(self.landmarks_frame)1617    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 order23        h, w, c = img.shape24        # rescale25        img = cv.resize(img, (64, 64))26        img = (np.float32(img) /255.0 - 0.5) / 0.527        landmarks = np.zeros(10, dtype=np.float32)28        for i in range(1, len(contents), 2):29            landmarks[i - 1] = np.float32(contents[i]) / w30            landmarks[i] = np.float32(contents[i + 1]) / h31        landmarks = landmarks.astype('float32').reshape(-1, 2)32        # H, W C to C, H, W33        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("https://www.gofarlic.com/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 loader10for i_batch, sample_batched in enumerate(dataloader):11    print(i_batch, sample_batched['image'].size(), sample_batched['landmarks'].size())

运行显示如下：

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

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

好消息！

小白学视觉知识星球

开始面向外开放啦👇👇👇

下载1：OpenCV-Contrib扩展模块中文版教程 在「小白学视觉」公众号后台回复：扩展模块中文教程，即可下载全网第一份OpenCV扩展模块教程中文版，涵盖扩展模块安装、SFM算法、立体视觉、目标跟踪、生物视觉、超分辨率处理等二十多章内容。下载2：Python视觉实战项目52讲在「小白学视觉」公众号后台回复：Python视觉实战项目，即可下载包括图像分割、口罩检测、车道线检测、车辆计数、添加眼线、车牌识别、字符识别、情绪检测、文本内容提取、面部识别等31个视觉实战项目，助力快速学校计算机视觉。下载3：OpenCV实战项目20讲在「小白学视觉」公众号后台回复：OpenCV实战项目20讲，即可下载含有20个基于OpenCV实现20个实战项目，实现OpenCV学习进阶。  交流群 欢迎加入公众号读者群一起和同行交流，目前有SLAM、三维视觉、传感器、自动驾驶、计算摄影、检测、分割、识别、医学影像、GAN、算法竞赛等微信群（以后会逐渐细分），请扫描下面微信号加群，备注：”昵称+学校/公司+研究方向“，例如：”张三 + 上海交大 + 视觉SLAM“。请按照格式备注，否则不予通过。添加成功后会根据研究方向邀请进入相关微信群。请勿在群内发送广告，否则会请出群，谢谢理解~





计算机书童 为大家分享计算机、机器人领域的顶会顶刊论

免责声明：本文系网络转载或改编，未找到原创作者，版权归原作者所有。如涉及版权，请联系删