如何从 Pytorch 图像列表开始加载数据集

Question

我有一个服务从另一个服务接收二进制格式的图像（我们称之为服务 B）：

from PIL import Image

img_list = []
img_bin = get_image_from_service_B()
image = Image.open(io.BytesIO(img_bin)) # Convert bytes to image using PIL

---

当图像通过 PIL 成功转换后，它也会被附加到图像列表中。

img_list.append(image)    

---

当我有足够的图像时，我想使用 Pytorch 加载图像列表，就好像它是数据集一样

if img_list.__len__() == 500:
     ### Load dataset and do a transform operation on the data

---

在该软件的早期版本中，要求只是从文件夹中检索图像，因此加载所有图像非常简单

my_dataset = datasets.ImageFolder("path/to/images/folder/", transform=transform)
dataset_iterator = DataLoader(my_dataset, batch_size=1)

现在我的问题是如何执行转换并从列表加载数据集。

Answer 1

您可以简单地编写自定义数据集：

class MyDataset(torch.utils.data.Dataset):
    def __init__(self, img_list, augmentations):
        super(MyDataset, self).__init__()
        self.img_list = img_list
        self.augmentations = augmentations

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

    def __getitem__(self, idx):
        img = self.img_list[idx]
        return self.augmentations(img)
  

现在，您可以将此自定义数据集插入其中DataLoader，然后就完成了。