| import os |
| import numpy as np |
| from PIL import Image |
| import torch.utils.data as data |
|
|
| DATA_ROOTS = 'data/CUBirds' |
|
|
| |
| |
|
|
| class CUBirds(data.Dataset): |
| def __init__(self, root=DATA_ROOTS, train=True, image_transforms=None): |
| super().__init__() |
| self.root = root |
| self.train = train |
| self.image_transforms = image_transforms |
| paths, labels = self.load_images() |
| self.paths, self.labels = paths, labels |
|
|
| def load_images(self): |
| image_info_path = os.path.join(self.root, 'images.txt') |
| with open(image_info_path, 'r') as f: |
| image_info = [line.split('\n')[0].split(' ', 1) for line in f.readlines()] |
| image_info = dict(image_info) |
|
|
| |
| label_info_path = os.path.join(self.root, 'image_class_labels.txt') |
| with open(label_info_path, 'r') as f: |
| label_info = [line.split('\n')[0].split(' ', 1) for line in f.readlines()] |
| label_info = dict(label_info) |
|
|
| |
| train_test_info_path = os.path.join(self.root, 'train_test_split.txt') |
| with open(train_test_info_path, 'r') as f: |
| train_test_info = [line.split('\n')[0].split(' ', 1) for line in f.readlines()] |
| train_test_info = dict(train_test_info) |
|
|
| all_paths, all_labels = [], [] |
| for index, image_path in image_info.items(): |
| label = label_info[index] |
| split = int(train_test_info[index]) |
| if self.train: |
| if split == 1: |
| all_paths.append(image_path) |
| all_labels.append(label) |
| else: |
| if split == 0: |
| all_paths.append(image_path) |
| all_labels.append(label) |
| return all_paths, all_labels |
|
|
| def __len__(self): |
| return len(self.paths) |
|
|
| def __getitem__(self, index): |
| path = os.path.join(self.root, 'images', self.paths[index]) |
| label = int(self.labels[index]) - 1 |
| image = Image.open(path).convert(mode='RGB') |
| if self.image_transforms: |
| image = self.image_transforms(image) |
| return image, label |
