pytorch构建自己的数据集

当ImageFolder无法满足需求时,可以自定义dataset类。核心在于getitem函数,它根据索引返回图像数据和标签。通过索引获取图片路径,使用PIL加载并用torchvision转换为Tensor,调整大小并归一化。标签从文件名中解析并编码。
摘要由CSDN通过智能技术生成

ImageFolder可以使用的情况是,比如猫狗识别中,train一个文件夹下已经将猫和狗分为两个不同的文件夹了,那我们可以直接使用ImageFolder来包装成数据集。
但有时候遇到的数据集是train一个文件夹下包含了所有的猫狗图片,那我们就无法使用ImageFolder函数,此时我们可以自己构造一个dataset类。
核心是getitem函数,这个函数的主要功能是根据样本的索引,返回索引对应的一张图片的图像数据X与对应的标签Y,也就是返回一个对应的训练样本。
具体而言,getitem的实现思路比较简单,将索引idx转换为图片的路径,然后用PIL的Image包来读取图片数据,然后将数据用torchvision的transforms转换成tensor并且进行Resize来统一大小(给出的图片尺寸不一致)与归一化,这样一来就可以得到图像数据了。因为训练集中图片的文件名上面带有猫狗的标签,所以标签可以通过对图片文件名split后得到然后转成0,1编码。

class MyDataset(Dataset):
    def __init__(self,data_path:str,train=True,transform=None):
        self.data_path = data_path
        self.train_flag = train
        if transform is None:([    ……
            ])
        else:
            self.transform = transform
        self.path_dir = os.listdir(data_path)
    def __getitem__(self, idx:int):
        img_path = self.path_dir[idx]
 
  • 0
    点赞
  • 5
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值