【Pytorch使用自制数据集,Dataloader】

数据集结构

话不多说,直接上核心代码

myDataset.py

python 复制代码
from collections import Counter
from torch.utils.data import Dataset
import os
from PIL import Image

class MyDataset(Dataset):
    """
    读取自制的数据集
    args:
        - image_dir: 图片的地址
        - label_dir: 标签的地址
        - name: 数据集的名称
        - transform: 数据集的预处理
    """
    def __init__(self, image_dir:str,  label_dir:str, name:str, transform=None):
        self.img_dir = os.path.join(image_dir, name)
        self.label_dir = os.path.join(label_dir, name)
        self.name = name
        self.image_path = os.listdir(self.img_dir)
        self.label_path = os.listdir(self.label_dir)
        self.transform = transform
    """
    读取数据集
    args:
        - index: 数据集的索引
    return:
        - image: 图片
        - label: 图片的标签
    """
    def __getitem__(self, index:int)->tuple:
        # 获取图片的地址
        image = self.image_path[index]
        image = os.path.join(self.img_dir, image)
        # 获取图像
        image = Image.open(image)
        # 如果不是彩色图像,将下面的注释解开可以转换成彩色图像,不过图片的模样改变很大
        # if image.mode!= 'RGB':
        #     image = image.convert('RGB')
        # 获取label的地址
        index_path = self.label_path[index]
        index_path = os.path.join(self.label_dir, index_path)
        label = self.parseTxt(index_path)
        if self.transform is not None:
            image = self.transform(image)
        return image, label
    
    """
    将txt文件解析成数字
    description:
        > 这里每个txt文件下可能有多个label,选出现最多的,如果你的txt里面只有一个label的话,想办法读取出来返回就行
    args:
        - label: txt文件的地址
    return:
        - label: 图片的标签
    """
    def parseTxt(self, label:str)->int:
        first_column = []
        with open(label, 'r') as f:
            for line in f.readlines():
                first_column.append(int(line.split()[0]))

        counter = Counter(first_column)
        return counter.most_common(1)[0][0]
    """
    获取数据集的长度
    """
    def __len__(self)->int:
        return len(self.image_path)

demo

train.py

python 复制代码
import torch
from torch.utils.data import DataLoader
from torchvision import transforms
# 导入加载数据集的类
from dataset import MyDataset
import os

root = os.path.join(os.getcwd(),'courseHomework','datasets')
transform = transforms.Compose([
    transforms.Resize((448, 448)),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    # transforms.Normalize((0.5), (0.5,))
])

train_dataset = MyDataset(root + '/images', root +'/labels', 'train', transform)
train_loader = DataLoader(train_dataset, batch_size=4, shuffle=False)

for step, data in enumerate(train_loader):
    imgs, labels = data
    print(imgs[0].shape)
    transforms.ToPILImage()(imgs[0]).show()
    break

大家结构和我不一样可以自由发挥

相关推荐
Jeremy_lf9 分钟前
【生成模型之三】ControlNet & Latent Diffusion Models论文详解
人工智能·深度学习·stable diffusion·aigc·扩散模型
桃花键神1 小时前
AI可信论坛亮点:合合信息分享视觉内容安全技术前沿
人工智能
野蛮的大西瓜1 小时前
开源呼叫中心中,如何将ASR与IVR菜单结合,实现动态的IVR交互
人工智能·机器人·自动化·音视频·信息与通信
CountingStars6192 小时前
目标检测常用评估指标(metrics)
人工智能·目标检测·目标跟踪
tangjunjun-owen2 小时前
第四节:GLM-4v-9b模型的tokenizer源码解读
人工智能·glm-4v-9b·多模态大模型教程
冰蓝蓝2 小时前
深度学习中的注意力机制:解锁智能模型的新视角
人工智能·深度学习
橙子小哥的代码世界2 小时前
【计算机视觉基础CV-图像分类】01- 从历史源头到深度时代:一文读懂计算机视觉的进化脉络、核心任务与产业蓝图
人工智能·计算机视觉
黄公子学安全2 小时前
Java的基础概念(一)
java·开发语言·python
新加坡内哥谈技术3 小时前
苏黎世联邦理工学院与加州大学伯克利分校推出MaxInfoRL:平衡内在与外在探索的全新强化学习框架
大数据·人工智能·语言模型
程序员一诺3 小时前
【Python使用】嘿马python高级进阶全体系教程第10篇:静态Web服务器-返回固定页面数据,1. 开发自己的静态Web服务器【附代码文档】
后端·python