| | 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.dataset_size = len(open('demo.txt').readlines()) |
| |
|
| | dir_I = '_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) |
| |
|
| | def __getitem__(self, index): |
| |
|
| | file_path ='demo.txt' |
| | im_name, c_name = linecache.getline(file_path, index+1).strip().split() |
| |
|
| | I_path = os.path.join(self.dir_I,im_name) |
| | I = Image.open(I_path).convert('RGB') |
| | |
| | I = I.resize((self.fine_width, self.fine_height)) |
| |
|
| | 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.dir_C,c_name) |
| | C = Image.open(C_path).convert('RGB') |
| | C_tensor = transform(C) |
| |
|
| | E_path = os.path.join(self.dir_E,c_name) |
| | E = Image.open(E_path).convert('L') |
| | E_tensor = transform_E(E) |
| |
|
| | input_dict = { 'image': I_tensor,'clothes': C_tensor, 'edge': E_tensor} |
| | return input_dict |
| |
|
| | def __len__(self): |
| | return self.dataset_size |
| |
|
| | def name(self): |
| | return 'AlignedDataset' |
| |
|
