从 Pandas 加载图像数据到 pytorch

Question

我是 pytorch 的新手，之前曾在 keras 和 fastai 上工作过。目前正在尝试图像回归任务，挑战是我必须从 Pandas 数据框中加载数据。数据帧结构：

ID   Path   Score
fig1  /folder/fig1.jpg  2
fig2  /folder/fig2.jpg  3
.....

我之前一直致力于将图像直接从文件夹加载到 pytorch，因为这是一项简单的分类任务，但现在有点卡住了。

我查看了 pytorch 论坛，但并没有安静地了解如何实施。任何帮助，将不胜感激。

Answer 1

数据集

您必须使用torch.utils.data.Dataset结构来定义它。以下是简单的方法pytorch（我正在使用pillow加载图像并将torchvision它们转换为torch.Tensor对象）：

import torch
import torchvision
from PIL import Image


class MyDataset(torch.utils.data.Dataset):
    def __init__(self, dataframe):
        self.dataframe = dataframe

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

    def __getitem__(self, index):
        row = self.dataframe.iloc[index]
        return (
            torchvision.transforms.functional.to_tensor(Image.open(row["Path"])),
            row["Score"],
        )

dataset = MyDataset(dataframe)

或者，您可以使用torchdata（免责声明：无耻的自我推销，因为我是作者...）它允许您解耦Path并Scores像这样：

import torchvision
from PIL import Image

import torchdata


class ImageDataset(torchdata.datasets.FilesDataset):
    def __getitem__(self, index):
        return Image.open(self.files[index])


class Labels(torchdata.Dataset):
    def __init__(self, scores):
        super().__init__()
        self.scores = scores

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

    def __getitem__(self, index):
        return self.scores[index]

# to_numpy for convenience
# I assume all your images are in /folder and have *.jpg extension
dataset = ImageDataset.from_folder("/folder", regex="*.jpg").map(
    torchvision.transforms.ToTensor()
) | Labels(dataframe["Score"].to_numpy())

（或者您可以像在常规中一样实现它，pytorch但继承torchdata.Dataset并调用super().__init__()构造函数）。

torchdata允许您轻松缓存图像或通过.map如图所示应用一些其他转换，查看 github 存储库以获取更多信息或在评论中询问。

数据加载器

无论您选择哪种方式，您都应该包装数据集torch.utils.data.DataLoader以创建批次并对其进行迭代，如下所示：

dataloader = torch.utils.data.DataLoader(dataset, batch_size=64, shuffle=True)

for images, scores in dataloader:
    # Rest of your code to train neural network or smth
    ...

使用这些图像并在循环中为您想要的内容打分。