pythorch 实现VOC数据集的加载

最新推荐文章于 2024-04-20 09:47:21 发布

糖不哭

最新推荐文章于 2024-04-20 09:47:21 发布

阅读量677

点赞数

文章标签：深度学习

本文链接：https://blog.csdn.net/tang330023555/article/details/119038550

版权

该博客主要介绍了如何使用PyTorch加载和预处理VOC2012数据集。通过定义`Data_Load`类，实现了对图像和标签的读取、随机裁剪等预处理操作，并使用`DataLoader`进行数据加载。此外，还使用了TensorBoard进行可视化。

摘要由CSDN通过智能技术生成

import os
import shutil

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


class Data_Load(Dataset):

    def __init__(self, root_dir, label_dir):
        self.root_dir = root_dir
        self.label_dir = label_dir
        self.img_path = []
        self.label_path = []
        file_list_img = open("E:\\pytorchlearning\\data\\VOCdevkit\\VOC2012\\ImageSets\\Segmentation\\trainval.txt",
                             "r")
        for line in file_list_img.readlines():
            self.img_path.append(list(map(str, line.split('\n'))))
        file_list_label = open("E:\\pytorchlearning\\data\\VOCdevkit\\VOC2012\\ImageSets\\Segmentation\\trainval.txt",
                               "r")
        for line in file_list_label.readlines():
            self.label_path.append(list(map(str, line.split('\n'))))
        self.transforms = transforms.Compose([
            torchvision.transforms.Resize((512, 512), interpolation=2),
            transforms.ToTensor(),  # 将图像转为Tensor
        ])

    def __getitem__(self, idx):
        # img_name = int(str(self.img_path[idx]))
        img_name = self.img_path[idx][0] + ".jpg"
        # print(img_name)
        img_item_path = os.path.join(self.root_dir, img_name)
        img = Image.open(img_item_path)

        label_name = self.label_path[idx][0] + ".png"
        label_item_path = os.path.join(self.label_dir, label_name)
        label = Image.open(label_item_path)
        # print(img)
        # imgs = []
        if img.size[0] >= 300 and img.size[1] >= 200:
            img, label = self.voc_rand_crop(img, label, 200, 300)

        return self.transforms(img), self.transforms(label)

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

    def voc_rand_crop(self, feature, label, height, width):
        """
        随机裁剪feature(PIL image) 和 label(PIL image).
        为了使裁剪的区域相同，不能直接使用RandomCrop，而要像下面这样做
        Get parameters for ``crop`` for a random crop.
        Args:
            img (PIL Image): Image to be cropped.
            output_size (tuple): Expected output size of the crop.
        Returns:
            tuple: params (i, j, h, w) to be passed to ``crop`` for random crop.
        """
        i, j, h, w = torchvision.transforms.RandomCrop.get_params(feature, output_size=(height, width))
        feature = torchvision.transforms.functional.crop(feature, i, j, h, w)
        label = torchvision.transforms.functional.crop(label, i, j, h, w)
        # print(feature)
        return feature, label


root_dir = "E:\\pytorchlearning\\data\VOCdevkit\\VOC2012\\JPEGImages"
label_dir = "E:\\pytorchlearning\\data\\VOCdevkit\\VOC2012\\SegmentationClass"
dataset = Data_Load(root_dir, label_dir)
data_loader = DataLoader(dataset, batch_size=4)

if os.path.exists("./log/log_dataloader"):
    shutil.rmtree("./log/log_dataloader")
    print("删除日志")

writer = SummaryWriter("./log/log_dataloader")
total_step = 0
for data in data_loader:
    imgs, label = data
    if total_step == 10:
        print(imgs.shape)
        print(label.shape)
        writer.add_images("train", imgs)
        writer.add_images("test", label)
    total_step += 1
writer.close()