读取一个batch的图像并且显示出来

1读取一个batch用于训练

我们在训练模型的时候,除了观察图像的标签和尺寸,最好能读取一个batch的图像显示出来,观察原始图像和grountruth是否对应,如果正确才能正式开始后续的训练。

下面以一个皮肤病分割的数据集加以演示。

2.导入所需要的包

python 复制代码
from torch.utils import data
from torchvision import transforms as T
from torchvision.transforms import functional as F
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
import os

os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'

3.使用dataset 和dataloader加载数据

python 复制代码
class ImageFolder(data.Dataset):
    def __init__(self, root, image_size=224, mode='train', augmentation_prob=0.4):
        """Initializes image paths and preprocessing module."""
        self.root = root
        # GT : Ground Truth
        self.GT_paths = root[:-1] + '_GT/'
        self.image_paths = list(map(lambda x: os.path.join(root, x), os.listdir(root)))
        self.image_size = image_size
        self.mode = mode
        self.RotationDegree = [0, 90, 180, 270]
        self.augmentation_prob = augmentation_prob
        print("image count in {} path :{}".format(self.mode, len(self.image_paths)))

    def __getitem__(self, index):
        """Reads an image from a file and preprocesses it and returns."""
        image_path = self.image_paths[index]
        # [:-len(".jpg")]是列表的索引,image_path.split('_')[-1]是选取路径中的最后一段字符,
        # [: -4]指的是截取第0个到第-4个元素,不包括第4个元素
        filename = image_path.split('_')[-1][:-len(".jpg")]
        GT_path = self.GT_paths + 'ISIC_' + filename + '_segmentation.png'

        image = Image.open(image_path)
        GT = Image.open(GT_path)
        # 计算图像的比例
        dada_transform = T.Compose([
            T.Resize((256,256)),
            T.ToTensor()
        ])

        image= dada_transform(image)
        GT = dada_transform(GT)
        return image,GT

    def __len__(self):
        """Returns the total number of font files."""
        return len(self.image_paths)


def get_loader(image_path, image_size, batch_size, num_workers=0, mode='train', augmentation_prob=0.4):
    """Builds and returns Dataloader."""

    dataset = ImageFolder(root=image_path, image_size=image_size, mode=mode, augmentation_prob=augmentation_prob)
    data_loader = data.DataLoader(dataset=dataset,
                                  batch_size=batch_size,
                                  shuffle=True,
                                  num_workers=num_workers)
    return data_loader

4 开始画图

python 复制代码
if __name__ == '__main__':
    train_path = './dataset/train/'
    image_size = 224
    batch_size =4
    num_workers = 0
    augmentation_prob = 0.4
    train_loader = get_loader(image_path=train_path,
                              image_size=image_size,
                              batch_size=batch_size,
                              num_workers=num_workers,
                              mode='train',
                              augmentation_prob=augmentation_prob)

#从train_loader中获取一个batch的图像和GT
    for step, (img,GT) in enumerate(train_loader):
        if step>0:
            break
    print(f"img:{img.shape}")
    print(f"GT:{GT.shape}")

    for ii in np.arange(4):
        plt.subplot(2, 4, ii + 1)
        image = img[ii, :, :, :].numpy().transpose(1, 2, 0)
        plt.imshow(image)
        plt.axis("off")
        plt.subplot(2, 4, ii + 5)
        GT_i = GT[ii, :, :, :].numpy().transpose(1, 2, 0)
        plt.imshow(GT_i)
        plt.axis("off")
    plt.show()
    plt.subplots_adjust(hspace=0.3)

程序运行的结果如图所示:

相关推荐
一 铭38 分钟前
AI领域新趋势:从提示(Prompt)工程到上下文(Context)工程
人工智能·语言模型·大模型·llm·prompt
顾道长生'2 小时前
(Arxiv-2025)通过动态 token 剔除实现无需训练的高效视频生成
计算机视觉·音视频·视频生成
麻雀无能为力4 小时前
CAU数据挖掘实验 表分析数据插件
人工智能·数据挖掘·中国农业大学
时序之心4 小时前
时空数据挖掘五大革新方向详解篇!
人工智能·数据挖掘·论文·时间序列
.30-06Springfield5 小时前
人工智能概念之七:集成学习思想(Bagging、Boosting、Stacking)
人工智能·算法·机器学习·集成学习
说私域6 小时前
基于开源AI智能名片链动2+1模式S2B2C商城小程序的超级文化符号构建路径研究
人工智能·小程序·开源
永洪科技6 小时前
永洪科技荣获商业智能品牌影响力奖,全力打造”AI+决策”引擎
大数据·人工智能·科技·数据分析·数据可视化·bi
shangyingying_16 小时前
关于小波降噪、小波增强、小波去雾的原理区分
人工智能·深度学习·计算机视觉
书玮嘎7 小时前
【WIP】【VLA&VLM——InternVL系列】
人工智能·深度学习
猫头虎7 小时前
猫头虎 AI工具分享:一个网页抓取、结构化数据提取、网页爬取、浏览器自动化操作工具:Hyperbrowser MCP
运维·人工智能·gpt·开源·自动化·文心一言·ai编程