|
import os |
|
import copy |
|
import numpy as np |
|
from PIL import Image |
|
from os.path import join |
|
from itertools import chain |
|
from collections import defaultdict |
|
|
|
import torch |
|
import torch.utils.data as data |
|
from torchvision import transforms |
|
|
|
DATA_ROOTS = 'data/DTD' |
|
|
|
|
|
|
|
|
|
class DTD(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): |
|
if self.train: |
|
train_info_path = os.path.join(self.root, 'labels', 'train1.txt') |
|
with open(train_info_path, 'r') as f: |
|
train_info = [line.split('\n')[0] for line in f.readlines()] |
|
|
|
val_info_path = os.path.join(self.root, 'labels', 'val1.txt') |
|
with open(val_info_path, 'r') as f: |
|
val_info = [line.split('\n')[0] for line in f.readlines()] |
|
split_info = train_info + val_info |
|
|
|
else: |
|
test_info_path = os.path.join(self.root, 'labels', 'test1.txt') |
|
with open(test_info_path, 'r') as f: |
|
split_info = [line.split('\n')[0] for line in f.readlines()] |
|
|
|
|
|
categories = [] |
|
for row in split_info: |
|
image_path = row |
|
category = image_path.split('/')[0] |
|
categories.append(category) |
|
categories = sorted(list(set(categories))) |
|
|
|
all_paths, all_labels = [], [] |
|
for row in split_info: |
|
image_path = row |
|
category = image_path.split('/')[0] |
|
label = categories.index(category) |
|
all_paths.append(join(self.root, 'images', image_path)) |
|
all_labels.append(label) |
|
return all_paths, all_labels |
|
|
|
def __len__(self): |
|
return len(self.paths) |
|
|
|
def __getitem__(self, index): |
|
path = self.paths[index] |
|
label = self.labels[index] |
|
image = Image.open(path).convert(mode='RGB') |
|
if self.image_transforms: |
|
image = self.image_transforms(image) |
|
return image, label |
