使用PyTorch加载图像的自定义数据集

时间:2019-08-25 19:07:02

标签: image-processing dataset pytorch dataloader

我正在使用coil-100数据集,该数据集包含100个对象的图像,通过将每个对象旋转5度,从固定摄像机拍摄的每个对象72个图像。以下是我正在使用的文件夹结构:

  

数据/火车/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

我使用了imageloader和dataloader类。训练和测试数据集已正确加载,我可以打印班级名称。

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.pngpost_resnet_train_data['obj01'][0]应该是obj01_0.png的resnet输出,依此类推。
我是Pytorch的初学者,在过去的两天里,我阅读了许多有关创建自定义数据集类的教程和stackoverflow问题,但无法弄清楚如何实现我想要的。 请帮忙!

1 个答案:

答案 0 :(得分:4)

假设您只计划在图像上重新发送一次,并保存输出以备后用,建议您编写自己的data set,该源于ImageFolder
将每个resnet输出保存到扩展名为.pth的图像文件所在的位置。

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