专栏首页贾志刚-OpenCV学堂轻松学Pytorch-自定义数据集制作与使用

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

大家好,这是轻松学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这个方法。完整的代码实现如下:

class FaceLandmarksDataset(Dataset):
    def __init__(self, txt_file):
        self.transform = transforms.Compose([transforms.ToTensor()])
        lines = []
        with open(txt_file) as read_file:
            for line in read_file:
                line = line.replace('\n', '')
                lines.append(line)
        self.landmarks_frame = lines

    def __len__(self):
        return len(self.landmarks_frame)

    def num_of_samples(self):
        return len(self.landmarks_frame)

    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()
        contents = self.landmarks_frame[idx].split('\t')
        image_path = contents[0]
        img = cv.imread(image_path)  # BGR order
        h, w, c = img.shape
        # rescale
        img = cv.resize(img, (64, 64))
        img = (np.float32(img) /255.0 - 0.5) / 0.5
        landmarks = np.zeros(10, dtype=np.float32)
        for i in range(1, len(contents), 2):
            landmarks[i - 1] = np.float32(contents[i]) / w
            landmarks[i] = np.float32(contents[i + 1]) / h
        landmarks = landmarks.astype('float32').reshape(-1, 2)
        # H, W C to C, H, W
        img = img.transpose((2, 0, 1))
        sample = {'image': torch.from_numpy(img), 'landmarks': torch.from_numpy(landmarks)}
        return sample

加载与显示

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

ds = FaceLandmarksDataset("D:/facedb/Face-Annotation-Tool/landmark_output.txt")
for i in range(len(ds)):
    sample = ds[i]
    print(i, sample['image'].size(), sample['landmarks'].size())
    if i == 3:
        break

dataloader = DataLoader(ds, batch_size=4, shuffle=True, num_workers=4)
# data loader
for i_batch, sample_batched in enumerate(dataloader):
    print(i_batch, sample_batched['image'].size(), sample_batched['landmarks'].size())

运行显示如下:

本文分享自微信公众号 - OpenCV学堂(CVSCHOOL),作者:gloomyfish

原文出处及转载信息见文内详细说明,如有侵权,请联系 yunjia_community@tencent.com 删除。

原始发表时间:2020-05-28

本文参与腾讯云自媒体分享计划,欢迎正在阅读的你也加入,一起分享。

我来说两句

0 条评论
登录 后参与评论

相关文章

  • 深度学习AI美颜系列----AI美发算法(美妆相机/天天P图染发特效)

    给照片或者视频中的人物头发换颜色,这个技术已经在手机app诸如天天P图,美图秀秀等应用中使用,并获得了不少用户的青睐。如何给照片或者视频中的人物头发换发色?换发...

    OpenCV学堂
  • 手把手教你用keras搭建GAN

    github:https://github.com/chenyang1999/KerasGAN/blob/master/gan/gan.py

    OpenCV学堂
  • 轻松学Pytorch – 年龄与性别预测

    大家好,上周太忙,没有更新Pytorch轻松学系列文章,但是我还是会坚定的继续走下去的,所谓有始有终,这个系列我会一直坚持写下去,希望大家继续支持我,积极给我反...

    OpenCV学堂
  • vue -- 动态加载组件 (tap 栏效果)

    在 vue 中,实现 Tab 切换主要有三种方式:使用动态组件,使用 vue-router 路由,使用第三方插件。

    小蔚
  • Vue.js如何划分组件

    王小婷
  • 环境配置:React Native 开发环境配置 For Android

    React Native 是FaceBook开源的一个项目,FaceBook希望可以用写 Web App 的方式去写 Native App。它可以让我们用JS和...

    非著名程序员
  • 阐述Session加载实体对象的过程

    Session加载实体对象的步骤是: ① Session在调用数据库查询功能之前,首先会在一级缓存中通过实体类型和主键进行查找,如果一级缓存查找命中且数据状态...

    唐怀瑟
  • Python的dict实现原理及与Java的比较探究

    Python内部很地方都使用着dict这种结构,在对象属性dict就是一个字典,所以对其效率要求很高。 dict采用了哈希表,最低能在 O(1)时间内完成搜索。...

    Python中文社区
  • JavaScript 中如何进行异步编程

    JavaScript语言的一大特点就是单线程,也就是说,同一个时间只能做一件事。那么,为什么JavaScript不能有多个线程呢?这样能提高效率啊。

    江米小枣
  • C语言字符串指针

    运行结果: http://c.biancheng.net http://c.biancheng.net 字符数组归根结底还是一个数组,上节讲到的关于指针...

    心跳包

扫码关注云+社区

领取腾讯云代金券