修复错误 img 应该是 PIL Image。得到 <class 'torch.Tensor'>

Question

我尝试创建自定义数据集，但在显示某些图像时出现错误。这是我的数据集类和转换：

transforms = transforms.Compose([transforms.Resize(224,224)])

class MyDataset(Dataset):
    def __init__(self, path, label,  transform=None):
        self.path = glob.glob(os.path.join(path, '*.jpg'))
        self.transform = transform
        self.label = label

    def __getitem__(self, index):
        img = io.imread(self.path[index])
        img = torch.tensor(img)
        labels = torch.tensor(int(self.label))
        if self.transform:
          img = self.transform(img)
        return (img,labels)
    
    def __len__(self):
        return len(self.path)

这里是错误行：

images, labels = next(iter(train_loader))

Answer 1

transforms.Resize 需要 PIL.Image 实例为 input 而您的 img 是 torch.Tensor。

这将解决您的问题（请参阅源代码中的 cmets）：

import torchvision
from PIL import Image

# In your transform you should cast PIL Image to tensor
# When no transforms on PIL Image are needed anymore
transforms = transforms.Compose([transforms.Resize(224, 224), transforms.ToTensor()])


class MyDataset(Dataset):
    def __init__(self, path, label, transform=None):
        self.path = glob.glob(os.path.join(path, "*.jpg"))
        self.transform = transform
        self.label = label

    def __getitem__(self, index):
        img = Image.open(self.path[index])
        labels = torch.tensor(int(self.label))
        if self.transform is not None:
            # Now you have PIL.Image instance for transforms
            img = self.transform(img)
        return (img, labels)

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