File size: 1,672 Bytes
8d9f691 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 |
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')
# new line
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'
|