You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
227 lines
8.9 KiB
227 lines
8.9 KiB
### Copyright (C) 2020 Roy Or-El. All rights reserved.
|
|
### Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
|
|
import os.path
|
|
import re
|
|
import torch
|
|
import random
|
|
import numpy as np
|
|
from data.base_dataset import BaseDataset
|
|
from data.dataset_utils import list_folder_images, get_transform
|
|
from util.preprocess_itw_im import preprocessInTheWildImage
|
|
from PIL import Image
|
|
from pdb import set_trace as st
|
|
|
|
CLASSES_UPPER_BOUNDS = [2, 6, 9, 14, 19, 29, 39, 49, 69, 120]
|
|
|
|
class MulticlassUnalignedDataset(BaseDataset):
|
|
def initialize(self, opt):
|
|
self.opt = opt
|
|
self.root = opt.dataroot
|
|
self.name_mapping = {}
|
|
self.prev_A = -1
|
|
self.prev_B = -1
|
|
self.class_A = -1
|
|
self.class_B = -1
|
|
self.get_samples = False
|
|
if not self.opt.isTrain:
|
|
self.in_the_wild = opt.in_the_wild
|
|
else:
|
|
self.in_the_wild = False
|
|
|
|
# find all existing classes in root
|
|
if not self.in_the_wild:
|
|
self.tempClassNames = []
|
|
subDirs = next(os.walk(self.root))[1] # a quick way to get all subdirectories
|
|
for currDir in subDirs:
|
|
if self.opt.isTrain:
|
|
prefix = 'train'
|
|
else:
|
|
prefix = 'test'
|
|
if prefix in currDir: # we assume that the class name starts with the prefix
|
|
len_prefix = len(prefix)
|
|
className = currDir[len_prefix:]
|
|
self.tempClassNames += [className]
|
|
|
|
# sort classes
|
|
if len(self.opt.sort_order) > 0:
|
|
self.classNames = []
|
|
for i, nextClass in enumerate(self.opt.sort_order):
|
|
for currClass in self.tempClassNames:
|
|
if nextClass == currClass:
|
|
self.classNames += [currClass]
|
|
curr_class_num = self.assign_age_class(currClass)
|
|
self.name_mapping[currClass] = curr_class_num
|
|
else:
|
|
self.classNames = sorted(self.tempClassNames)
|
|
for i, currClass in enumerate(self.classNames):
|
|
curr_class_num = self.assign_age_class(currClass)
|
|
self.name_mapping[currClass] = curr_class_num
|
|
else:
|
|
self.classNames = []
|
|
for i, nextClass in enumerate(self.opt.sort_order):
|
|
self.classNames += [nextClass]
|
|
curr_class_num = self.assign_age_class(nextClass)
|
|
self.name_mapping[nextClass] = curr_class_num
|
|
|
|
self.active_classes_mapping = {}
|
|
|
|
for i, name in enumerate(self.classNames):
|
|
self.active_classes_mapping[i] = self.name_mapping[name]
|
|
|
|
self.numClasses = len(self.classNames)
|
|
opt.numClasses = self.numClasses
|
|
opt.classNames = self.classNames
|
|
|
|
# set class counter for test mode
|
|
if self.opt.isTrain is False:
|
|
opt.batchSize = self.numClasses
|
|
self.class_counter = 0
|
|
self.img_counter = 0
|
|
|
|
# arrange directories
|
|
if not self.in_the_wild:
|
|
self.dirs = []
|
|
self.img_paths = []
|
|
self.parsing_paths = []
|
|
self.sizes = []
|
|
|
|
for currClass in self.classNames:
|
|
self.dirs += [os.path.join(self.root, opt.phase + currClass)]
|
|
imgs, parsings = list_folder_images(self.dirs[-1], self.opt)
|
|
self.img_paths += [imgs]
|
|
self.parsing_paths += [parsings]
|
|
self.sizes += [len(self.img_paths[-1])]
|
|
|
|
opt.dataset_size = self.__len__()
|
|
|
|
self.transform = get_transform(opt)
|
|
|
|
if (not self.opt.isTrain) and self.in_the_wild:
|
|
self.preprocessor = preprocessInTheWildImage(out_size=opt.fineSize)
|
|
|
|
def set_sample_mode(self, mode=False):
|
|
self.get_samples = mode
|
|
self.class_counter = 0
|
|
self.img_counter = 0
|
|
|
|
def assign_age_class(self, class_name):
|
|
ages = [int(s) for s in re.split('-|_', class_name) if s.isdigit()]
|
|
max_age = ages[-1]
|
|
for i in range(len(CLASSES_UPPER_BOUNDS)):
|
|
if max_age <= CLASSES_UPPER_BOUNDS[i]:
|
|
break
|
|
|
|
return i
|
|
|
|
def mask_image(self, img, parsings):
|
|
labels_to_mask = [0,14,15,16,18]
|
|
for idx in labels_to_mask:
|
|
img[parsings == idx] = 128
|
|
|
|
return img
|
|
|
|
def get_item_from_path(self, path):
|
|
path_dir, im_name = os.path.split(path)
|
|
img = Image.open(path).convert('RGB')
|
|
img = np.array(img.getdata(), dtype=np.uint8).reshape(img.size[1], img.size[0], 3)
|
|
|
|
if self.in_the_wild:
|
|
img, parsing = self.preprocessor.forward(img)
|
|
else:
|
|
parsing_path = os.path.join(path_dir, 'parsings', im_name[:-4] + '.png')
|
|
parsing = Image.open(parsing_path).convert('RGB')
|
|
parsing = np.array(parsing.getdata(), dtype=np.uint8).reshape(parsing.size[1], parsing.size[0], 3)
|
|
|
|
img = Image.fromarray(self.mask_image(img, parsing))
|
|
img = self.transform(img).unsqueeze(0)
|
|
|
|
return {'Imgs': img,
|
|
'Paths': [path],
|
|
'Classes': torch.zeros(1, dtype=torch.int),
|
|
'Valid': True}
|
|
|
|
def __getitem__(self, index):
|
|
if self.opt.isTrain and not self.get_samples:
|
|
condition = True
|
|
self.class_A_idx = random.randint(0,self.numClasses - 1)
|
|
self.class_A = self.active_classes_mapping[self.class_A_idx]
|
|
while condition:
|
|
self.class_B_idx = random.randint(0,self.numClasses - 1)
|
|
self.class_B = self.active_classes_mapping[self.class_B_idx]
|
|
condition = self.class_A == self.class_B
|
|
|
|
index_A = random.randint(0, self.sizes[self.class_A_idx] - 1)
|
|
index_B = random.randint(0, self.sizes[self.class_B_idx] - 1)
|
|
|
|
A_img_path = self.img_paths[self.class_A_idx][index_A]
|
|
A_img = Image.open(A_img_path).convert('RGB')
|
|
A_img = np.array(A_img.getdata(), dtype=np.uint8).reshape(A_img.size[1], A_img.size[0], 3)
|
|
|
|
B_img_path = self.img_paths[self.class_B_idx][index_B]
|
|
B_img = Image.open(B_img_path).convert('RGB')
|
|
B_img = np.array(B_img.getdata(), dtype=np.uint8).reshape(B_img.size[1], B_img.size[0], 3)
|
|
|
|
A_parsing_path = self.parsing_paths[self.class_A_idx][index_A]
|
|
A_parsing = Image.open(A_parsing_path).convert('RGB')
|
|
A_parsing = np.array(A_parsing.getdata(), dtype=np.uint8).reshape(A_parsing.size[1], A_parsing.size[0], 3)
|
|
A_img = Image.fromarray(self.mask_image(A_img, A_parsing))
|
|
|
|
B_parsing_path = self.parsing_paths[self.class_B_idx][index_B]
|
|
B_parsing = Image.open(B_parsing_path).convert('RGB')
|
|
B_parsing = np.array(B_parsing.getdata(), dtype=np.uint8).reshape(B_parsing.size[1], B_parsing.size[0], 3)
|
|
B_img = Image.fromarray(self.mask_image(B_img, B_parsing))
|
|
|
|
# numpy conversions are an annoying hack to form a PIL image with more than 3 channels
|
|
A_img = self.transform(A_img)
|
|
B_img = self.transform(B_img)
|
|
|
|
return {'A': A_img, 'B': B_img,
|
|
"A_class": self.class_A_idx, "B_class": self.class_B_idx,
|
|
'A_paths': A_img_path, 'B_paths': B_img_path}
|
|
|
|
else: # in test mode, load one image from each class
|
|
i = self.class_counter % self.numClasses
|
|
self.class_counter += 1
|
|
|
|
if self.get_samples:
|
|
ind = random.randint(0, self.sizes[i] - 1)
|
|
else:
|
|
ind = self.img_counter if self.img_counter < self.sizes[i] else -1
|
|
|
|
if i == self.numClasses - 1:
|
|
self.img_counter += 1
|
|
|
|
if ind > -1:
|
|
valid = True
|
|
paths = self.img_paths[i][ind]
|
|
img = Image.open(self.img_paths[i][ind]).convert('RGB')
|
|
img = np.array(img.getdata(), dtype=np.uint8).reshape(img.size[1], img.size[0], 3)
|
|
|
|
parsing_path = self.parsing_paths[i][ind]
|
|
parsing = Image.open(parsing_path).convert('RGB')
|
|
parsing = np.array(parsing.getdata(), dtype=np.uint8).reshape(parsing.size[1], parsing.size[0], 3)
|
|
img = Image.fromarray(self.mask_image(img, parsing))
|
|
|
|
img = self.transform(img)
|
|
|
|
else:
|
|
img = torch.zeros(3, self.opt.fineSize, self.opt.fineSize)
|
|
paths = ''
|
|
valid = False
|
|
|
|
return {'Imgs': img,
|
|
'Paths': paths,
|
|
'Classes': i,
|
|
'Valid': valid}
|
|
|
|
def __len__(self):
|
|
if self.opt.isTrain:
|
|
return round(sum(self.sizes) / 2) # this determines how many iterations we make per epoch
|
|
elif self.in_the_wild:
|
|
return 0
|
|
else:
|
|
return max(self.sizes) * self.numClasses
|
|
|
|
def name(self):
|
|
return 'MulticlassUnalignedDataset'
|