| import os.path |
| from data.base_dataset import BaseDataset, get_params, get_transform |
| from PIL import Image |
| import linecache |
|
|
| class AlignedDataset(BaseDataset): |
| def initialize(self, opt): |
| self.opt = opt |
| self.root = opt.dataroot |
|
|
| self.fine_height=256 |
| self.fine_width=192 |
|
|
| self.text = './test_pairs.txt' |
|
|
| dir_I = '_ma_img' |
| self.dir_I = os.path.join(opt.dataroot, opt.phase + dir_I) |
|
|
| dir_C = '_clothes' |
| self.dir_C = os.path.join(opt.dataroot, opt.phase + dir_C) |
|
|
| dir_E = '_edge' |
| self.dir_E = os.path.join(opt.dataroot, opt.phase + dir_E) |
|
|
| self.im_name = [] |
| self.c_name = [] |
| self.e_name = [] |
| self.get_file_name() |
| |
| self.dataset_size = len(self.im_name) |
|
|
| def get_file_name(self): |
|
|
| with open(self.text, 'r') as f: |
| for line in f.readlines(): |
| im_name, c_name = line.strip().split() |
| self.im_name.append(os.path.join(self.dir_I, im_name)) |
| self.c_name.append(os.path.join(self.dir_C, c_name)) |
| self.e_name.append(os.path.join(self.dir_E, c_name)) |
|
|
| def __getitem__(self, index): |
|
|
| |
| |
|
|
| I_path = os.path.join(self.im_name[index]) |
| I = Image.open(I_path).convert('RGB') |
|
|
| params = get_params(self.opt, I.size) |
| transform = get_transform(self.opt, params) |
| transform_E = get_transform(self.opt, params, method=Image.NEAREST, normalize=False) |
|
|
| I_tensor = transform(I) |
| |
| C_path = os.path.join(self.c_name[index]) |
| |
| C = Image.open(C_path).convert('RGB') |
| C_tensor = transform(C) |
|
|
| E_path = os.path.join(self.e_name[index]) |
| E = Image.open(E_path).convert('L') |
| E_tensor = transform_E(E) |
|
|
| input_dict = { 'image': I_tensor,'clothes': C_tensor, 'edge': E_tensor, 'p_name':self.im_name[index].split('/')[-1]} |
| return input_dict |
|
|
| def __len__(self): |
| return self.dataset_size |
|
|
| def name(self): |
| return 'AlignedDataset' |
|
|
