Pytorch的Dataloader加载图片并显示[自学笔记]

AI Chen

已于 2022-05-09 15:37:25 修改

阅读量4.7k

点赞数 7

分类专栏： Pytorch学习文章标签： pytorch 深度学习 python Dataloader

于 2022-05-09 15:31:20 首次发布

本文链接：https://blog.csdn.net/qq_39400324/article/details/124666479

版权

Pytorch学习专栏收录该内容

15 篇文章

订阅专栏

这篇博客介绍如何使用PyTorch构建一个数据类，加载和显示语义分割数据集。数据集包含两个子文件夹：SegmentationClass存储mask.png，JPEGImages存储.jpg图片。通过定义转化函数保持图片尺寸，并使用Dataset和DataLoader加载数据。示例代码展示了如何读取和显示数据加载器中的一组图片和其对应的mask。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

(自学笔记，都是抄的，长期更新)

利用pytorch写一个数据类，可以加载数据图片并显示。在datasets文件夹下有两个子文件夹，分别是SegmentationClass用以存放语义分割数据的mask.png，JPEGImages用以存放语义分割的数据图片.jpg。

在这里插入图片描述

import os
import torch
from torchvision import transforms
from torch.utils.data import Dataset,DataLoader
from  PIL import  Image

# 将图片转换为tensor
transform=transforms.Compose([
    transforms.ToTensor()
])
# 将tensor转化为图片
show_tensor_pic=transforms.ToPILImage()

#保持图片最长的边进行缩放，默认缩放到(256,256)
def keep_png_size(path,size=(256,256)):
    img=Image.open(path)
    temp=max(img.size)
    mask=Image.new('P',(temp,temp))
    mask.paste(img,(0,0))
    mask=mask.resize(size)
    return mask
#保持图片最长的边进行缩放，默认缩放到(26,256)
def keep_jpg_size(path,size=(256,256)):
    img=Image.open(path)
    temp=max(img.size)
    mask=Image.new('RGB',(temp,temp))
    mask.paste(img,(0,0))
    mask=mask.resize(size)
    return mask


class MyDataset(Dataset):
    def __init__(self,path):
        self.path=path#输入图片路径
        #保存所有mask的名字
        self.name=os.listdir(os.path.join(path,'SegmentationClass'))

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

    def __getitem__(self, item):
        # 根据索引获取一张mask的名字
        segment_name=self.name[item]

        # 获取这张mask的路径
        segment_path=os.path.join(self.path,'SegmentationClass',segment_name)
        # 根据mask的名字取索引图片的名字
        image_path=os.path.join(self.path,'JPEGImages',segment_name.replace('png','jpg'))

        # 将mask压缩到一定的尺寸
        segment_image=keep_png_size(segment_path)
        # 将图片压缩到一定的尺寸
        image=keep_jpg_size(image_path)

        return transform(image),transform(segment_image)

if __name__ == '__main__':
	#建立一个dataloader，2个图片为一组，随机打乱
    data_loader=DataLoader(MyDataset('datasets'),batch_size=2,shuffle=True)
	#读取一组图片和标签
    img, seg=next(iter(data_loader))
	#显示该组的第一张图片和标签
    show_tensor_pic(img[0]).show()
    show_tensor_pic(seg[0]).show()