data

data/__init__.py

@@ -0,0 +1 @@
+# data_init

data/aligned_dataset_test.py

@@ -0,0 +1,71 @@
+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()
+        #import ipdb; ipdb.set_trace()
+        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):        
+
+        #file_path ='demo.txt'
+        #im_name, c_name = linecache.getline(file_path, index+1).strip().split()
+
+        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)
+        #import ipdb; ipdb.set_trace()
+        C_path = os.path.join(self.c_name[index])
+        #print(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'

data/base_data_loader.py

@@ -0,0 +1,14 @@
+
+class BaseDataLoader():
+    def __init__(self):
+        pass
+    
+    def initialize(self, opt):
+        self.opt = opt
+        pass
+
+    def load_data():
+        return None
+
+        
+        

data/base_dataset.py

@@ -0,0 +1,113 @@
+import torch.utils.data as data
+from PIL import Image
+import torchvision.transforms as transforms
+import numpy as np
+import random
+
+class BaseDataset(data.Dataset):
+    def __init__(self):
+        super(BaseDataset, self).__init__()
+
+    def name(self):
+        return 'BaseDataset'
+
+    def initialize(self, opt):
+        pass
+
+def get_params(opt, size):
+    w, h = size
+    new_h = h
+    new_w = w
+    if opt.resize_or_crop == 'resize_and_crop':
+        new_h = new_w = opt.loadSize            
+    elif opt.resize_or_crop == 'scale_width_and_crop':
+        new_w = opt.loadSize
+        new_h = opt.loadSize * h // w
+
+    x = random.randint(0, np.maximum(0, new_w - opt.fineSize))
+    y = random.randint(0, np.maximum(0, new_h - opt.fineSize))
+
+    flip = 0
+    return {'crop_pos': (x, y), 'flip': flip}
+
+def get_transform_resize(opt, params, method=Image.BICUBIC, normalize=True):
+    transform_list = []
+    transform_list.append(transforms.Lambda(lambda img: __scale_width(img, opt.loadSize, method)))
+    osize = [256,192]
+    transform_list.append(transforms.Scale(osize, method))
+    if 'crop' in opt.resize_or_crop:
+        transform_list.append(transforms.Lambda(lambda img: __crop(img, params['crop_pos'], opt.fineSize)))
+
+    if opt.resize_or_crop == 'none':
+        base = float(2 ** opt.n_downsample_global)
+        if opt.netG == 'local':
+            base *= (2 ** opt.n_local_enhancers)
+        transform_list.append(transforms.Lambda(lambda img: __make_power_2(img, base, method)))
+
+    if opt.isTrain and not opt.no_flip:
+        transform_list.append(transforms.Lambda(lambda img: __flip(img, params['flip'])))
+
+    transform_list += [transforms.ToTensor()]
+
+    if normalize:
+        transform_list += [transforms.Normalize((0.5, 0.5, 0.5),
+                                                (0.5, 0.5, 0.5))]
+    return transforms.Compose(transform_list)
+
+def get_transform(opt, params, method=Image.BICUBIC, normalize=True):
+    transform_list = []
+    if 'resize' in opt.resize_or_crop:
+        osize = [opt.loadSize, opt.loadSize]
+        transform_list.append(transforms.Scale(osize, method))   
+    elif 'scale_width' in opt.resize_or_crop:
+        transform_list.append(transforms.Lambda(lambda img: __scale_width(img, opt.loadSize, method)))
+        osize = [256,192]
+        transform_list.append(transforms.Scale(osize, method))  
+    if 'crop' in opt.resize_or_crop:
+        transform_list.append(transforms.Lambda(lambda img: __crop(img, params['crop_pos'], opt.fineSize)))
+
+    if opt.resize_or_crop == 'none':
+        base = float(16)
+        transform_list.append(transforms.Lambda(lambda img: __make_power_2(img, base, method)))
+
+    if opt.isTrain and not opt.no_flip:
+        transform_list.append(transforms.Lambda(lambda img: __flip(img, params['flip'])))
+
+    transform_list += [transforms.ToTensor()]
+
+    if normalize:
+        transform_list += [transforms.Normalize((0.5, 0.5, 0.5),
+                                                (0.5, 0.5, 0.5))]
+    return transforms.Compose(transform_list)
+
+def normalize():    
+    return transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
+
+def __make_power_2(img, base, method=Image.BICUBIC):
+    ow, oh = img.size        
+    h = int(round(oh / base) * base)
+    w = int(round(ow / base) * base)
+    if (h == oh) and (w == ow):
+        return img
+    return img.resize((w, h), method)
+
+def __scale_width(img, target_width, method=Image.BICUBIC):
+    ow, oh = img.size
+    if (ow == target_width):
+        return img    
+    w = target_width
+    h = int(target_width * oh / ow)    
+    return img.resize((w, h), method)
+
+def __crop(img, pos, size):
+    ow, oh = img.size
+    x1, y1 = pos
+    tw = th = size
+    if (ow > tw or oh > th):        
+        return img.crop((x1, y1, x1 + tw, y1 + th))
+    return img
+
+def __flip(img, flip):
+    if flip:
+        return img.transpose(Image.FLIP_LEFT_RIGHT)
+    return img

