Pytorch: как я могу использовать изображение в качестве метки в загрузчике данных?

Question

Я хочу выполнить некоторую реконструкцию изображения, используя авто-кодеры в pytorch, однако я не нашел способа использовать изображение в качестве метки для входного изображения (изображение метки отличается от оригинала)

Я попробовал метод папки с изображениями, но я думаю, что это для классификации, и в настоящее время я не могу найти одно решение.Должен ли я создать собственный набор данных для этого ...

Заранее спасибо!

Answer 1

Напишите свой пользовательский Dataset, ниже приведен простой пример.

import torch.utils.data.Dataset as Dataset

class CustomDataset(Dataset):

    def __init__(self, input_imgs, label_imgs, transform):

        self.input_imgs = input_imgs
        self.label_imgs = label_imgs
        self.transform = transform

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


    def __getitem__(self, idx):
        input_img, label_img = self.input_imgs[idx], self.label_imgs[idx]
        return self.transform(input_img), self.transform(label_img)

А затем передайте его в Dataloader:

dataloader = DataLoader(CustomDataset)