【Pytorch使用自制数据集,Dataloader】
2023-12-13 03:45:25
数据集结构
话不多说,直接上核心代码
myDataset.py
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
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
大家结构和我不一样可以自由发挥
文章来源:https://blog.csdn.net/weixin_47466118/article/details/134814218
本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。 如若内容造成侵权/违法违规/事实不符,请联系我的编程经验分享网邮箱:veading@qq.com进行投诉反馈,一经查实,立即删除!
本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。 如若内容造成侵权/违法违规/事实不符,请联系我的编程经验分享网邮箱:veading@qq.com进行投诉反馈,一经查实,立即删除!