data/custom_dataset_data_loader_test.py

@@ -0,0 +1,31 @@
+import torch.utils.data
+from data.base_data_loader import BaseDataLoader
+
+
+def CreateDataset(opt):
+    dataset = None
+    from data.aligned_dataset_test import AlignedDataset
+    dataset = AlignedDataset()
+
+    print("dataset [%s] was created" % (dataset.name()))
+    dataset.initialize(opt)
+    return dataset
+
+class CustomDatasetDataLoader(BaseDataLoader):
+    def name(self):
+        return 'CustomDatasetDataLoader'
+
+    def initialize(self, opt):
+        BaseDataLoader.initialize(self, opt)
+        self.dataset = CreateDataset(opt)
+        self.dataloader = torch.utils.data.DataLoader(
+            self.dataset,
+            batch_size=opt.batchSize,
+            shuffle = False,
+            num_workers=int(opt.nThreads))
+
+    def load_data(self):
+        return self.dataloader
+
+    def __len__(self):
+        return min(len(self.dataset), self.opt.max_dataset_size)

data/data_loader_test.py

@@ -0,0 +1,7 @@
+
+def CreateDataLoader(opt):
+    from data.custom_dataset_data_loader_test import CustomDatasetDataLoader
+    data_loader = CustomDatasetDataLoader()
+    print(data_loader.name())
+    data_loader.initialize(opt)
+    return data_loader

data/image_folder.py

@@ -0,0 +1,73 @@
+import torch.utils.data as data
+from PIL import Image
+import os
+
+IMG_EXTENSIONS = [
+    '.jpg', '.JPG', '.jpeg', '.JPEG',
+    '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP', '.tiff'
+]
+
+
+def is_image_file(filename):
+    return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
+
+def make_dataset(dir):
+    images = []
+    assert os.path.isdir(dir), '%s is not a valid directory' % dir
+
+    f = dir.split('/')[-1].split('_')[-1]
+    print (dir, f)
+    dirs= os.listdir(dir)
+    for img in dirs:
+
+        path = os.path.join(dir, img)
+        #print(path)
+        images.append(path)
+    return images
+
+def make_dataset_test(dir):
+    images = []
+    assert os.path.isdir(dir), '%s is not a valid directory' % dir
+
+    f = dir.split('/')[-1].split('_')[-1]
+    for i in range(len([name for name in os.listdir(dir) if os.path.isfile(os.path.join(dir, name))])):
+        if f == 'label' or f == 'labelref':
+            img = str(i) + '.png'
+        else:
+            img = str(i) + '.jpg'
+        path = os.path.join(dir, img)
+        images.append(path)
+    return images
+
+def default_loader(path):
+    return Image.open(path).convert('RGB')
+
+
+class ImageFolder(data.Dataset):
+
+    def __init__(self, root, transform=None, return_paths=False,
+                 loader=default_loader):
+        imgs = make_dataset(root)
+        if len(imgs) == 0:
+            raise(RuntimeError("Found 0 images in: " + root + "\n"
+                               "Supported image extensions are: " +
+                               ",".join(IMG_EXTENSIONS)))
+
+        self.root = root
+        self.imgs = imgs
+        self.transform = transform
+        self.return_paths = return_paths
+        self.loader = loader
+
+    def __getitem__(self, index):
+        path = self.imgs[index]
+        img = self.loader(path)
+        if self.transform is not None:
+            img = self.transform(img)
+        if self.return_paths:
+            return img, path
+        else:
+            return img
+
+    def __len__(self):
+        return len(self.imgs)