使用 PyTorch 加载自定义图像数据集 [英] Loading custom dataset of images using PyTorch

问题描述

我正在使用线圈 100 数据集，该数据集包含 100 个对象的图像，每个对象的 72 个图像是从固定相机拍摄的，每个图像将对象旋转 5 度.以下是我正在使用的文件夹结构:

I'm using the coil-100 dataset which has images of 100 objects, 72 images per object taken from a fixed camera by turning the object 5 degrees per image. Following is the folder structure I'm using:

数据/火车/obj1/obj01_0.png, obj01_5.png ... obj01_355.png
.
.
数据/火车/obj85/obj85_0.png, obj85_5.png ... obj85_355.png
.
.
数据/测试/obj86/obj86_0.ong, obj86_5.png ... obj86_355.png
.
.
数据/测试/obj100/obj100_0.ong, obj100_5.png ... obj100_355.png

data/train/obj1/obj01_0.png, obj01_5.png ... obj01_355.png
.
.
data/train/obj85/obj85_0.png, obj85_5.png ... obj85_355.png
.
.
data/test/obj86/obj86_0.ong, obj86_5.png ... obj86_355.png
.
.
data/test/obj100/obj100_0.ong, obj100_5.png ... obj100_355.png

我使用了 imageloader 和 dataloader 类.训练和测试数据集正确加载，我可以打印类名.

I have used the imageloader and dataloader classes. The train and test datasets loaded properly and I can print the class names.

train_path = 'data/train/'
test_path = 'data/test/'
data_transforms = {
    transforms.Compose([
    transforms.Resize(224, 224),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ])
}

train_data = torchvision.datasets.ImageFolder(
    root=train_path,
    transform= data_transforms
)
test_data = torchvision.datasets.ImageFolder(
    root = test_path,
    transform = data_transforms
)
train_loader = torch.utils.data.DataLoader(
    train_data,
    batch_size=None,
    num_workers=1,
    shuffle=False
)
test_loader = torch.utils.data.DataLoader(
    test_data,
    batch_size=None,
    num_workers=1,
    shuffle=False
)

print(len(train_data))
print(len(test_data))
classes = train_data.class_to_idx
print("detected classes: ", classes)

在我的模型中，我希望通过预训练的 resnet 传递每张图像，并从 resnet 的输出中制作一个数据集以馈送到双向 LSTM.为此，我需要通过类名和索引访问图像.例如.pre_resnet_train_data['obj01'][0] 应该是 obj01_0.png 并且 post_resnet_train_data['obj01'][0] 应该是 resnet 输出obj01_0.png 等等.
我是 Pytorch 的初学者，在过去的两天里，我阅读了许多关于创建自定义数据集类的教程和 stackoverflow 问题，但不知道如何实现我想要的.请帮忙！

In my model I wish to pass every image through pretrained resnet and make a dataset from the output of resnet to feed into a biderectional LSTM. For which I need to access the images by classname and index. for ex. pre_resnet_train_data['obj01'][0] should be obj01_0.png and post_resnet_train_data['obj01'][0] should be the resnet output of obj01_0.png and so on.
I'm a beginner in Pytorch and for the past 2 days, I have read many tutorials and stackoverflow questions about creating a custom dataset class but couldn't figure out how to achieve what I want. please help!

推荐答案

假设你只打算在图像上运行一次重新发送并保存输出以供以后使用，我建议你自己编写 数据集，源自 ImageFolder.
将每个 resnet 输出保存在与图像文件相同的位置，扩展名为 .pth.

Assuming you only plan on running resent on the images once and save the output for later use, I suggest you write your own data set, derived from ImageFolder.
Save each resnet output at the same location as the image file with .pth extension.

class MyDataset(torchvision.datasets.ImageFolder):
  def __init__(self, root, transform):
    super(MyDataset, self).__init__(root, transform)

  def __getitem__(self, index):
    # override ImageFolder's method
    """
    Args:
      index (int): Index
    Returns:
      tuple: (sample, resnet, target) where target is class_index of the target class.
    """
    path, target = self.samples[index]
    sample = self.loader(path)
    if self.transform is not None:
      sample = self.transform(sample)
    if self.target_transform is not None:
      target = self.target_transform(target)
    # this is where you load your resnet data
    resnet_path = os.path.join(os.path.splitext(path)[0], '.pth')  # replace image extension with .pth
    resnet = torch.load(resnet_path)  # load the stored features
    return sample, resnet, target

这篇关于使用 PyTorch 加载自定义图像数据集的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！