Разделить набор данных на основе имен файлов в наборе данных Pytorch - PullRequest
Новая
       23

Разделить набор данных на основе имен файлов в наборе данных Pytorch

0 голосов
/ 24 сентября 2018

Есть ли способ разделить набор данных на обучение и тестирование на основе имен файлов.У меня есть папка, содержащая две папки: вход и выход.Входная папка содержит изображения, а выходные - метки для этого изображения.Имена файлов во входной папке выглядят примерно так: input01_train.png и input01_test.png, как показано ниже. 

                          Dataset
                          /     \
                     Input       Output
                      |             |
           input01_train.png   output01_train.png
                    .                 .
                    .                 .
           input01_test.png    output01_test.png

У меня есть код, который разделяет только набор данных на входы и метки, а не тест и тренировка. 

class CancerDataset(Dataset):
  def __init__(self, dataset_folder):#,label_folder):
    self.dataset_folder = torchvision.datasets.ImageFolder(dataset_folder ,transform = transforms.Compose([transforms.Resize(512),transforms.ToTensor()]))
    self.label_folder = torchvision.datasets.ImageFolder(dataset_folder ,transform = transforms.Compose([transforms.Resize(512),transforms.ToTensor()]))

  def __getitem__(self,index):
    img = self.dataset_folder[index]
    label = self.label_folder[index]
    return img,label

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

trainset = CancerDataset(dataset_folder = '/content/drive/My Drive/cancer_data/')
trainsetloader = DataLoader(trainset,batch_size = 1, shuffle = True,num_workers = 0,pin_memory = True)

Я бы хотел разделить поезд и набор тестов по их именам, если это возможно.

1 Ответ

0 голосов
/ 18 октября 2018

Вы можете загрузить изображения самостоятельно в __getitem__, выбрав только те из них, которые содержат _train.png или _test.png. 

class CancerDataset(Dataset):
    def __init__(self, datafolder, datatype='train', transform = transforms.Compose([transforms.Resize(512),transforms.ToTensor()]):
        self.datafolder = datafolder
        self.image_files_list = [s for s in os.listdir(datafolder) if
                                 '_%s.png' % datatype in s]
        # Same for the labels files
        self.label_files_list = ...
        self.transform = transform

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

    def __getitem__(self, idx):
        img_name = os.path.join(self.datafolder,
                                self.image_files_list[idx])
        image = Image.open(img_name)
        image = self.transform(image)
        # Same for the labels files
        label = .... # Load in etc
        label = self.transform(label)
        return image, label

Теперь вы можете создать два набора данных (trainsetи testset). 

trainset = CancerDataset(dataset_folder = '/content/drive/My Drive/cancer_data/', datatype='train')
testset = CancerDataset(dataset_folder = '/content/drive/My Drive/cancer_data/', datatype='test')
...