【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

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

相关推荐
凯禾瑞华养老实训室1 小时前
人才教育导向下:老年生活照护实训室助力提升学生老年照护服务能力
人工智能
luckys.one1 小时前
第9篇:Freqtrade量化交易之config.json 基础入门与初始化
javascript·数据库·python·mysql·算法·json·区块链
湫兮之风2 小时前
Opencv: cv::LUT()深入解析图像块快速查表变换
人工智能·opencv·计算机视觉
大翻哥哥3 小时前
Python 2025:量化金融与智能交易的新纪元
开发语言·python·金融
Christo33 小时前
TFS-2018《On the convergence of the sparse possibilistic c-means algorithm》
人工智能·算法·机器学习·数据挖掘
qq_508823403 小时前
金融量化指标--2Alpha 阿尔法
大数据·人工智能
黑金IT3 小时前
`.cursorrules` 与 `.cursorcontext`:Cursor AI 编程助手时代下的“双轨配置”指南
人工智能
zhousenshan4 小时前
Python爬虫常用框架
开发语言·爬虫·python
dlraba8024 小时前
基于 OpenCV 的信用卡数字识别:从原理到实现
人工智能·opencv·计算机视觉
IMER SIMPLE4 小时前
人工智能-python-深度学习-经典神经网络AlexNet
人工智能·python·深度学习