init

APDrawingGAN2/data/__init__.py

@@ -0,0 +1,75 @@
+import importlib
+import torch.utils.data
+from data.base_data_loader import BaseDataLoader
+from data.base_dataset import BaseDataset
+
+
+def find_dataset_using_name(dataset_name):
+    # Given the option --dataset_mode [datasetname],
+    # the file "data/datasetname_dataset.py"
+    # will be imported.
+    dataset_filename = "data." + dataset_name + "_dataset"
+    datasetlib = importlib.import_module(dataset_filename)
+
+    # In the file, the class called DatasetNameDataset() will
+    # be instantiated. It has to be a subclass of BaseDataset,
+    # and it is case-insensitive.
+    dataset = None
+    target_dataset_name = dataset_name.replace('_', '') + 'dataset'
+    for name, cls in datasetlib.__dict__.items():
+        if name.lower() == target_dataset_name.lower() \
+           and issubclass(cls, BaseDataset):
+            dataset = cls
+
+    if dataset is None:
+        print("In %s.py, there should be a subclass of BaseDataset with class name that matches %s in lowercase." % (dataset_filename, target_dataset_name))
+        exit(0)
+
+    return dataset
+
+
+def get_option_setter(dataset_name):
+    dataset_class = find_dataset_using_name(dataset_name)
+    return dataset_class.modify_commandline_options
+
+
+def create_dataset(opt):
+    dataset = find_dataset_using_name(opt.dataset_mode)
+    instance = dataset()
+    instance.initialize(opt)
+    print("dataset [%s] was created" % (instance.name()))
+    return instance
+
+
+def CreateDataLoader(opt):
+    data_loader = CustomDatasetDataLoader()
+    data_loader.initialize(opt)
+    return data_loader
+
+
+# Wrapper class of Dataset class that performs
+# multi-threaded data loading
+class CustomDatasetDataLoader(BaseDataLoader):
+    def name(self):
+        return 'CustomDatasetDataLoader'
+
+    def initialize(self, opt):
+        BaseDataLoader.initialize(self, opt)
+        self.dataset = create_dataset(opt)
+        self.dataloader = torch.utils.data.DataLoader(
+            self.dataset,
+            batch_size=opt.batch_size,
+            shuffle=not opt.serial_batches,#in training, serial_batches by default is false, shuffle=true
+            num_workers=int(opt.num_threads))
+
+    def load_data(self):
+        return self
+
+    def __len__(self):
+        return min(len(self.dataset), self.opt.max_dataset_size)
+
+    def __iter__(self):
+        for i, data in enumerate(self.dataloader):
+            if i * self.opt.batch_size >= self.opt.max_dataset_size:
+                break
+            yield data

APDrawingGAN2/data/aligned_dataset.py

@@ -0,0 +1,288 @@
+import os.path
+import random
+import torchvision.transforms as transforms
+import torch
+from data.base_dataset import BaseDataset
+from data.image_folder import make_dataset
+from PIL import Image
+import numpy as np
+import cv2
+import csv
+
+def getfeats(featpath):
+	trans_points = np.empty([5,2],dtype=np.int64) 
+	with open(featpath, 'r') as csvfile:
+		reader = csv.reader(csvfile, delimiter=' ')
+		for ind,row in enumerate(reader):
+			trans_points[ind,:] = row
+	return trans_points
+
+def tocv2(ts):
+    img = (ts.numpy()/2+0.5)*255
+    img = img.astype('uint8')
+    img = np.transpose(img,(1,2,0))
+    img = img[:,:,::-1]#rgb->bgr
+    return img
+
+def dt(img):
+    if(img.shape[2]==3):
+        img = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
+    #convert to BW
+    ret1,thresh1 = cv2.threshold(img,127,255,cv2.THRESH_BINARY)
+    ret2,thresh2 = cv2.threshold(img,127,255,cv2.THRESH_BINARY_INV)
+    dt1 = cv2.distanceTransform(thresh1,cv2.DIST_L2,5)
+    dt2 = cv2.distanceTransform(thresh2,cv2.DIST_L2,5)
+    dt1 = dt1/dt1.max()#->[0,1]
+    dt2 = dt2/dt2.max()
+    return dt1, dt2
+
+def getSoft(size,xb,yb,boundwidth=5.0):
+    xarray = np.tile(np.arange(0,size[1]),(size[0],1))
+    yarray = np.tile(np.arange(0,size[0]),(size[1],1)).transpose()
+    cxdists = []
+    cydists = []
+    for i in range(len(xb)):
+        xba = np.tile(xb[i],(size[1],1)).transpose()
+        yba = np.tile(yb[i],(size[0],1))
+        cxdists.append(np.abs(xarray-xba))
+        cydists.append(np.abs(yarray-yba))
+    xdist = np.minimum.reduce(cxdists)
+    ydist = np.minimum.reduce(cydists)
+    manhdist = np.minimum.reduce([xdist,ydist])
+    im = (manhdist+1) / (boundwidth+1) * 1.0
+    im[im>=1.0] = 1.0
+    return im
+
+class AlignedDataset(BaseDataset):
+    @staticmethod
+    def modify_commandline_options(parser, is_train):
+        return parser
+
+    def initialize(self, opt):
+        self.opt = opt
+        self.root = opt.dataroot
+        imglist = 'datasets/apdrawing_list/%s/%s.txt' % (opt.phase, opt.dataroot)
+        if os.path.exists(imglist):
+            lines = open(imglist, 'r').read().splitlines()
+            lines = sorted(lines)
+            self.AB_paths = [line.split()[0] for line in lines]
+            if len(lines[0].split()) == 2:
+                self.B_paths = [line.split()[1] for line in lines]
+        else:
+            self.dir_AB = os.path.join(opt.dataroot, opt.phase)
+            self.AB_paths = sorted(make_dataset(self.dir_AB))
+        assert(opt.resize_or_crop == 'resize_and_crop')
+
+    def __getitem__(self, index):
+        AB_path = self.AB_paths[index]
+        AB = Image.open(AB_path).convert('RGB')
+        w, h = AB.size
+        if w/h == 2:
+            w2 = int(w / 2)
+            A = AB.crop((0, 0, w2, h)).resize((self.opt.loadSize, self.opt.loadSize), Image.BICUBIC)
+            B = AB.crop((w2, 0, w, h)).resize((self.opt.loadSize, self.opt.loadSize), Image.BICUBIC)
+        else: # if w/h != 2, need B_paths
+            A = AB.resize((self.opt.loadSize, self.opt.loadSize), Image.BICUBIC)
+            B = Image.open(self.B_paths[index]).convert('RGB')
+            B = B.resize((self.opt.loadSize, self.opt.loadSize), Image.BICUBIC)
+        A = transforms.ToTensor()(A)
+        B = transforms.ToTensor()(B)
+        w_offset = random.randint(0, max(0, self.opt.loadSize - self.opt.fineSize - 1))
+        h_offset = random.randint(0, max(0, self.opt.loadSize - self.opt.fineSize - 1))
+
+        A = A[:, h_offset:h_offset + self.opt.fineSize, w_offset:w_offset + self.opt.fineSize]#C,H,W
+        B = B[:, h_offset:h_offset + self.opt.fineSize, w_offset:w_offset + self.opt.fineSize]
+
+        A = transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))(A)
+        B = transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))(B)
+
+        if self.opt.which_direction == 'BtoA':
+            input_nc = self.opt.output_nc
+            output_nc = self.opt.input_nc
+        else:
+            input_nc = self.opt.input_nc
+            output_nc = self.opt.output_nc
+
+        flipped = False
+        if (not self.opt.no_flip) and random.random() < 0.5:
+            flipped = True
+            idx = [i for i in range(A.size(2) - 1, -1, -1)]
+            idx = torch.LongTensor(idx)
+            A = A.index_select(2, idx)
+            B = B.index_select(2, idx)
+
+        if input_nc == 1:  # RGB to gray
+            tmp = A[0, ...] * 0.299 + A[1, ...] * 0.587 + A[2, ...] * 0.114
+            A = tmp.unsqueeze(0)
+
+        if output_nc == 1:  # RGB to gray
+            tmp = B[0, ...] * 0.299 + B[1, ...] * 0.587 + B[2, ...] * 0.114
+            B = tmp.unsqueeze(0)
+        
+        item = {'A': A, 'B': B,
+                'A_paths': AB_path, 'B_paths': AB_path}
+
+        if self.opt.use_local:
+            regions = ['eyel','eyer','nose','mouth']
+            basen = os.path.basename(AB_path)[:-4]+'.txt'
+            if self.opt.region_enm in [0,1]:
+                featdir = self.opt.lm_dir
+                featpath = os.path.join(featdir,basen)
+                feats = getfeats(featpath)
+                if flipped:
+                    for i in range(5):
+                        feats[i,0] = self.opt.fineSize - feats[i,0] - 1
+                    tmp = [feats[0,0],feats[0,1]]
+                    feats[0,:] = [feats[1,0],feats[1,1]]
+                    feats[1,:] = tmp
+                mouth_x = int((feats[3,0]+feats[4,0])/2.0)
+                mouth_y = int((feats[3,1]+feats[4,1])/2.0)
+                ratio = self.opt.fineSize / 256
+                EYE_H = self.opt.EYE_H * ratio
+                EYE_W = self.opt.EYE_W * ratio
+                NOSE_H = self.opt.NOSE_H * ratio
+                NOSE_W = self.opt.NOSE_W * ratio
+                MOUTH_H = self.opt.MOUTH_H * ratio
+                MOUTH_W = self.opt.MOUTH_W * ratio
+                center = torch.IntTensor([[feats[0,0],feats[0,1]-4*ratio],[feats[1,0],feats[1,1]-4*ratio],[feats[2,0],feats[2,1]-NOSE_H/2+16*ratio],[mouth_x,mouth_y]])
+                item['center'] = center
+                rhs = [int(EYE_H),int(EYE_H),int(NOSE_H),int(MOUTH_H)]
+                rws = [int(EYE_W),int(EYE_W),int(NOSE_W),int(MOUTH_W)]
+                if self.opt.soft_border:
+                    soft_border_mask4 = []
+                    for i in range(4):
+                        xb = [np.zeros(rhs[i]),np.ones(rhs[i])*(rws[i]-1)]
+                        yb = [np.zeros(rws[i]),np.ones(rws[i])*(rhs[i]-1)]
+                        soft_border_mask = getSoft([rhs[i],rws[i]],xb,yb)
+                        soft_border_mask4.append(torch.Tensor(soft_border_mask).unsqueeze(0))
+                        item['soft_'+regions[i]+'_mask'] = soft_border_mask4[i]
+                for i in range(4):
+                    item[regions[i]+'_A'] = A[:,center[i,1]-rhs[i]/2:center[i,1]+rhs[i]/2,center[i,0]-rws[i]/2:center[i,0]+rws[i]/2]
+                    item[regions[i]+'_B'] = B[:,center[i,1]-rhs[i]/2:center[i,1]+rhs[i]/2,center[i,0]-rws[i]/2:center[i,0]+rws[i]/2]
+                    if self.opt.soft_border:
+                        item[regions[i]+'_A'] = item[regions[i]+'_A'] * soft_border_mask4[i].repeat(int(input_nc/output_nc),1,1)
+                        item[regions[i]+'_B'] = item[regions[i]+'_B'] * soft_border_mask4[i]
+            if self.opt.compactmask:
+                cmasks0 = []
+                cmasks = []
+                for i in range(4):
+                    if flipped and i in [0,1]:
+                        cmaskpath = os.path.join(self.opt.cmask_dir,regions[1-i],basen[:-4]+'.png')
+                    else:
+                        cmaskpath = os.path.join(self.opt.cmask_dir,regions[i],basen[:-4]+'.png')
+                    im_cmask = Image.open(cmaskpath)
+                    cmask0 = transforms.ToTensor()(im_cmask)
+                    if flipped:
+                        cmask0 = cmask0.index_select(2, idx)
+                    if output_nc == 1 and cmask0.shape[0] == 3:
+                        tmp = cmask0[0, ...] * 0.299 + cmask0[1, ...] * 0.587 + cmask0[2, ...] * 0.114
+                        cmask0 = tmp.unsqueeze(0)
+                    cmask0 = (cmask0 >= 0.5).float()
+                    cmasks0.append(cmask0)
+                    cmask = cmask0.clone()
+                    if self.opt.region_enm in [0,1]:
+                        cmask = cmask[:,center[i,1]-rhs[i]/2:center[i,1]+rhs[i]/2,center[i,0]-rws[i]/2:center[i,0]+rws[i]/2]
+                    elif self.opt.region_enm in [2]: # need to multiply cmask
+                        item[regions[i]+'_A'] = (A/2+0.5) * cmask * 2 - 1
+                        item[regions[i]+'_B'] = (B/2+0.5) * cmask * 2 - 1
+                    cmasks.append(cmask)
+                item['cmaskel'] = cmasks[0]
+                item['cmasker'] = cmasks[1]
+                item['cmask'] = cmasks[2]
+                item['cmaskmo'] = cmasks[3]
+            if self.opt.hair_local:
+                mask = torch.ones(B.shape)
+                if self.opt.region_enm == 0:
+                    for i in range(4):
+                        mask[:,center[i,1]-rhs[i]/2:center[i,1]+rhs[i]/2,center[i,0]-rws[i]/2:center[i,0]+rws[i]/2] = 0
+                    if self.opt.soft_border:
+                        imgsize = self.opt.fineSize
+                        maskn = mask[0].numpy()
+                        masks = [np.ones([imgsize,imgsize]),np.ones([imgsize,imgsize]),np.ones([imgsize,imgsize]),np.ones([imgsize,imgsize])]
+                        masks[0][1:] = maskn[:-1]
+                        masks[1][:-1] = maskn[1:]
+                        masks[2][:,1:] = maskn[:,:-1]
+                        masks[3][:,:-1] = maskn[:,1:]
+                        masks2 = [maskn-e for e in masks]
+                        bound = np.minimum.reduce(masks2)
+                        bound = -bound
+                        xb = []
+                        yb = []
+                        for i in range(4):
+                            xbi = [center[i,0]-rws[i]/2, center[i,0]+rws[i]/2-1]
+                            ybi = [center[i,1]-rhs[i]/2, center[i,1]+rhs[i]/2-1]
+                            for j in range(2):
+                                maskx = bound[:,xbi[j]]
+                                masky = bound[ybi[j],:]
+                                tmp_a = torch.from_numpy(maskx)*xbi[j].double()
+                                tmp_b = torch.from_numpy(1-maskx)
+                                xb += [tmp_b*10000 + tmp_a]
+
+                                tmp_a = torch.from_numpy(masky)*ybi[j].double()
+                                tmp_b = torch.from_numpy(1-masky)
+                                yb += [tmp_b*10000 + tmp_a]
+                        soft = 1-getSoft([imgsize,imgsize],xb,yb)
+                        soft = torch.Tensor(soft).unsqueeze(0)
+                        mask = (torch.ones(mask.shape)-mask)*soft + mask
+                elif self.opt.region_enm == 1:
+                    for i in range(4):
+                        cmask0 = cmasks0[i]
+                        rec = torch.zeros(B.shape)
+                        rec[:,center[i,1]-rhs[i]/2:center[i,1]+rhs[i]/2,center[i,0]-rws[i]/2:center[i,0]+rws[i]/2] = 1
+                        mask = mask * (torch.ones(B.shape) - cmask0 * rec)
+                elif self.opt.region_enm == 2:
+                    for i in range(4):
+                        cmask0 = cmasks0[i]
+                        mask = mask * (torch.ones(B.shape) - cmask0)
+                hair_A = (A/2+0.5) * mask.repeat(int(input_nc/output_nc),1,1) * 2 - 1
+                hair_B = (B/2+0.5) * mask * 2 - 1
+                item['hair_A'] = hair_A
+                item['hair_B'] = hair_B
+                item['mask'] = mask # mask out eyes, nose, mouth
+                if self.opt.bg_local:
+                    bgdir = self.opt.bg_dir
+                    bgpath = os.path.join(bgdir,basen[:-4]+'.png')
+                    im_bg = Image.open(bgpath)
+                    mask2 = transforms.ToTensor()(im_bg) # mask out background
+                    if flipped:
+                        mask2 = mask2.index_select(2, idx)
+                    mask2 = (mask2 >= 0.5).float()
+                    hair_A = (A/2+0.5) * mask.repeat(int(input_nc/output_nc),1,1) * mask2.repeat(int(input_nc/output_nc),1,1) * 2 - 1
+                    hair_B = (B/2+0.5) * mask * mask2 * 2 - 1
+                    bg_A = (A/2+0.5) * (torch.ones(mask2.shape)-mask2).repeat(int(input_nc/output_nc),1,1) * 2 - 1
+                    bg_B = (B/2+0.5) * (torch.ones(mask2.shape)-mask2) * 2 - 1
+                    item['hair_A'] = hair_A
+                    item['hair_B'] = hair_B
+                    item['bg_A'] = bg_A
+                    item['bg_B'] = bg_B
+                    item['mask'] = mask
+                    item['mask2'] = mask2
+        
+        if (self.opt.isTrain and self.opt.chamfer_loss):
+            if self.opt.which_direction == 'AtoB':
+                img = tocv2(B)
+            else:
+                img = tocv2(A)
+            dt1, dt2 = dt(img)
+            dt1 = torch.from_numpy(dt1)
+            dt2 = torch.from_numpy(dt2)
+            dt1 = dt1.unsqueeze(0)
+            dt2 = dt2.unsqueeze(0)
+            item['dt1gt'] = dt1
+            item['dt2gt'] = dt2
+        
+        if self.opt.isTrain and self.opt.emphasis_conti_face:
+            face_mask_path = os.path.join(self.opt.facemask_dir,basen[:-4]+'.png')
+            face_mask = Image.open(face_mask_path)
+            face_mask = transforms.ToTensor()(face_mask) # [0,1]
+            if flipped:
+                face_mask = face_mask.index_select(2, idx)
+            item['face_mask'] = face_mask
+
+        return item
+
+    def __len__(self):
+        return len(self.AB_paths)
+
+    def name(self):
+        return 'AlignedDataset'

APDrawingGAN2/data/base_data_loader.py

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

APDrawingGAN2/data/base_dataset.py

@@ -0,0 +1,103 @@
+import torch.utils.data as data
+from PIL import Image
+import torchvision.transforms as transforms
+
+
+class BaseDataset(data.Dataset):
+    def __init__(self):
+        super(BaseDataset, self).__init__()
+
+    def name(self):
+        return 'BaseDataset'
+
+    @staticmethod
+    def modify_commandline_options(parser, is_train):
+        return parser
+
+    def initialize(self, opt):
+        pass
+
+    def __len__(self):
+        return 0
+
+
+def get_transform(opt):
+    transform_list = []
+    if opt.resize_or_crop == 'resize_and_crop':
+        osize = [opt.loadSize, opt.fineSize]
+        transform_list.append(transforms.Resize(osize, Image.BICUBIC))
+        transform_list.append(transforms.RandomCrop(opt.fineSize))
+    elif opt.resize_or_crop == 'crop':
+        transform_list.append(transforms.RandomCrop(opt.fineSize))
+    elif opt.resize_or_crop == 'scale_width':
+        transform_list.append(transforms.Lambda(
+            lambda img: __scale_width(img, opt.fineSize)))
+    elif opt.resize_or_crop == 'scale_width_and_crop':
+        transform_list.append(transforms.Lambda(
+            lambda img: __scale_width(img, opt.loadSize)))
+        transform_list.append(transforms.RandomCrop(opt.fineSize))
+    elif opt.resize_or_crop == 'none':
+        transform_list.append(transforms.Lambda(
+            lambda img: __adjust(img)))
+    else:
+        raise ValueError('--resize_or_crop %s is not a valid option.' % opt.resize_or_crop)
+
+    if opt.isTrain and not opt.no_flip:
+        transform_list.append(transforms.RandomHorizontalFlip())
+
+    transform_list += [transforms.ToTensor(),
+                       transforms.Normalize((0.5, 0.5, 0.5),
+                                            (0.5, 0.5, 0.5))]
+    return transforms.Compose(transform_list)
+
+# just modify the width and height to be multiple of 4
+def __adjust(img):
+    ow, oh = img.size
+
+    # the size needs to be a multiple of this number, 
+    # because going through generator network may change img size
+    # and eventually cause size mismatch error
+    mult = 4 
+    if ow % mult == 0 and oh % mult == 0:
+        return img
+    w = (ow - 1) // mult
+    w = (w + 1) * mult
+    h = (oh - 1) // mult
+    h = (h + 1) * mult
+
+    if ow != w or oh != h:
+        __print_size_warning(ow, oh, w, h)
+        
+    return img.resize((w, h), Image.BICUBIC)
+
+
+def __scale_width(img, target_width):
+    ow, oh = img.size
+    
+    # the size needs to be a multiple of this number, 
+    # because going through generator network may change img size
+    # and eventually cause size mismatch error    
+    mult = 4
+    assert target_width % mult == 0, "the target width needs to be multiple of %d." % mult
+    if (ow == target_width and oh % mult == 0):
+        return img
+    w = target_width
+    target_height = int(target_width * oh / ow)
+    m = (target_height - 1) // mult
+    h = (m + 1) * mult
+
+    if target_height != h:
+        __print_size_warning(target_width, target_height, w, h)
+    
+    return img.resize((w, h), Image.BICUBIC)
+
+
+def __print_size_warning(ow, oh, w, h):
+    if not hasattr(__print_size_warning, 'has_printed'):
+        print("The image size needs to be a multiple of 4. "
+              "The loaded image size was (%d, %d), so it was adjusted to "
+              "(%d, %d). This adjustment will be done to all images "
+              "whose sizes are not multiples of 4" % (ow, oh, w, h))
+        __print_size_warning.has_printed = True
+
+

APDrawingGAN2/data/image_folder.py

@@ -0,0 +1,68 @@
+###############################################################################
+# Code from
+# https://github.com/pytorch/vision/blob/master/torchvision/datasets/folder.py
+# Modified the original code so that it also loads images from the current
+# directory as well as the subdirectories
+###############################################################################
+
+import torch.utils.data as data
+
+from PIL import Image
+import os
+import os.path
+
+IMG_EXTENSIONS = [
+    '.jpg', '.JPG', '.jpeg', '.JPEG',
+    '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP',
+]
+
+
+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
+
+    for root, _, fnames in sorted(os.walk(dir)):
+        for fname in fnames:
+            if is_image_file(fname):
+                path = os.path.join(root, fname)
+                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)

APDrawingGAN2/data/single_dataset.py

@@ -0,0 +1,176 @@
+import os.path
+from data.base_dataset import BaseDataset, get_transform
+from data.image_folder import make_dataset
+from PIL import Image
+import numpy as np
+import csv
+import torch
+import torchvision.transforms as transforms
+
+def getfeats(featpath):
+	trans_points = np.empty([5,2],dtype=np.int64) 
+	with open(featpath, 'r') as csvfile:
+		reader = csv.reader(csvfile, delimiter=' ')
+		for ind,row in enumerate(reader):
+			trans_points[ind,:] = row
+	return trans_points
+
+def getSoft(size,xb,yb,boundwidth=5.0):
+    xarray = np.tile(np.arange(0,size[1]),(size[0],1))
+    yarray = np.tile(np.arange(0,size[0]),(size[1],1)).transpose()
+    cxdists = []
+    cydists = []
+    for i in range(len(xb)):
+        xba = np.tile(xb[i],(size[1],1)).transpose()
+        yba = np.tile(yb[i],(size[0],1))
+        cxdists.append(np.abs(xarray-xba))
+        cydists.append(np.abs(yarray-yba))
+    xdist = np.minimum.reduce(cxdists)
+    ydist = np.minimum.reduce(cydists)
+    manhdist = np.minimum.reduce([xdist,ydist])
+    im = (manhdist+1) / (boundwidth+1) * 1.0
+    im[im>=1.0] = 1.0
+    return im
+
+class SingleDataset(BaseDataset):
+    @staticmethod
+    def modify_commandline_options(parser, is_train):
+        return parser
+
+    def initialize(self, opt):
+        self.opt = opt
+        self.root = opt.dataroot
+        self.dir_A = os.path.join(opt.dataroot)
+        imglist = 'datasets/apdrawing_list/%s/%s.txt' % (opt.phase, opt.dataroot)
+        if os.path.exists(imglist):
+            lines = open(imglist, 'r').read().splitlines()
+            self.A_paths = sorted(lines)
+        else:
+            self.A_paths = make_dataset(self.dir_A)
+            self.A_paths = sorted(self.A_paths)
+        self.transform = get_transform(opt) # this function uses NO_FLIP; aligned dataset do not use this, aligned dataset manually transform
+
+    def __getitem__(self, index):
+        A_path = self.A_paths[index]
+        A_img = Image.open(A_path).convert('RGB')
+        A = self.transform(A_img)
+        if self.opt.which_direction == 'BtoA':
+            input_nc = self.opt.output_nc
+            output_nc = self.opt.input_nc
+        else:
+            input_nc = self.opt.input_nc
+            output_nc = self.opt.output_nc
+
+        if input_nc == 1:  # RGB to gray
+            tmp = A[0, ...] * 0.299 + A[1, ...] * 0.587 + A[2, ...] * 0.114
+            A = tmp.unsqueeze(0)
+
+        item = {'A': A, 'A_paths': A_path}
+
+        if self.opt.use_local:
+            regions = ['eyel','eyer','nose','mouth']
+            basen = os.path.basename(A_path)[:-4]+'.txt'
+            featdir = self.opt.lm_dir
+            featpath = os.path.join(featdir,basen)
+            feats = getfeats(featpath)
+            mouth_x = int((feats[3,0]+feats[4,0])/2.0)
+            mouth_y = int((feats[3,1]+feats[4,1])/2.0)
+            ratio = self.opt.fineSize / 256
+            EYE_H = self.opt.EYE_H * ratio
+            EYE_W = self.opt.EYE_W * ratio
+            NOSE_H = self.opt.NOSE_H * ratio
+            NOSE_W = self.opt.NOSE_W * ratio
+            MOUTH_H = self.opt.MOUTH_H * ratio
+            MOUTH_W = self.opt.MOUTH_W * ratio
+            center = torch.IntTensor([[feats[0,0],feats[0,1]-4*ratio],[feats[1,0],feats[1,1]-4*ratio],[feats[2,0],feats[2,1]-NOSE_H/2+16*ratio],[mouth_x,mouth_y]])
+            item['center'] = center
+            rhs = [int(EYE_H),int(EYE_H),int(NOSE_H),int(MOUTH_H)]
+            rws = [int(EYE_W),int(EYE_W),int(NOSE_W),int(MOUTH_W)]
+            if self.opt.soft_border:
+                soft_border_mask4 = []
+                for i in range(4):
+                    xb = [np.zeros(rhs[i]),np.ones(rhs[i])*(rws[i]-1)]
+                    yb = [np.zeros(rws[i]),np.ones(rws[i])*(rhs[i]-1)]
+                    soft_border_mask = getSoft([rhs[i],rws[i]],xb,yb)
+                    soft_border_mask4.append(torch.Tensor(soft_border_mask).unsqueeze(0))
+                    item['soft_'+regions[i]+'_mask'] = soft_border_mask4[i]
+            for i in range(4):
+                item[regions[i]+'_A'] = A[:,center[i,1]-rhs[i]/2:center[i,1]+rhs[i]/2,center[i,0]-rws[i]/2:center[i,0]+rws[i]/2]
+                if self.opt.soft_border:
+                    item[regions[i]+'_A'] = item[regions[i]+'_A'] * soft_border_mask4[i].repeat(int(input_nc/output_nc),1,1)
+            if self.opt.compactmask:
+                cmasks0 = []
+                cmasks = []
+                for i in range(4):
+                    cmaskpath = os.path.join(self.opt.cmask_dir,regions[i],basen[:-4]+'.png')
+                    im_cmask = Image.open(cmaskpath)
+                    cmask0 = transforms.ToTensor()(im_cmask)
+                    if output_nc == 1 and cmask0.shape[0] == 3:
+                        tmp = cmask0[0, ...] * 0.299 + cmask0[1, ...] * 0.587 + cmask0[2, ...] * 0.114
+                        cmask0 = tmp.unsqueeze(0)
+                    cmask0 = (cmask0 >= 0.5).float()
+                    cmasks0.append(cmask0)
+                    cmask = cmask0.clone()
+                    cmask = cmask[:,center[i,1]-rhs[i]/2:center[i,1]+rhs[i]/2,center[i,0]-rws[i]/2:center[i,0]+rws[i]/2]
+                    cmasks.append(cmask)
+                item['cmaskel'] = cmasks[0]
+                item['cmasker'] = cmasks[1]
+                item['cmask'] = cmasks[2]
+                item['cmaskmo'] = cmasks[3]
+            if self.opt.hair_local:
+                output_nc = self.opt.output_nc
+                mask = torch.ones([output_nc,A.shape[1],A.shape[2]])
+                for i in range(4):
+                    mask[:,center[i,1]-rhs[i]/2:center[i,1]+rhs[i]/2,center[i,0]-rws[i]/2:center[i,0]+rws[i]/2] = 0
+                if self.opt.soft_border:
+                    imgsize = self.opt.fineSize
+                    maskn = mask[0].numpy()
+                    masks = [np.ones([imgsize,imgsize]),np.ones([imgsize,imgsize]),np.ones([imgsize,imgsize]),np.ones([imgsize,imgsize])]
+                    masks[0][1:] = maskn[:-1]
+                    masks[1][:-1] = maskn[1:]
+                    masks[2][:,1:] = maskn[:,:-1]
+                    masks[3][:,:-1] = maskn[:,1:]
+                    masks2 = [maskn-e for e in masks]
+                    bound = np.minimum.reduce(masks2)
+                    bound = -bound
+                    xb = []
+                    yb = []
+                    for i in range(4):
+                        xbi = [center[i,0]-rws[i]/2, center[i,0]+rws[i]/2-1]
+                        ybi = [center[i,1]-rhs[i]/2, center[i,1]+rhs[i]/2-1]
+                        for j in range(2):
+                            maskx = bound[:,xbi[j]]
+                            masky = bound[ybi[j],:]
+                            tmp_a = torch.from_numpy(maskx)*xbi[j].double()
+                            tmp_b = torch.from_numpy(1-maskx)
+                            xb += [tmp_b*10000 + tmp_a]
+
+                            tmp_a = torch.from_numpy(masky)*ybi[j].double()
+                            tmp_b = torch.from_numpy(1-masky)
+                            yb += [tmp_b*10000 + tmp_a]
+                    soft = 1-getSoft([imgsize,imgsize],xb,yb)
+                    soft = torch.Tensor(soft).unsqueeze(0)
+                    mask = (torch.ones(mask.shape)-mask)*soft + mask
+                hair_A = (A/2+0.5) * mask.repeat(int(input_nc/output_nc),1,1) * 2 - 1
+                item['hair_A'] = hair_A
+                item['mask'] = mask
+                if self.opt.bg_local:
+                    bgdir = self.opt.bg_dir
+                    bgpath = os.path.join(bgdir,basen[:-4]+'.png')
+                    im_bg = Image.open(bgpath)
+                    mask2 = transforms.ToTensor()(im_bg) # mask out background
+                    mask2 = (mask2 >= 0.5).float()
+                    hair_A = (A/2+0.5) * mask.repeat(int(input_nc/output_nc),1,1) * mask2.repeat(int(input_nc/output_nc),1,1) * 2 - 1
+                    bg_A = (A/2+0.5) * (torch.ones(mask2.shape)-mask2).repeat(int(input_nc/output_nc),1,1) * 2 - 1
+                    item['hair_A'] = hair_A
+                    item['bg_A'] = bg_A
+                    item['mask'] = mask
+                    item['mask2'] = mask2
+
+        return item
+
+    def __len__(self):
+        return len(self.A_paths)
+
+    def name(self):
+        return 'SingleImageDataset'

APDrawingGAN2/docs/tips.md

@@ -0,0 +1,8 @@
+## Training/test Tips
+- Flags: see `options/train_options.py` and `options/base_options.py` for the training flags; see `options/test_options.py` and `options/base_options.py` for the test flags. The default values of these options are somtimes adjusted in the model files.
+
+- CPU/GPU (default `--gpu_ids 0`): set`--gpu_ids -1` to use CPU mode; set `--gpu_ids 0,1,2` for multi-GPU mode. You need a large batch size (e.g. `--batch_size 32`) to benefit from multiple GPUs.
+
+- Visualization: during training, the current results can be viewed using two methods. First, if you set `--display_id` > 0, the results and loss plot will appear on a local graphics web server launched by [visdom](https://github.com/facebookresearch/visdom). To do this, you should have `visdom` installed and a server running by the command `python -m visdom.server`. The default server URL is `http://localhost:8097`. `display_id` corresponds to the window ID that is displayed on the `visdom` server. The `visdom` display functionality is turned on by default. To avoid the extra overhead of communicating with `visdom` set `--display_id -1`. Second, the intermediate results are saved to `[opt.checkpoints_dir]/[opt.name]/web/` as an HTML file. To avoid this, set `--no_html`.
+
+- Fine-tuning/Resume training: to fine-tune a pre-trained model, or resume the previous training, use the `--continue_train` flag. The program will then load the model based on `which_epoch`. By default, the program will initialize the epoch count as 1. Set `--epoch_count <int>` to specify a different starting epoch count.

APDrawingGAN2/models/__init__.py

@@ -0,0 +1,39 @@
+import importlib
+from models.base_model import BaseModel
+
+
+def find_model_using_name(model_name):
+    # Given the option --model [modelname],
+    # the file "models/modelname_model.py"
+    # will be imported.
+    model_filename = "models." + model_name + "_model"
+    modellib = importlib.import_module(model_filename)
+
+    # In the file, the class called ModelNameModel() will
+    # be instantiated. It has to be a subclass of BaseModel,
+    # and it is case-insensitive.
+    model = None
+    target_model_name = model_name.replace('_', '') + 'model'
+    for name, cls in modellib.__dict__.items():
+        if name.lower() == target_model_name.lower() \
+           and issubclass(cls, BaseModel):
+            model = cls
+
+    if model is None:
+        print("In %s.py, there should be a subclass of BaseModel with class name that matches %s in lowercase." % (model_filename, target_model_name))
+        exit(0)
+
+    return model
+
+
+def get_option_setter(model_name):
+    model_class = find_model_using_name(model_name)
+    return model_class.modify_commandline_options
+
+
+def create_model(opt):
+    model = find_model_using_name(opt.model)
+    instance = model()
+    instance.initialize(opt)
+    print("model [%s] was created" % (instance.name()))
+    return instance

APDrawingGAN2/models/apdrawingpp_style_model.py

@@ -0,0 +1,692 @@
+import torch
+from util.image_pool import ImagePool
+from .base_model import BaseModel
+from . import networks
+import os
+import math
+
+W = 11
+aa = int(math.floor(512./W))
+res = 512 - W*aa
+
+
+def padpart(A,part,centers,opt,device):
+    IMAGE_SIZE = opt.fineSize
+    bs,nc,_,_ = A.shape
+    ratio = IMAGE_SIZE / 256
+    NOSE_W = opt.NOSE_W * ratio
+    NOSE_H = opt.NOSE_H * ratio
+    EYE_W = opt.EYE_W * ratio
+    EYE_H = opt.EYE_H * ratio
+    MOUTH_W = opt.MOUTH_W * ratio
+    MOUTH_H = opt.MOUTH_H * ratio
+    A_p = torch.ones((bs,nc,IMAGE_SIZE,IMAGE_SIZE)).to(device)
+    padvalue = -1 # black
+    for i in range(bs):
+        center = centers[i]
+        if part == 'nose':
+            A_p[i] = torch.nn.ConstantPad2d((center[2,0] - NOSE_W / 2, IMAGE_SIZE - (center[2,0]+NOSE_W/2), center[2,1] - NOSE_H / 2, IMAGE_SIZE - (center[2,1]+NOSE_H/2)),padvalue)(A[i])
+        elif part == 'eyel':
+            A_p[i] = torch.nn.ConstantPad2d((center[0,0] - EYE_W / 2, IMAGE_SIZE - (center[0,0]+EYE_W/2), center[0,1] - EYE_H / 2, IMAGE_SIZE - (center[0,1]+EYE_H/2)),padvalue)(A[i])
+        elif part == 'eyer':
+            A_p[i] = torch.nn.ConstantPad2d((center[1,0] - EYE_W / 2, IMAGE_SIZE - (center[1,0]+EYE_W/2), center[1,1] - EYE_H / 2, IMAGE_SIZE - (center[1,1]+EYE_H/2)),padvalue)(A[i])
+        elif part == 'mouth':
+            A_p[i] = torch.nn.ConstantPad2d((center[3,0] - MOUTH_W / 2, IMAGE_SIZE - (center[3,0]+MOUTH_W/2), center[3,1] - MOUTH_H / 2, IMAGE_SIZE - (center[3,1]+MOUTH_H/2)),padvalue)(A[i])
+    return A_p
+
+import numpy as np
+def nonlinearDt(dt,type='atan',xmax=torch.Tensor([10.0])):#dt in [0,1], first multiply xmax(>1), then remap to [0,1]
+    if type == 'atan':
+        nldt = torch.atan(dt*xmax) / torch.atan(xmax)
+    elif type == 'sigmoid':
+        nldt = (torch.sigmoid(dt*xmax)-0.5) / (torch.sigmoid(xmax)-0.5)
+    elif type == 'tanh':
+        nldt = torch.tanh(dt*xmax) / torch.tanh(xmax)
+    elif type == 'pow':
+        nldt = torch.pow(dt*xmax,2) / torch.pow(xmax,2)
+    elif type == 'exp':
+        if xmax.item()>1:
+            xmax = xmax / 3
+        nldt = (torch.exp(dt*xmax)-1) / (torch.exp(xmax)-1)
+    #print("remap dt:", type, xmax.item())
+    return nldt
+
+class APDrawingPPStyleModel(BaseModel):
+    def name(self):
+        return 'APDrawingPPStyleModel'
+
+    @staticmethod
+    def modify_commandline_options(parser, is_train=True):
+
+        # changing the default values to match the pix2pix paper
+        # (https://phillipi.github.io/pix2pix/)
+        parser.set_defaults(pool_size=0, no_lsgan=True, norm='batch')# no_lsgan=True, use_lsgan=False
+        parser.set_defaults(dataset_mode='aligned')
+        parser.set_defaults(auxiliary_root='auxiliaryeye2o')
+        parser.set_defaults(use_local=True, hair_local=True, bg_local=True)
+        parser.set_defaults(discriminator_local=True, gan_loss_strategy=2)
+        parser.set_defaults(chamfer_loss=True, dt_nonlinear='exp', lambda_chamfer=0.35, lambda_chamfer2=0.35)
+        parser.set_defaults(nose_ae=True, others_ae=True, compactmask=True, MOUTH_H=56)
+        parser.set_defaults(soft_border=1, batch_size=1, save_epoch_freq=25)
+        parser.add_argument('--nnG_hairc', type=int, default=6, help='nnG for hair classifier')
+        parser.add_argument('--use_resnet', action='store_true', help='use resnet for generator')
+        parser.add_argument('--regarch', type=int, default=4, help='architecture for netRegressor')
+        if is_train:
+            parser.add_argument('--lambda_L1', type=float, default=100.0, help='weight for L1 loss')
+            parser.add_argument('--lambda_local', type=float, default=25.0, help='weight for Local loss')
+            parser.set_defaults(netG_dt='unet_512')
+            parser.set_defaults(netG_line='unet_512')
+
+        return parser
+
+    def initialize(self, opt):
+        BaseModel.initialize(self, opt)
+        self.isTrain = opt.isTrain
+        # specify the training losses you want to print out. The program will call base_model.get_current_losses
+        self.loss_names = ['G_GAN', 'G_L1', 'D_real', 'D_fake']
+        if self.isTrain and self.opt.no_l1_loss:
+            self.loss_names = ['G_GAN', 'D_real', 'D_fake']
+        if self.isTrain and self.opt.use_local and not self.opt.no_G_local_loss:
+            self.loss_names.append('G_local')
+            self.loss_names.append('G_hair_local')
+            self.loss_names.append('G_bg_local')
+        if self.isTrain and self.opt.discriminator_local:
+            self.loss_names.append('D_real_local')
+            self.loss_names.append('D_fake_local')
+            self.loss_names.append('G_GAN_local')
+        if self.isTrain and self.opt.chamfer_loss:
+            self.loss_names.append('G_chamfer')
+            self.loss_names.append('G_chamfer2')
+        if self.isTrain and self.opt.continuity_loss:
+            self.loss_names.append('G_continuity')
+        self.loss_names.append('G')
+        print('loss_names', self.loss_names)
+        # specify the images you want to save/display. The program will call base_model.get_current_visuals
+        self.visual_names = ['real_A', 'fake_B', 'real_B']
+        if self.opt.use_local:
+            self.visual_names += ['fake_B0', 'fake_B1']
+            self.visual_names += ['fake_B_hair', 'real_B_hair', 'real_A_hair']
+            self.visual_names += ['fake_B_bg', 'real_B_bg', 'real_A_bg']
+            if self.opt.region_enm in [0,1]:
+                if self.opt.nose_ae:
+                    self.visual_names += ['fake_B_nose_v','fake_B_nose_v1','fake_B_nose_v2','cmask1no']
+                if self.opt.others_ae:
+                    self.visual_names += ['fake_B_eyel_v','fake_B_eyel_v1','fake_B_eyel_v2','cmask1el']
+                    self.visual_names += ['fake_B_eyer_v','fake_B_eyer_v1','fake_B_eyer_v2','cmask1er']
+                    self.visual_names += ['fake_B_mouth_v','fake_B_mouth_v1','fake_B_mouth_v2','cmask1mo']
+            elif self.opt.region_enm in [2]:
+                self.visual_names += ['fake_B_nose','fake_B_eyel','fake_B_eyer','fake_B_mouth']
+        if self.isTrain and self.opt.chamfer_loss:
+            self.visual_names += ['dt1', 'dt2']
+            self.visual_names += ['dt1gt', 'dt2gt']
+        if self.isTrain and self.opt.soft_border:
+            self.visual_names += ['mask']
+        if not self.isTrain and self.opt.save2:
+            self.visual_names = ['real_A', 'fake_B']
+        print('visuals', self.visual_names)
+        # specify the models you want to save to the disk. The program will call base_model.save_networks and base_model.load_networks
+        self.auxiliary_model_names = []
+        if self.isTrain:
+            self.model_names = ['G', 'D']
+            if self.opt.discriminator_local:
+                self.model_names += ['DLEyel','DLEyer','DLNose','DLMouth','DLHair','DLBG']
+            # auxiliary nets for loss calculation
+            if self.opt.chamfer_loss:
+                self.auxiliary_model_names += ['DT1', 'DT2']
+                self.auxiliary_model_names += ['Line1', 'Line2']
+            if self.opt.continuity_loss:
+                self.auxiliary_model_names += ['Regressor']
+        else:  # during test time, only load Gs
+            self.model_names = ['G']
+            if self.opt.test_continuity_loss:
+                self.auxiliary_model_names += ['Regressor']
+        if self.opt.use_local:
+            self.model_names += ['GLEyel','GLEyer','GLNose','GLMouth','GLHair','GLBG','GCombine']
+            self.auxiliary_model_names += ['CLm','CLh']
+            # auxiliary nets for local output refinement
+            if self.opt.nose_ae:
+                self.auxiliary_model_names += ['AE']
+            if self.opt.others_ae:
+                self.auxiliary_model_names += ['AEel','AEer','AEmowhite','AEmoblack']
+        print('model_names', self.model_names)
+        print('auxiliary_model_names', self.auxiliary_model_names)
+        # load/define networks
+        self.netG = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG, opt.norm,
+                                      not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids,
+                                      opt.nnG)
+        print('netG', opt.netG)
+
+        if self.isTrain:
+            use_sigmoid = opt.no_lsgan
+            self.netD = networks.define_D(opt.input_nc + opt.output_nc, opt.ndf, opt.netD,
+                                          opt.n_layers_D, opt.norm, use_sigmoid, opt.init_type, opt.init_gain, self.gpu_ids)
+            print('netD', opt.netD, opt.n_layers_D)
+            if self.opt.discriminator_local:
+                self.netDLEyel = networks.define_D(opt.input_nc + opt.output_nc, opt.ndf, opt.netD,
+                                          opt.n_layers_D, opt.norm, use_sigmoid, opt.init_type, opt.init_gain, self.gpu_ids)
+                self.netDLEyer = networks.define_D(opt.input_nc + opt.output_nc, opt.ndf, opt.netD,
+                                          opt.n_layers_D, opt.norm, use_sigmoid, opt.init_type, opt.init_gain, self.gpu_ids)
+                self.netDLNose = networks.define_D(opt.input_nc + opt.output_nc, opt.ndf, opt.netD,
+                                          opt.n_layers_D, opt.norm, use_sigmoid, opt.init_type, opt.init_gain, self.gpu_ids)
+                self.netDLMouth = networks.define_D(opt.input_nc + opt.output_nc, opt.ndf, opt.netD,
+                                          opt.n_layers_D, opt.norm, use_sigmoid, opt.init_type, opt.init_gain, self.gpu_ids)
+                self.netDLHair = networks.define_D(opt.input_nc + opt.output_nc, opt.ndf, opt.netD,
+                                          opt.n_layers_D, opt.norm, use_sigmoid, opt.init_type, opt.init_gain, self.gpu_ids)
+                self.netDLBG = networks.define_D(opt.input_nc + opt.output_nc, opt.ndf, opt.netD,
+                                          opt.n_layers_D, opt.norm, use_sigmoid, opt.init_type, opt.init_gain, self.gpu_ids)
+                
+        
+        if self.opt.use_local:
+            netlocal1 = 'partunet' if self.opt.use_resnet == 0 else 'resnet_nblocks'
+            netlocal2 = 'partunet2' if self.opt.use_resnet == 0 else 'resnet_6blocks'
+            netlocal2_style = 'partunet2style' if self.opt.use_resnet == 0 else 'resnet_style2_6blocks'
+            self.netGLEyel = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, netlocal1, opt.norm,
+                                      not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids, nnG=3)
+            self.netGLEyer = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, netlocal1, opt.norm,
+                                      not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids, nnG=3)
+            self.netGLNose = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, netlocal1, opt.norm,
+                                      not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids, nnG=3)
+            self.netGLMouth = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, netlocal1, opt.norm,
+                                        not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids, nnG=3)
+            self.netGLHair = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, netlocal2_style, opt.norm,
+                                      not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids, nnG=4,
+                                      extra_channel=3)
+            self.netGLBG = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, netlocal2, opt.norm,
+                                    not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids, nnG=4)
+            # by default combiner_type is combiner, which uses resnet
+            print('combiner_type', self.opt.combiner_type)
+            self.netGCombine = networks.define_G(2*opt.output_nc, opt.output_nc, opt.ngf, self.opt.combiner_type, opt.norm,
+                                      not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids, 2)
+            # auxiliary classifiers for mouth and hair
+            ratio = self.opt.fineSize / 256
+            self.MOUTH_H = int(self.opt.MOUTH_H * ratio)
+            self.MOUTH_W = int(self.opt.MOUTH_W * ratio)
+            self.netCLm = networks.define_G(opt.input_nc, 2, opt.ngf, 'classifier', opt.norm,
+                                      not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids,
+                                      nnG = 3, ae_h = self.MOUTH_H, ae_w = self.MOUTH_W)
+            self.netCLh = networks.define_G(opt.input_nc, 3, opt.ngf, 'classifier', opt.norm,
+                                      not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids,
+                                      nnG = opt.nnG_hairc, ae_h = opt.fineSize, ae_w = opt.fineSize)
+        
+
+        if self.isTrain:
+            self.fake_AB_pool = ImagePool(opt.pool_size)
+            # define loss functions
+            self.criterionGAN = networks.GANLoss(use_lsgan=not opt.no_lsgan).to(self.device)
+            self.criterionL1 = torch.nn.L1Loss()
+
+            # initialize optimizers
+            self.optimizers = []
+            if not self.opt.use_local:
+                print('G_params 1 components')
+                self.optimizer_G = torch.optim.Adam(self.netG.parameters(),
+                                                lr=opt.lr, betas=(opt.beta1, 0.999))
+            else:
+                G_params = list(self.netG.parameters()) + list(self.netGLEyel.parameters()) + list(self.netGLEyer.parameters()) + list(self.netGLNose.parameters()) + list(self.netGLMouth.parameters()) + list(self.netGCombine.parameters()) + list(self.netGLHair.parameters()) + list(self.netGLBG.parameters())
+                print('G_params 8 components')
+                self.optimizer_G = torch.optim.Adam(G_params,
+                                                lr=opt.lr, betas=(opt.beta1, 0.999))
+            
+            if not self.opt.discriminator_local:
+                print('D_params 1 components')
+                self.optimizer_D = torch.optim.Adam(self.netD.parameters(),
+                                                lr=opt.lr, betas=(opt.beta1, 0.999))
+            else:#self.opt.discriminator_local == True
+                D_params = list(self.netD.parameters()) + list(self.netDLEyel.parameters()) +list(self.netDLEyer.parameters()) + list(self.netDLNose.parameters()) + list(self.netDLMouth.parameters()) + list(self.netDLHair.parameters()) + list(self.netDLBG.parameters())
+                print('D_params 7 components')
+                self.optimizer_D = torch.optim.Adam(D_params,
+                                                lr=opt.lr, betas=(opt.beta1, 0.999))
+            self.optimizers.append(self.optimizer_G)
+            self.optimizers.append(self.optimizer_D)
+        
+        # ==================================auxiliary nets (loaded, parameters fixed)=============================
+        if self.opt.use_local and self.opt.nose_ae:
+            ratio = self.opt.fineSize / 256
+            NOSE_H = self.opt.NOSE_H * ratio
+            NOSE_W = self.opt.NOSE_W * ratio
+            self.netAE = networks.define_G(opt.output_nc, opt.output_nc, opt.ngf, self.opt.nose_ae_net, 'batch',
+                                      not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids,
+                                       latent_dim=self.opt.ae_latentno, ae_h=NOSE_H, ae_w=NOSE_W)
+            self.set_requires_grad(self.netAE, False)            
+        if self.opt.use_local and self.opt.others_ae:
+            ratio = self.opt.fineSize / 256
+            EYE_H = self.opt.EYE_H * ratio
+            EYE_W = self.opt.EYE_W * ratio
+            MOUTH_H = self.opt.MOUTH_H * ratio
+            MOUTH_W = self.opt.MOUTH_W * ratio
+            self.netAEel = networks.define_G(opt.output_nc, opt.output_nc, opt.ngf, self.opt.nose_ae_net, 'batch',
+                                      not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids,
+                                      latent_dim=self.opt.ae_latenteye, ae_h=EYE_H, ae_w=EYE_W)
+            self.netAEer = networks.define_G(opt.output_nc, opt.output_nc, opt.ngf, self.opt.nose_ae_net, 'batch',
+                                      not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids,
+                                      latent_dim=self.opt.ae_latenteye, ae_h=EYE_H, ae_w=EYE_W)
+            self.netAEmowhite = networks.define_G(opt.output_nc, opt.output_nc, opt.ngf, self.opt.nose_ae_net, 'batch',
+                                      not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids,
+                                      latent_dim=self.opt.ae_latentmo, ae_h=MOUTH_H, ae_w=MOUTH_W)
+            self.netAEmoblack = networks.define_G(opt.output_nc, opt.output_nc, opt.ngf, self.opt.nose_ae_net, 'batch',
+                                      not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids,
+                                      latent_dim=self.opt.ae_latentmo, ae_h=MOUTH_H, ae_w=MOUTH_W)
+            self.set_requires_grad(self.netAEel, False)
+            self.set_requires_grad(self.netAEer, False)
+            self.set_requires_grad(self.netAEmowhite, False)
+            self.set_requires_grad(self.netAEmoblack, False)
+            
+        
+        if self.isTrain and self.opt.continuity_loss:
+            self.nc = 1
+            self.netRegressor = networks.define_G(self.nc, 1, opt.ngf, 'regressor', opt.norm,
+                                      not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids_p,
+                                      nnG = opt.regarch)
+            self.set_requires_grad(self.netRegressor, False)
+
+        if self.isTrain and self.opt.chamfer_loss:
+            self.nc = 1
+            self.netDT1 = networks.define_G(self.nc, self.nc, opt.ngf, opt.netG_dt, opt.norm,
+                                      not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids_p)
+            self.netDT2 = networks.define_G(self.nc, self.nc, opt.ngf, opt.netG_dt, opt.norm,
+                                      not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids_p)
+            self.set_requires_grad(self.netDT1, False)
+            self.set_requires_grad(self.netDT2, False)
+            self.netLine1 = networks.define_G(self.nc, self.nc, opt.ngf, opt.netG_line, opt.norm,
+                                    not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids_p)
+            self.netLine2 = networks.define_G(self.nc, self.nc, opt.ngf, opt.netG_line, opt.norm,
+                                    not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids_p)
+            self.set_requires_grad(self.netLine1, False)
+            self.set_requires_grad(self.netLine2, False)
+        
+        # ==================================for test (nets loaded, parameters fixed)=============================
+        if  not self.isTrain and self.opt.test_continuity_loss:
+            self.nc = 1
+            self.netRegressor = networks.define_G(self.nc, 1, opt.ngf, 'regressor', opt.norm,
+                                      not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids,
+                                      nnG = opt.regarch)
+            self.set_requires_grad(self.netRegressor, False)
+        
+
+    def set_input(self, input):
+        AtoB = self.opt.which_direction == 'AtoB'
+        self.real_A = input['A' if AtoB else 'B'].to(self.device)
+        self.real_B = input['B' if AtoB else 'A'].to(self.device)
+        self.image_paths = input['A_paths' if AtoB else 'B_paths']
+        self.batch_size = len(self.image_paths)
+        if self.opt.use_local:
+            self.real_A_eyel = input['eyel_A'].to(self.device)
+            self.real_A_eyer = input['eyer_A'].to(self.device)
+            self.real_A_nose = input['nose_A'].to(self.device)
+            self.real_A_mouth = input['mouth_A'].to(self.device)
+            self.real_B_eyel = input['eyel_B'].to(self.device)
+            self.real_B_eyer = input['eyer_B'].to(self.device)
+            self.real_B_nose = input['nose_B'].to(self.device)
+            self.real_B_mouth = input['mouth_B'].to(self.device)
+            if self.opt.region_enm in [0,1]:
+                self.center = input['center']
+            if self.opt.soft_border:
+                self.softel = input['soft_eyel_mask'].to(self.device)
+                self.softer = input['soft_eyer_mask'].to(self.device)
+                self.softno = input['soft_nose_mask'].to(self.device)
+                self.softmo = input['soft_mouth_mask'].to(self.device)
+            if self.opt.compactmask:
+                self.cmask = input['cmask'].to(self.device)
+                self.cmask1 = self.cmask*2-1#[0,1]->[-1,1]
+                self.cmaskel = input['cmaskel'].to(self.device)
+                self.cmask1el = self.cmaskel*2-1
+                self.cmasker = input['cmasker'].to(self.device)
+                self.cmask1er = self.cmasker*2-1
+                self.cmaskmo = input['cmaskmo'].to(self.device)
+                self.cmask1mo = self.cmaskmo*2-1
+            self.real_A_hair = input['hair_A'].to(self.device)
+            self.real_B_hair = input['hair_B'].to(self.device)
+            self.mask = input['mask'].to(self.device) # mask for non-eyes,nose,mouth
+            self.mask2 = input['mask2'].to(self.device) # mask for non-bg
+            self.real_A_bg = input['bg_A'].to(self.device)
+            self.real_B_bg = input['bg_B'].to(self.device)
+        if (self.isTrain and self.opt.chamfer_loss):
+            self.dt1gt = input['dt1gt'].to(self.device)
+            self.dt2gt = input['dt2gt'].to(self.device)
+        if self.isTrain and self.opt.emphasis_conti_face:
+            self.face_mask = input['face_mask'].cuda(self.gpu_ids_p[0])
+        
+    def getonehot(self,outputs,classes):
+        [maxv,index] = torch.max(outputs,1)
+        y = torch.unsqueeze(index,1)
+        onehot = torch.FloatTensor(self.batch_size,classes).to(self.device)
+        onehot.zero_()
+        onehot.scatter_(1,y,1)
+        return onehot
+
+    def forward(self):
+        if not self.opt.use_local:
+            self.fake_B = self.netG(self.real_A)
+        else:
+            self.fake_B0 = self.netG(self.real_A)
+            # EYES, MOUTH
+            outputs1 = self.netCLm(self.real_A_mouth)
+            onehot1 = self.getonehot(outputs1,2)
+
+            if not self.opt.others_ae:
+                fake_B_eyel = self.netGLEyel(self.real_A_eyel)
+                fake_B_eyer = self.netGLEyer(self.real_A_eyer)
+                fake_B_mouth = self.netGLMouth(self.real_A_mouth)
+            else: # use AE that only constains compact region, need cmask!
+                self.fake_B_eyel1 = self.netGLEyel(self.real_A_eyel)
+                self.fake_B_eyer1 = self.netGLEyer(self.real_A_eyer)
+                self.fake_B_mouth1 = self.netGLMouth(self.real_A_mouth)
+                self.fake_B_eyel2,_ = self.netAEel(self.fake_B_eyel1)
+                self.fake_B_eyer2,_ = self.netAEer(self.fake_B_eyer1)
+                # USE 2 AEs
+                self.fake_B_mouth2 = torch.FloatTensor(self.batch_size,self.opt.output_nc,self.MOUTH_H,self.MOUTH_W).to(self.device)
+                for i in range(self.batch_size):
+                    if onehot1[i][0] == 1:
+                        self.fake_B_mouth2[i],_ = self.netAEmowhite(self.fake_B_mouth1[i].unsqueeze(0))
+                        #print('AEmowhite')
+                    elif onehot1[i][1] == 1:
+                        self.fake_B_mouth2[i],_ = self.netAEmoblack(self.fake_B_mouth1[i].unsqueeze(0))
+                        #print('AEmoblack')
+                fake_B_eyel = self.add_with_mask(self.fake_B_eyel2,self.fake_B_eyel1,self.cmaskel)
+                fake_B_eyer = self.add_with_mask(self.fake_B_eyer2,self.fake_B_eyer1,self.cmasker)
+                fake_B_mouth = self.add_with_mask(self.fake_B_mouth2,self.fake_B_mouth1,self.cmaskmo)
+            # NOSE
+            if not self.opt.nose_ae:
+                fake_B_nose = self.netGLNose(self.real_A_nose)
+            else: # use AE that only constains compact region, need cmask!
+                self.fake_B_nose1 = self.netGLNose(self.real_A_nose)
+                self.fake_B_nose2,_ = self.netAE(self.fake_B_nose1)
+                fake_B_nose = self.add_with_mask(self.fake_B_nose2,self.fake_B_nose1,self.cmask)
+
+            # for visuals and later local loss
+            if self.opt.region_enm in [0,1]:
+                self.fake_B_nose = fake_B_nose
+                self.fake_B_eyel = fake_B_eyel
+                self.fake_B_eyer = fake_B_eyer
+                self.fake_B_mouth = fake_B_mouth
+                # for soft border of 4 rectangle facial feature
+                if self.opt.region_enm == 0 and self.opt.soft_border:
+                    self.fake_B_nose = self.masked(fake_B_nose, self.softno)
+                    self.fake_B_eyel = self.masked(fake_B_eyel, self.softel)
+                    self.fake_B_eyer = self.masked(fake_B_eyer, self.softer)
+                    self.fake_B_mouth = self.masked(fake_B_mouth, self.softmo)
+            elif self.opt.region_enm in [2]: # need to multiply cmask
+                self.fake_B_nose = self.masked(fake_B_nose,self.cmask)
+                self.fake_B_eyel = self.masked(fake_B_eyel,self.cmaskel)
+                self.fake_B_eyer = self.masked(fake_B_eyer,self.cmasker)
+                self.fake_B_mouth = self.masked(fake_B_mouth,self.cmaskmo)
+            
+            # HAIR, BG AND PARTCOMBINE
+            outputs2 = self.netCLh(self.real_A_hair)
+            onehot2 = self.getonehot(outputs2,3)
+
+            if not self.isTrain:
+                opt = self.opt
+                if opt.imagefolder == 'images':
+                    file_name = os.path.join(opt.results_dir, opt.name, '%s_%s' % (opt.phase, opt.which_epoch), 'styleonehot.txt')
+                else:
+                    file_name = os.path.join(opt.results_dir, opt.name, '%s_%s' % (opt.phase, opt.which_epoch), opt.imagefolder, 'styleonehot.txt')
+                message = '%s [%d %d] [%d %d %d]' % (self.image_paths[0], onehot1[0][0], onehot1[0][1], 
+                onehot2[0][0], onehot2[0][1], onehot2[0][2])
+                with open(file_name, 'a+') as s_file:
+                    s_file.write(message)
+                    s_file.write('\n')
+
+            fake_B_hair = self.netGLHair(self.real_A_hair,onehot2)
+            fake_B_bg = self.netGLBG(self.real_A_bg)
+            self.fake_B_hair = self.masked(fake_B_hair,self.mask*self.mask2)
+            self.fake_B_bg = self.masked(fake_B_bg,self.inverse_mask(self.mask2))
+            if not self.opt.compactmask:
+                self.fake_B1 = self.partCombiner2_bg(fake_B_eyel,fake_B_eyer,fake_B_nose,fake_B_mouth,fake_B_hair,fake_B_bg,self.mask*self.mask2,self.inverse_mask(self.mask2),self.opt.comb_op)
+            else:
+                self.fake_B1 = self.partCombiner2_bg(fake_B_eyel,fake_B_eyer,fake_B_nose,fake_B_mouth,fake_B_hair,fake_B_bg,self.mask*self.mask2,self.inverse_mask(self.mask2),self.opt.comb_op,self.opt.region_enm,self.cmaskel,self.cmasker,self.cmask,self.cmaskmo)
+            
+            self.fake_B = self.netGCombine(torch.cat([self.fake_B0,self.fake_B1],1))
+
+            # for AE visuals
+            if self.opt.region_enm in [0,1]:
+                if self.opt.nose_ae:
+                    self.fake_B_nose_v = padpart(self.fake_B_nose, 'nose', self.center, self.opt, self.device)
+                    self.fake_B_nose_v1 = padpart(self.fake_B_nose1, 'nose', self.center, self.opt, self.device)
+                    self.fake_B_nose_v2 = padpart(self.fake_B_nose2, 'nose', self.center, self.opt, self.device)
+                    self.cmask1no = padpart(self.cmask1, 'nose', self.center, self.opt, self.device)
+                if self.opt.others_ae:
+                    self.fake_B_eyel_v = padpart(self.fake_B_eyel, 'eyel', self.center, self.opt, self.device)
+                    self.fake_B_eyel_v1 = padpart(self.fake_B_eyel1, 'eyel', self.center, self.opt, self.device)
+                    self.fake_B_eyel_v2 = padpart(self.fake_B_eyel2, 'eyel', self.center, self.opt, self.device)
+                    self.cmask1el = padpart(self.cmask1el, 'eyel', self.center, self.opt, self.device)
+                    self.fake_B_eyer_v = padpart(self.fake_B_eyer, 'eyer', self.center, self.opt, self.device)
+                    self.fake_B_eyer_v1 = padpart(self.fake_B_eyer1, 'eyer', self.center, self.opt, self.device)
+                    self.fake_B_eyer_v2 = padpart(self.fake_B_eyer2, 'eyer', self.center, self.opt, self.device)
+                    self.cmask1er = padpart(self.cmask1er, 'eyer', self.center, self.opt, self.device)
+                    self.fake_B_mouth_v = padpart(self.fake_B_mouth, 'mouth', self.center, self.opt, self.device)
+                    self.fake_B_mouth_v1 = padpart(self.fake_B_mouth1, 'mouth', self.center, self.opt, self.device)
+                    self.fake_B_mouth_v2 = padpart(self.fake_B_mouth2, 'mouth', self.center, self.opt, self.device)
+                    self.cmask1mo = padpart(self.cmask1mo, 'mouth', self.center, self.opt, self.device)
+            
+            if not self.isTrain and self.opt.test_continuity_loss:
+                self.ContinuityForTest(real=1)
+    
+        
+    def backward_D(self):
+        # Fake
+        # stop backprop to the generator by detaching fake_B
+        fake_AB = self.fake_AB_pool.query(torch.cat((self.real_A, self.fake_B), 1))
+        #print('fake_AB', fake_AB.shape) # (1,4,512,512)
+        pred_fake = self.netD(fake_AB.detach())# by detach, not affect G's gradient
+        self.loss_D_fake = self.criterionGAN(pred_fake, False)
+        if self.opt.discriminator_local:
+            fake_AB_parts = self.getLocalParts(fake_AB)
+            local_names = ['DLEyel','DLEyer','DLNose','DLMouth','DLHair','DLBG']
+            self.loss_D_fake_local = 0
+            for i in range(len(fake_AB_parts)):
+                net = getattr(self, 'net' + local_names[i])
+                pred_fake_tmp = net(fake_AB_parts[i].detach())
+                addw = self.getaddw(local_names[i])
+                self.loss_D_fake_local = self.loss_D_fake_local + self.criterionGAN(pred_fake_tmp, False) * addw
+            self.loss_D_fake = self.loss_D_fake + self.loss_D_fake_local
+
+        # Real
+        real_AB = torch.cat((self.real_A, self.real_B), 1)
+        pred_real = self.netD(real_AB)
+        self.loss_D_real = self.criterionGAN(pred_real, True)
+        if self.opt.discriminator_local:
+            real_AB_parts = self.getLocalParts(real_AB)
+            local_names = ['DLEyel','DLEyer','DLNose','DLMouth','DLHair','DLBG']
+            self.loss_D_real_local = 0
+            for i in range(len(real_AB_parts)):
+                net = getattr(self, 'net' + local_names[i])
+                pred_real_tmp = net(real_AB_parts[i])
+                addw = self.getaddw(local_names[i])
+                self.loss_D_real_local = self.loss_D_real_local + self.criterionGAN(pred_real_tmp, True) * addw
+            self.loss_D_real = self.loss_D_real + self.loss_D_real_local
+
+        # Combined loss
+        self.loss_D = (self.loss_D_fake + self.loss_D_real) * 0.5
+
+        self.loss_D.backward()
+
+    def backward_G(self):
+        # First, G(A) should fake the discriminator
+        fake_AB = torch.cat((self.real_A, self.fake_B), 1)
+        pred_fake = self.netD(fake_AB) # (1,4,512,512)->(1,1,30,30)
+        self.loss_G_GAN = self.criterionGAN(pred_fake, True)
+        if self.opt.discriminator_local:
+            fake_AB_parts = self.getLocalParts(fake_AB)
+            local_names = ['DLEyel','DLEyer','DLNose','DLMouth','DLHair','DLBG']
+            self.loss_G_GAN_local = 0 # G_GAN_local is then added into G_GAN
+            for i in range(len(fake_AB_parts)):
+                net = getattr(self, 'net' + local_names[i])
+                pred_fake_tmp = net(fake_AB_parts[i])
+                addw = self.getaddw(local_names[i])
+                self.loss_G_GAN_local = self.loss_G_GAN_local + self.criterionGAN(pred_fake_tmp, True) * addw
+            if self.opt.gan_loss_strategy == 1:
+                self.loss_G_GAN = (self.loss_G_GAN + self.loss_G_GAN_local) / (len(fake_AB_parts) + 1)
+            elif self.opt.gan_loss_strategy == 2:
+                self.loss_G_GAN_local = self.loss_G_GAN_local * 0.25
+                self.loss_G_GAN = self.loss_G_GAN + self.loss_G_GAN_local
+
+        # Second, G(A) = B
+        if not self.opt.no_l1_loss:
+            self.loss_G_L1 = self.criterionL1(self.fake_B, self.real_B) * self.opt.lambda_L1
+        
+        if self.opt.use_local and not self.opt.no_G_local_loss:
+            local_names = ['eyel','eyer','nose','mouth']
+            self.loss_G_local = 0
+            for i in range(len(local_names)):
+                fakeblocal = getattr(self, 'fake_B_' + local_names[i])
+                realblocal = getattr(self, 'real_B_' + local_names[i])
+                addw = self.getaddw(local_names[i])
+                self.loss_G_local = self.loss_G_local + self.criterionL1(fakeblocal,realblocal) * self.opt.lambda_local * addw
+            self.loss_G_hair_local = self.criterionL1(self.fake_B_hair, self.real_B_hair) * self.opt.lambda_local * self.opt.addw_hair
+            self.loss_G_bg_local = self.criterionL1(self.fake_B_bg, self.real_B_bg) * self.opt.lambda_local * self.opt.addw_bg
+
+        # Third, chamfer matching (assume chamfer_2way and chamfer_only_line is true)
+        if self.opt.chamfer_loss:
+            if self.fake_B.shape[1] == 3:
+                tmp = self.fake_B[:,0,...]*0.299+self.fake_B[:,1,...]*0.587+self.fake_B[:,2,...]*0.114
+                fake_B_gray = tmp.unsqueeze(1)
+            else:
+                fake_B_gray = self.fake_B
+            if self.real_B.shape[1] == 3:
+                tmp = self.real_B[:,0,...]*0.299+self.real_B[:,1,...]*0.587+self.real_B[:,2,...]*0.114
+                real_B_gray = tmp.unsqueeze(1)
+            else:
+                real_B_gray = self.real_B
+            
+            gpu_p = self.opt.gpu_ids_p[0]
+            gpu = self.opt.gpu_ids[0]
+            if gpu_p != gpu:
+                fake_B_gray = fake_B_gray.cuda(gpu_p)
+                real_B_gray = real_B_gray.cuda(gpu_p)
+
+            # d_CM(a_i,G(p_i))
+            self.dt1 = self.netDT1(fake_B_gray)
+            self.dt2 = self.netDT2(fake_B_gray)
+            dt1 = self.dt1/2.0+0.5#[-1,1]->[0,1]
+            dt2 = self.dt2/2.0+0.5
+            if self.opt.dt_nonlinear != '':
+                dt_xmax = torch.Tensor([self.opt.dt_xmax]).cuda(gpu_p)
+                dt1 = nonlinearDt(dt1, self.opt.dt_nonlinear, dt_xmax)
+                dt2 = nonlinearDt(dt2, self.opt.dt_nonlinear, dt_xmax)
+                #print('dt1dt2',torch.min(dt1).item(),torch.max(dt1).item(),torch.min(dt2).item(),torch.max(dt2).item())
+            
+            bs = real_B_gray.shape[0]
+            real_B_gray_line1 = self.netLine1(real_B_gray)
+            real_B_gray_line2 = self.netLine2(real_B_gray)
+            self.loss_G_chamfer = (dt1[(real_B_gray<0)&(real_B_gray_line1<0)].sum() + dt2[(real_B_gray>=0)&(real_B_gray_line2>=0)].sum()) / bs * self.opt.lambda_chamfer
+            if gpu_p != gpu:
+                self.loss_G_chamfer = self.loss_G_chamfer.cuda(gpu) 
+
+            # d_CM(G(p_i),a_i)
+            if gpu_p != gpu:
+                dt1gt = self.dt1gt.cuda(gpu_p)
+                dt2gt = self.dt2gt.cuda(gpu_p)
+            else:
+                dt1gt = self.dt1gt
+                dt2gt = self.dt2gt
+            if self.opt.dt_nonlinear != '':
+                dt1gt = nonlinearDt(dt1gt, self.opt.dt_nonlinear, dt_xmax)
+                dt2gt = nonlinearDt(dt2gt, self.opt.dt_nonlinear, dt_xmax)
+                #print('dt1gtdt2gt',torch.min(dt1gt).item(),torch.max(dt1gt).item(),torch.min(dt2gt).item(),torch.max(dt2gt).item())
+            self.dt1gt = (self.dt1gt-0.5)*2
+            self.dt2gt = (self.dt2gt-0.5)*2
+
+            fake_B_gray_line1 = self.netLine1(fake_B_gray)
+            fake_B_gray_line2 = self.netLine2(fake_B_gray)
+            self.loss_G_chamfer2 = (dt1gt[(fake_B_gray<0)&(fake_B_gray_line1<0)].sum() + dt2gt[(fake_B_gray>=0)&(fake_B_gray_line2>=0)].sum()) / bs * self.opt.lambda_chamfer2
+            if gpu_p != gpu:
+                self.loss_G_chamfer2 = self.loss_G_chamfer2.cuda(gpu)
+
+        # Fourth, line continuity loss, constrained on synthesized drawing
+        if self.opt.continuity_loss:
+            # Patch-based
+            self.get_patches()
+            self.outputs = self.netRegressor(self.fake_B_patches)
+            if not self.opt.emphasis_conti_face:
+                self.loss_G_continuity = (1.0-torch.mean(self.outputs)).cuda(gpu) * self.opt.lambda_continuity
+            else:
+                self.loss_G_continuity = torch.mean((1.0-self.outputs)*self.conti_weights).cuda(gpu) * self.opt.lambda_continuity
+
+
+
+        self.loss_G = self.loss_G_GAN
+        if 'G_L1' in self.loss_names:
+            self.loss_G = self.loss_G + self.loss_G_L1
+        if 'G_local' in self.loss_names:
+            self.loss_G = self.loss_G + self.loss_G_local
+        if 'G_hair_local' in self.loss_names:
+            self.loss_G = self.loss_G + self.loss_G_hair_local
+        if 'G_bg_local' in self.loss_names:
+            self.loss_G = self.loss_G + self.loss_G_bg_local
+        if 'G_chamfer' in self.loss_names:
+            self.loss_G = self.loss_G + self.loss_G_chamfer
+        if 'G_chamfer2' in self.loss_names:
+            self.loss_G = self.loss_G + self.loss_G_chamfer2
+        if 'G_continuity' in self.loss_names:
+            self.loss_G = self.loss_G + self.loss_G_continuity
+
+        self.loss_G.backward()
+
+    def optimize_parameters(self):
+        self.forward()
+        # update D
+        self.set_requires_grad(self.netD, True)
+        
+        if self.opt.discriminator_local:
+            self.set_requires_grad(self.netDLEyel, True)
+            self.set_requires_grad(self.netDLEyer, True)
+            self.set_requires_grad(self.netDLNose, True)
+            self.set_requires_grad(self.netDLMouth, True)
+            self.set_requires_grad(self.netDLHair, True)
+            self.set_requires_grad(self.netDLBG, True)
+        self.optimizer_D.zero_grad()
+        self.backward_D()
+        self.optimizer_D.step()
+
+        # update G
+        self.set_requires_grad(self.netD, False)
+        if self.opt.discriminator_local:
+            self.set_requires_grad(self.netDLEyel, False)
+            self.set_requires_grad(self.netDLEyer, False)
+            self.set_requires_grad(self.netDLNose, False)
+            self.set_requires_grad(self.netDLMouth, False)
+            self.set_requires_grad(self.netDLHair, False)
+            self.set_requires_grad(self.netDLBG, False)
+        self.optimizer_G.zero_grad()
+        self.backward_G()
+        self.optimizer_G.step()
+
+    def get_patches(self):
+        gpu_p = self.opt.gpu_ids_p[0]
+        gpu = self.opt.gpu_ids[0]
+        if gpu_p != gpu:
+            self.fake_B = self.fake_B.cuda(gpu_p)
+        # [1,1,512,512]->[bs,1,11,11]
+        patches = []
+        if self.isTrain and self.opt.emphasis_conti_face:
+            weights = []
+            W2 = int(W/2)
+        t = np.random.randint(res,size=2)
+        for i in range(aa):
+            for j in range(aa):
+                p = self.fake_B[:,:,t[0]+i*W:t[0]+(i+1)*W,t[1]+j*W:t[1]+(j+1)*W]
+                whitenum = torch.sum(p>=0.0)
+                #if whitenum < 5 or whitenum > W*W-5:
+                if whitenum < 1 or whitenum > W*W-1:
+                    continue
+                patches.append(p)
+                if self.isTrain and self.opt.emphasis_conti_face:
+                    weights.append(self.face_mask[:,:,t[0]+i*W+W2,t[1]+j*W+W2])
+        self.fake_B_patches = torch.cat(patches, dim=0)
+        if self.isTrain and self.opt.emphasis_conti_face:
+            self.conti_weights = torch.cat(weights, dim=0)+1 #0->1,1->2
+    
+    def get_patches_real(self):
+        # [1,1,512,512]->[bs,1,11,11]
+        patches = []
+        t = np.random.randint(res,size=2)
+        for i in range(aa):
+            for j in range(aa):
+                p = self.real_B[:,:,t[0]+i*W:t[0]+(i+1)*W,t[1]+j*W:t[1]+(j+1)*W]
+                whitenum = torch.sum(p>=0.0)
+                #if whitenum < 5 or whitenum > W*W-5:
+                if whitenum < 1 or whitenum > W*W-1:  
+                    continue
+                patches.append(p)
+        self.real_B_patches = torch.cat(patches, dim=0)

APDrawingGAN2/models/base_model.py

@@ -0,0 +1,545 @@
+import os
+import torch
+from collections import OrderedDict
+from . import networks
+
+
+class BaseModel():
+
+    # modify parser to add command line options,
+    # and also change the default values if needed
+    @staticmethod
+    def modify_commandline_options(parser, is_train):
+        return parser
+
+    def name(self):
+        return 'BaseModel'
+
+    def initialize(self, opt):
+        self.opt = opt
+        self.gpu_ids = opt.gpu_ids
+        self.gpu_ids_p = opt.gpu_ids_p
+        self.isTrain = opt.isTrain
+        self.device = torch.device('cuda:{}'.format(self.gpu_ids[0])) if self.gpu_ids else torch.device('cpu')
+        self.device_p = torch.device('cuda:{}'.format(self.gpu_ids_p[0])) if self.gpu_ids else torch.device('cpu')
+        self.save_dir = os.path.join(opt.checkpoints_dir, opt.name)
+        self.auxiliary_dir = os.path.join(opt.checkpoints_dir, opt.auxiliary_root)
+        if opt.resize_or_crop != 'scale_width':
+            torch.backends.cudnn.benchmark = True
+        self.loss_names = []
+        self.model_names = []
+        self.visual_names = []
+        self.image_paths = []
+
+    def set_input(self, input):
+        self.input = input
+
+    def forward(self):
+        pass
+
+    # load and print networks; create schedulers
+    def setup(self, opt, parser=None):
+        if self.isTrain:
+            self.schedulers = [networks.get_scheduler(optimizer, opt) for optimizer in self.optimizers]
+
+        if not self.isTrain or opt.continue_train:
+            self.load_networks(opt.which_epoch)
+        if len(self.auxiliary_model_names) > 0:
+            self.load_auxiliary_networks()
+        self.print_networks(opt.verbose)
+
+    # make models eval mode during test time
+    def eval(self):
+        for name in self.model_names:
+            if isinstance(name, str):
+                net = getattr(self, 'net' + name)
+                net.eval()
+
+    # used in test time, wrapping `forward` in no_grad() so we don't save
+    # intermediate steps for backprop
+    def test(self):
+        with torch.no_grad():
+            self.forward()
+
+    # get image paths
+    def get_image_paths(self):
+        return self.image_paths
+
+    def optimize_parameters(self):
+        pass
+
+    # update learning rate (called once every epoch)
+    def update_learning_rate(self):
+        for scheduler in self.schedulers:
+            scheduler.step()
+        lr = self.optimizers[0].param_groups[0]['lr']
+        print('learning rate = %.7f' % lr)
+
+    # return visualization images. train.py will display these images, and save the images to a html
+    def get_current_visuals(self):
+        visual_ret = OrderedDict()
+        for name in self.visual_names:
+            if isinstance(name, str):
+                visual_ret[name] = getattr(self, name)
+        return visual_ret
+
+    # return traning losses/errors. train.py will print out these errors as debugging information
+    def get_current_losses(self):
+        errors_ret = OrderedDict()
+        for name in self.loss_names:
+            if isinstance(name, str):
+                # float(...) works for both scalar tensor and float number
+                errors_ret[name] = float(getattr(self, 'loss_' + name))
+        return errors_ret
+
+    # save models to the disk
+    def save_networks(self, which_epoch):
+        for name in self.model_names:
+            if isinstance(name, str):
+                save_filename = '%s_net_%s.pth' % (which_epoch, name)
+                save_path = os.path.join(self.save_dir, save_filename)
+                net = getattr(self, 'net' + name)
+
+                if len(self.gpu_ids) > 0 and torch.cuda.is_available():
+                    torch.save(net.module.cpu().state_dict(), save_path)
+                    net.cuda(self.gpu_ids[0])
+                else:
+                    torch.save(net.cpu().state_dict(), save_path)
+    
+    def save_networks2(self, which_epoch):
+        gen_name = os.path.join(self.save_dir, '%s_net_gen.pt' % (which_epoch))
+        dis_name = os.path.join(self.save_dir, '%s_net_dis.pt' % (which_epoch))
+        dict_gen = {}
+        dict_dis = {}
+        for name in self.model_names:
+            if isinstance(name, str):
+                net = getattr(self, 'net' + name)
+
+                if len(self.gpu_ids) > 0 and torch.cuda.is_available():
+                    state_dict = net.module.cpu().state_dict()
+                    net.cuda(self.gpu_ids[0])
+                else:
+                    state_dict = net.cpu().state_dict()
+                
+                if name[0] == 'G':
+                    dict_gen[name] = state_dict
+                elif name[0] == 'D':
+                    dict_dis[name] = state_dict
+                else:
+                    save_filename = '%s_net_%s.pth' % (which_epoch, name)
+                    save_path = os.path.join(self.save_dir, save_filename)
+                    torch.save(state_dict, save_path)
+        if dict_gen:
+            torch.save(dict_gen, gen_name)
+        if dict_dis:
+            torch.save(dict_dis, dis_name)
+
+    def __patch_instance_norm_state_dict(self, state_dict, module, keys, i=0):
+        key = keys[i]
+        if i + 1 == len(keys):  # at the end, pointing to a parameter/buffer
+            if module.__class__.__name__.startswith('InstanceNorm') and \
+                    (key == 'running_mean' or key == 'running_var'):
+                if getattr(module, key) is None:
+                    state_dict.pop('.'.join(keys))
+            if module.__class__.__name__.startswith('InstanceNorm') and \
+               (key == 'num_batches_tracked'):
+                state_dict.pop('.'.join(keys))
+        else:
+            self.__patch_instance_norm_state_dict(state_dict, getattr(module, key), keys, i + 1)
+
+    # load models from the disk
+    def load_networks(self, which_epoch):
+        gen_name = os.path.join(self.save_dir, '%s_net_gen.pt' % (which_epoch))
+        if os.path.exists(gen_name):
+            self.load_networks2(which_epoch)
+            return
+        for name in self.model_names:
+            if isinstance(name, str):
+                load_filename = '%s_net_%s.pth' % (which_epoch, name)
+                load_path = os.path.join(self.save_dir, load_filename)
+                net = getattr(self, 'net' + name)
+                if isinstance(net, torch.nn.DataParallel):
+                    net = net.module
+                print('loading the model from %s' % load_path)
+                # if you are using PyTorch newer than 0.4 (e.g., built from
+                # GitHub source), you can remove str() on self.device
+                state_dict = torch.load(load_path, map_location=str(self.device))
+                if hasattr(state_dict, '_metadata'):
+                    del state_dict._metadata
+
+                # patch InstanceNorm checkpoints prior to 0.4
+                for key in list(state_dict.keys()):  # need to copy keys here because we mutate in loop
+                    self.__patch_instance_norm_state_dict(state_dict, net, key.split('.'))
+                net.load_state_dict(state_dict)
+    
+    def load_networks2(self, which_epoch):
+        gen_name = os.path.join(self.save_dir, '%s_net_gen.pt' % (which_epoch))
+        gen_state_dict = torch.load(gen_name, map_location=str(self.device))
+        if self.isTrain and self.opt.model != 'apdrawing_style_nogan':
+            dis_name = os.path.join(self.save_dir, '%s_net_dis.pt' % (which_epoch))
+            dis_state_dict = torch.load(dis_name, map_location=str(self.device))
+        for name in self.model_names:
+            if isinstance(name, str):
+                net = getattr(self, 'net' + name)
+                if isinstance(net, torch.nn.DataParallel):
+                    net = net.module
+                if name[0] == 'G':
+                    print('loading the model %s from %s' % (name,gen_name))
+                    state_dict = gen_state_dict[name]
+                elif name[0] == 'D':
+                    print('loading the model %s from %s' % (name,gen_name))
+                    state_dict = dis_state_dict[name]
+                
+                if hasattr(state_dict, '_metadata'):
+                    del state_dict._metadata
+                # patch InstanceNorm checkpoints prior to 0.4
+                for key in list(state_dict.keys()):  # need to copy keys here because we mutate in loop
+                    self.__patch_instance_norm_state_dict(state_dict, net, key.split('.'))
+                net.load_state_dict(state_dict)
+    
+    # load auxiliary net models from the disk
+    def load_auxiliary_networks(self):
+        for name in self.auxiliary_model_names:
+            if isinstance(name, str):
+                if 'AE' in name and self.opt.ae_small:
+                    load_filename = '%s_net_%s_small.pth' % ('latest', name)
+                elif 'Regressor' in name:
+                    load_filename = '%s_net_%s%d.pth' % ('latest', name, self.opt.regarch)
+                else:
+                    load_filename = '%s_net_%s.pth' % ('latest', name)
+                load_path = os.path.join(self.auxiliary_dir, load_filename)
+                net = getattr(self, 'net' + name)
+                if isinstance(net, torch.nn.DataParallel):
+                    net = net.module
+                print('loading the model from %s' % load_path)
+                # if you are using PyTorch newer than 0.4 (e.g., built from
+                # GitHub source), you can remove str() on self.device
+                if name in ['DT1', 'DT2', 'Line1', 'Line2', 'Continuity1', 'Continuity2', 'Regressor', 'Regressorhair', 'Regressorface']:
+                    state_dict = torch.load(load_path, map_location=str(self.device_p))
+                else:
+                    state_dict = torch.load(load_path, map_location=str(self.device))
+                if hasattr(state_dict, '_metadata'):
+                    del state_dict._metadata
+
+                # patch InstanceNorm checkpoints prior to 0.4
+                for key in list(state_dict.keys()):  # need to copy keys here because we mutate in loop
+                    self.__patch_instance_norm_state_dict(state_dict, net, key.split('.'))
+                net.load_state_dict(state_dict)
+
+    # print network information
+    def print_networks(self, verbose):
+        print('---------- Networks initialized -------------')
+        for name in self.model_names:
+            if isinstance(name, str):
+                net = getattr(self, 'net' + name)
+                num_params = 0
+                for param in net.parameters():
+                    num_params += param.numel()
+                if verbose:
+                    print(net)
+                print('[Network %s] Total number of parameters : %.3f M' % (name, num_params / 1e6))
+        print('-----------------------------------------------')
+
+    # set requies_grad=Fasle to avoid computation
+    def set_requires_grad(self, nets, requires_grad=False):
+        if not isinstance(nets, list):
+            nets = [nets]
+        for net in nets:
+            if net is not None:
+                for param in net.parameters():
+                    param.requires_grad = requires_grad
+
+    # =============================================================================================================
+    def inverse_mask(self, mask):
+        return torch.ones(mask.shape).to(self.device)-mask
+    
+    def masked(self, A,mask):
+        return (A/2+0.5)*mask*2-1
+    
+    def add_with_mask(self, A,B,mask):
+        return ((A/2+0.5)*mask+(B/2+0.5)*(torch.ones(mask.shape).to(self.device)-mask))*2-1
+    
+    def addone_with_mask(self, A,mask):
+        return ((A/2+0.5)*mask+(torch.ones(mask.shape).to(self.device)-mask))*2-1
+
+    def partCombiner(self, eyel, eyer, nose, mouth, average_pos=False, comb_op = 1, region_enm = 0, cmaskel = None, cmasker = None, cmaskno = None, cmaskmo = None):
+        '''
+        x         y
+        100.571   123.429
+        155.429   123.429
+        128.000   155.886
+        103.314   185.417
+        152.686   185.417
+        this is the mean locaiton of 5 landmarks (for 256x256)
+        Pad2d Left,Right,Top,Down
+        '''
+        if comb_op == 0:
+            # use max pooling, pad black for eyes etc
+            padvalue = -1
+            if region_enm in [1,2]:
+                eyel = eyel * cmaskel
+                eyer = eyer * cmasker
+                nose = nose * cmaskno
+                mouth = mouth * cmaskmo
+        else:
+            # use min pooling, pad white for eyes etc
+            padvalue = 1
+            if region_enm in [1,2]:
+                eyel = self.addone_with_mask(eyel, cmaskel)
+                eyer = self.addone_with_mask(eyer, cmasker)
+                nose = self.addone_with_mask(nose, cmaskno)
+                mouth = self.addone_with_mask(mouth, cmaskmo)
+        if region_enm in [0,1]: # need to pad
+            IMAGE_SIZE = self.opt.fineSize
+            ratio = IMAGE_SIZE / 256
+            EYE_W = self.opt.EYE_W * ratio
+            EYE_H = self.opt.EYE_H * ratio
+            NOSE_W = self.opt.NOSE_W * ratio
+            NOSE_H = self.opt.NOSE_H * ratio
+            MOUTH_W = self.opt.MOUTH_W * ratio
+            MOUTH_H = self.opt.MOUTH_H * ratio
+            bs,nc,_,_ = eyel.shape
+            eyel_p = torch.ones((bs,nc,IMAGE_SIZE,IMAGE_SIZE)).to(self.device)
+            eyer_p = torch.ones((bs,nc,IMAGE_SIZE,IMAGE_SIZE)).to(self.device)
+            nose_p = torch.ones((bs,nc,IMAGE_SIZE,IMAGE_SIZE)).to(self.device)
+            mouth_p = torch.ones((bs,nc,IMAGE_SIZE,IMAGE_SIZE)).to(self.device)
+            for i in range(bs):
+                if not average_pos:
+                    center = self.center[i]#x,y
+                else:# if average_pos = True
+                    center = torch.tensor([[101,123-4],[155,123-4],[128,156-NOSE_H/2+16],[128,185]])
+                eyel_p[i] = torch.nn.ConstantPad2d((int(center[0,0] - EYE_W / 2 - 1), int(IMAGE_SIZE - (center[0,0]+EYE_W/2-1)), int(center[0,1] - EYE_H / 2 - 1),int(IMAGE_SIZE - (center[0,1]+EYE_H/2 - 1))),-1)(eyel[i])
+                eyer_p[i] = torch.nn.ConstantPad2d((int(center[1,0] - EYE_W / 2 - 1), int(IMAGE_SIZE - (center[1,0]+EYE_W/2-1)), int(center[1,1] - EYE_H / 2 - 1), int(IMAGE_SIZE - (center[1,1]+EYE_H/2 - 1))),-1)(eyer[i])
+                nose_p[i] = torch.nn.ConstantPad2d((int(center[2,0] - NOSE_W / 2 - 1), int(IMAGE_SIZE - (center[2,0]+NOSE_W/2-1)), int(center[2,1] - NOSE_H / 2 - 1), int(IMAGE_SIZE - (center[2,1]+NOSE_H/2 - 1))),-1)(nose[i])
+                mouth_p[i] = torch.nn.ConstantPad2d((int(center[3,0] - MOUTH_W / 2 - 1), int(IMAGE_SIZE - (center[3,0]+MOUTH_W/2-1)), int(center[3,1] - MOUTH_H / 2 - 1), int(IMAGE_SIZE - (center[3,1]+MOUTH_H/2 - 1))),-1)(mouth[i])
+        elif region_enm in [2]:
+            eyel_p = eyel
+            eyer_p = eyer
+            nose_p = nose
+            mouth_p = mouth
+        if comb_op == 0:
+            # use max pooling
+            eyes = torch.max(eyel_p, eyer_p)
+            eye_nose = torch.max(eyes, nose_p)
+            result = torch.max(eye_nose, mouth_p)
+        else:
+            # use min pooling
+            eyes = torch.min(eyel_p, eyer_p)
+            eye_nose = torch.min(eyes, nose_p)
+            result = torch.min(eye_nose, mouth_p)
+        return result
+    
+    def partCombiner2(self, eyel, eyer, nose, mouth, hair, mask, comb_op = 1, region_enm = 0, cmaskel = None, cmasker = None, cmaskno = None, cmaskmo = None):
+        if comb_op == 0:
+            # use max pooling, pad black for eyes etc
+            padvalue = -1
+            hair = self.masked(hair, mask)
+            if region_enm in [1,2]:
+                eyel = eyel * cmaskel
+                eyer = eyer * cmasker
+                nose = nose * cmaskno
+                mouth = mouth * cmaskmo
+        else:
+            # use min pooling, pad white for eyes etc
+            padvalue = 1
+            hair = self.addone_with_mask(hair, mask)
+            if region_enm in [1,2]:
+                eyel = self.addone_with_mask(eyel, cmaskel)
+                eyer = self.addone_with_mask(eyer, cmasker)
+                nose = self.addone_with_mask(nose, cmaskno)
+                mouth = self.addone_with_mask(mouth, cmaskmo)
+        if region_enm in [0,1]: # need to pad
+            IMAGE_SIZE = self.opt.fineSize
+            ratio = IMAGE_SIZE / 256
+            EYE_W = self.opt.EYE_W * ratio
+            EYE_H = self.opt.EYE_H * ratio
+            NOSE_W = self.opt.NOSE_W * ratio
+            NOSE_H = self.opt.NOSE_H * ratio
+            MOUTH_W = self.opt.MOUTH_W * ratio
+            MOUTH_H = self.opt.MOUTH_H * ratio
+            bs,nc,_,_ = eyel.shape
+            eyel_p = torch.ones((bs,nc,IMAGE_SIZE,IMAGE_SIZE)).to(self.device)
+            eyer_p = torch.ones((bs,nc,IMAGE_SIZE,IMAGE_SIZE)).to(self.device)
+            nose_p = torch.ones((bs,nc,IMAGE_SIZE,IMAGE_SIZE)).to(self.device)
+            mouth_p = torch.ones((bs,nc,IMAGE_SIZE,IMAGE_SIZE)).to(self.device)
+            for i in range(bs):
+                center = self.center[i]#x,y
+                eyel_p[i] = torch.nn.ConstantPad2d((center[0,0] - EYE_W / 2, IMAGE_SIZE - (center[0,0]+EYE_W/2), center[0,1] - EYE_H / 2, IMAGE_SIZE - (center[0,1]+EYE_H/2)),padvalue)(eyel[i])
+                eyer_p[i] = torch.nn.ConstantPad2d((center[1,0] - EYE_W / 2, IMAGE_SIZE - (center[1,0]+EYE_W/2), center[1,1] - EYE_H / 2, IMAGE_SIZE - (center[1,1]+EYE_H/2)),padvalue)(eyer[i])
+                nose_p[i] = torch.nn.ConstantPad2d((center[2,0] - NOSE_W / 2, IMAGE_SIZE - (center[2,0]+NOSE_W/2), center[2,1] - NOSE_H / 2, IMAGE_SIZE - (center[2,1]+NOSE_H/2)),padvalue)(nose[i])
+                mouth_p[i] = torch.nn.ConstantPad2d((center[3,0] - MOUTH_W / 2, IMAGE_SIZE - (center[3,0]+MOUTH_W/2), center[3,1] - MOUTH_H / 2, IMAGE_SIZE - (center[3,1]+MOUTH_H/2)),padvalue)(mouth[i])
+        elif region_enm in [2]:
+            eyel_p = eyel
+            eyer_p = eyer
+            nose_p = nose
+            mouth_p = mouth
+        if comb_op == 0:
+            # use max pooling
+            eyes = torch.max(eyel_p, eyer_p)
+            eye_nose = torch.max(eyes, nose_p)
+            eye_nose_mouth = torch.max(eye_nose, mouth_p)
+            result = torch.max(hair,eye_nose_mouth)
+        else:
+            # use min pooling
+            eyes = torch.min(eyel_p, eyer_p)
+            eye_nose = torch.min(eyes, nose_p)
+            eye_nose_mouth = torch.min(eye_nose, mouth_p)
+            result = torch.min(hair,eye_nose_mouth)
+        return result
+    
+    def partCombiner2_bg(self, eyel, eyer, nose, mouth, hair, bg, maskh, maskb, comb_op = 1, region_enm = 0, cmaskel = None, cmasker = None, cmaskno = None, cmaskmo = None):
+        if comb_op == 0:
+            # use max pooling, pad black for eyes etc
+            padvalue = -1
+            hair = self.masked(hair, maskh)
+            bg = self.masked(bg, maskb)
+            if region_enm in [1,2]:
+                eyel = eyel * cmaskel
+                eyer = eyer * cmasker
+                nose = nose * cmaskno
+                mouth = mouth * cmaskmo
+        else:
+            # use min pooling, pad white for eyes etc
+            padvalue = 1
+            hair = self.addone_with_mask(hair, maskh)
+            bg = self.addone_with_mask(bg, maskb)
+            if region_enm in [1,2]:
+                eyel = self.addone_with_mask(eyel, cmaskel)
+                eyer = self.addone_with_mask(eyer, cmasker)
+                nose = self.addone_with_mask(nose, cmaskno)
+                mouth = self.addone_with_mask(mouth, cmaskmo)
+        if region_enm in [0,1]: # need to pad to full size
+            IMAGE_SIZE = self.opt.fineSize
+            ratio = IMAGE_SIZE / 256
+            EYE_W = self.opt.EYE_W * ratio
+            EYE_H = self.opt.EYE_H * ratio
+            NOSE_W = self.opt.NOSE_W * ratio
+            NOSE_H = self.opt.NOSE_H * ratio
+            MOUTH_W = self.opt.MOUTH_W * ratio
+            MOUTH_H = self.opt.MOUTH_H * ratio
+            bs,nc,_,_ = eyel.shape
+            eyel_p = torch.ones((bs,nc,IMAGE_SIZE,IMAGE_SIZE)).to(self.device)
+            eyer_p = torch.ones((bs,nc,IMAGE_SIZE,IMAGE_SIZE)).to(self.device)
+            nose_p = torch.ones((bs,nc,IMAGE_SIZE,IMAGE_SIZE)).to(self.device)
+            mouth_p = torch.ones((bs,nc,IMAGE_SIZE,IMAGE_SIZE)).to(self.device)
+            for i in range(bs):
+                center = self.center[i]#x,y
+                eyel_p[i] = torch.nn.ConstantPad2d((center[0,0] - EYE_W / 2, IMAGE_SIZE - (center[0,0]+EYE_W/2), center[0,1] - EYE_H / 2, IMAGE_SIZE - (center[0,1]+EYE_H/2)),padvalue)(eyel[i])
+                eyer_p[i] = torch.nn.ConstantPad2d((center[1,0] - EYE_W / 2, IMAGE_SIZE - (center[1,0]+EYE_W/2), center[1,1] - EYE_H / 2, IMAGE_SIZE - (center[1,1]+EYE_H/2)),padvalue)(eyer[i])
+                nose_p[i] = torch.nn.ConstantPad2d((center[2,0] - NOSE_W / 2, IMAGE_SIZE - (center[2,0]+NOSE_W/2), center[2,1] - NOSE_H / 2, IMAGE_SIZE - (center[2,1]+NOSE_H/2)),padvalue)(nose[i])
+                mouth_p[i] = torch.nn.ConstantPad2d((center[3,0] - MOUTH_W / 2, IMAGE_SIZE - (center[3,0]+MOUTH_W/2), center[3,1] - MOUTH_H / 2, IMAGE_SIZE - (center[3,1]+MOUTH_H/2)),padvalue)(mouth[i])
+        elif region_enm in [2]:
+            eyel_p = eyel
+            eyer_p = eyer
+            nose_p = nose
+            mouth_p = mouth
+        if comb_op == 0:
+            eyes = torch.max(eyel_p, eyer_p)
+            eye_nose = torch.max(eyes, nose_p)
+            eye_nose_mouth = torch.max(eye_nose, mouth_p)
+            eye_nose_mouth_hair = torch.max(hair,eye_nose_mouth)
+            result = torch.max(bg,eye_nose_mouth_hair)
+        else:
+            eyes = torch.min(eyel_p, eyer_p)
+            eye_nose = torch.min(eyes, nose_p)
+            eye_nose_mouth = torch.min(eye_nose, mouth_p)
+            eye_nose_mouth_hair = torch.min(hair,eye_nose_mouth)
+            result = torch.min(bg,eye_nose_mouth_hair)
+        return result
+    
+    def partCombiner3(self, face, hair, maskf, maskh, comb_op = 1):
+        if comb_op == 0:
+            # use max pooling, pad black etc
+            padvalue = -1
+            face = self.masked(face, maskf)
+            hair = self.masked(hair, maskh)
+        else:
+            # use min pooling, pad white etc
+            padvalue = 1
+            face = self.addone_with_mask(face, maskf)
+            hair = self.addone_with_mask(hair, maskh)
+        if comb_op == 0:
+            result = torch.max(face,hair)
+        else:
+            result = torch.min(face,hair)
+        return result
+
+
+    def tocv2(ts):
+        img = (ts.numpy()/2+0.5)*255
+        img = img.astype('uint8')
+        img = np.transpose(img,(1,2,0))
+        img = img[:,:,::-1]#rgb->bgr
+        return img
+    
+    def totor(img):
+        img = img[:,:,::-1]
+        tor = transforms.ToTensor()(img)
+        tor = transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))(tor)
+        return tor
+    
+    
+    def ContinuityForTest(self, real = 0):
+        # Patch-based
+        self.get_patches()
+        self.outputs = self.netRegressor(self.fake_B_patches)
+        line_continuity = torch.mean(self.outputs)
+        opt = self.opt
+        file_name = os.path.join(opt.results_dir, opt.name, '%s_%s' % (opt.phase, opt.which_epoch), 'continuity.txt')
+        message = '%s %.04f' % (self.image_paths[0], line_continuity)
+        with open(file_name, 'a+') as c_file:
+            c_file.write(message)
+            c_file.write('\n')
+        if real == 1:
+            self.get_patches_real()
+            self.outputs2 = self.netRegressor(self.real_B_patches)
+            line_continuity2 = torch.mean(self.outputs2)
+            file_name = os.path.join(opt.results_dir, opt.name, '%s_%s' % (opt.phase, opt.which_epoch), 'continuity-r.txt')
+            message = '%s %.04f' % (self.image_paths[0], line_continuity2)
+            with open(file_name, 'a+') as c_file:
+                c_file.write(message)
+                c_file.write('\n')
+    
+    def getLocalParts(self,fakeAB):
+        bs,nc,_,_ = fakeAB.shape #dtype torch.float32
+        ncr = int(nc / self.opt.output_nc)
+        if self.opt.region_enm in [0,1]:
+            ratio = self.opt.fineSize / 256
+            EYE_H = self.opt.EYE_H * ratio
+            EYE_W = self.opt.EYE_W * ratio
+            NOSE_H = self.opt.NOSE_H * ratio
+            NOSE_W = self.opt.NOSE_W * ratio
+            MOUTH_H = self.opt.MOUTH_H * ratio
+            MOUTH_W = self.opt.MOUTH_W * ratio
+            eyel = torch.ones((bs,nc,int(EYE_H),int(EYE_W))).to(self.device)
+            eyer = torch.ones((bs,nc,int(EYE_H),int(EYE_W))).to(self.device)
+            nose = torch.ones((bs,nc,int(NOSE_H),int(NOSE_W))).to(self.device)
+            mouth = torch.ones((bs,nc,int(MOUTH_H),int(MOUTH_W))).to(self.device)
+            for i in range(bs):
+                center = self.center[i]
+                eyel[i] = fakeAB[i,:,center[0,1]-EYE_H/2:center[0,1]+EYE_H/2,center[0,0]-EYE_W/2:center[0,0]+EYE_W/2]
+                eyer[i] = fakeAB[i,:,center[1,1]-EYE_H/2:center[1,1]+EYE_H/2,center[1,0]-EYE_W/2:center[1,0]+EYE_W/2]
+                nose[i] = fakeAB[i,:,center[2,1]-NOSE_H/2:center[2,1]+NOSE_H/2,center[2,0]-NOSE_W/2:center[2,0]+NOSE_W/2]
+                mouth[i] = fakeAB[i,:,center[3,1]-MOUTH_H/2:center[3,1]+MOUTH_H/2,center[3,0]-MOUTH_W/2:center[3,0]+MOUTH_W/2]
+        elif self.opt.region_enm in [2]:
+            eyel = (fakeAB/2+0.5) * self.cmaskel.repeat(1,ncr,1,1) * 2 - 1
+            eyer = (fakeAB/2+0.5) * self.cmasker.repeat(1,ncr,1,1) * 2 - 1
+            nose = (fakeAB/2+0.5) * self.cmask.repeat(1,ncr,1,1) * 2 - 1
+            mouth = (fakeAB/2+0.5) * self.cmaskmo.repeat(1,ncr,1,1) * 2 - 1
+        hair = (fakeAB/2+0.5) * self.mask.repeat(1,ncr,1,1) * self.mask2.repeat(1,ncr,1,1) * 2 - 1
+        bg = (fakeAB/2+0.5) * (torch.ones(fakeAB.shape).to(self.device)-self.mask2.repeat(1,ncr,1,1)) * 2 - 1
+        return eyel, eyer, nose, mouth, hair, bg
+    
+    def getaddw(self,local_name):
+        addw = 1
+        if local_name in ['DLEyel','DLEyer','eyel','eyer','DLFace','face']:
+            addw = self.opt.addw_eye
+        elif local_name in ['DLNose', 'nose']:
+            addw = self.opt.addw_nose
+        elif local_name in ['DLMouth', 'mouth']:
+            addw = self.opt.addw_mouth
+        elif local_name in ['DLHair', 'hair']:
+            addw = self.opt.addw_hair
+        elif local_name in ['DLBG', 'bg']:
+            addw = self.opt.addw_bg
+        return addw

APDrawingGAN2/models/networks.py

@@ -0,0 +1,1194 @@
+import torch
+import torch.nn as nn
+from torch.nn import init
+import functools
+from torch.optim import lr_scheduler
+
+###############################################################################
+# Helper Functions
+###############################################################################
+
+
+def get_norm_layer(norm_type='instance'):
+    if norm_type == 'batch':
+        norm_layer = functools.partial(nn.BatchNorm2d, affine=True)
+    elif norm_type == 'instance':
+        norm_layer = functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=True)
+    elif norm_type == 'none':
+        norm_layer = None
+    else:
+        raise NotImplementedError('normalization layer [%s] is not found' % norm_type)
+    return norm_layer
+
+
+def get_scheduler(optimizer, opt):
+    if opt.lr_policy == 'lambda':
+        def lambda_rule(epoch):
+            lr_l = 1.0 - max(0, epoch + 1 + opt.epoch_count - opt.niter) / float(opt.niter_decay + 1)
+            return lr_l
+        scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule)
+    elif opt.lr_policy == 'step':
+        scheduler = lr_scheduler.StepLR(optimizer, step_size=opt.lr_decay_iters, gamma=0.1)
+    elif opt.lr_policy == 'plateau':
+        scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.2, threshold=0.01, patience=5)
+    elif opt.lr_policy == 'cosine':
+        scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=opt.niter, eta_min=0)
+    else:
+        return NotImplementedError('learning rate policy [%s] is not implemented', opt.lr_policy)
+    return scheduler
+
+
+def init_weights(net, init_type='normal', gain=0.02):
+    def init_func(m):
+        classname = m.__class__.__name__
+        if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):
+            if init_type == 'normal':
+                init.normal_(m.weight.data, 0.0, gain)
+            elif init_type == 'xavier':
+                init.xavier_normal_(m.weight.data, gain=gain)
+            elif init_type == 'kaiming':
+                init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
+            elif init_type == 'orthogonal':
+                init.orthogonal_(m.weight.data, gain=gain)
+            else:
+                raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
+            if hasattr(m, 'bias') and m.bias is not None:
+                init.constant_(m.bias.data, 0.0)
+        elif classname.find('BatchNorm2d') != -1:
+            init.normal_(m.weight.data, 1.0, gain)
+            init.constant_(m.bias.data, 0.0)
+
+    print('initialize network with %s' % init_type)
+    net.apply(init_func)
+
+
+def init_net(net, init_type='normal', init_gain=0.02, gpu_ids=[]):
+    if len(gpu_ids) > 0:
+        assert(torch.cuda.is_available())
+        net.to(gpu_ids[0])
+        net = torch.nn.DataParallel(net, gpu_ids)
+    init_weights(net, init_type, gain=init_gain)
+    return net
+
+
+def define_G(input_nc, output_nc, ngf, netG, norm='batch', use_dropout=False, init_type='normal', init_gain=0.02, gpu_ids=[], nnG=9, multiple=2, latent_dim=1024, ae_h=96, ae_w=96, extra_channel=2, nres=1):
+    net = None
+    norm_layer = get_norm_layer(norm_type=norm)
+
+    if netG == 'autoencoder':
+        net = AutoEncoder(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
+    elif netG == 'autoencoderfc':
+        net = AutoEncoderWithFC(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout,
+        multiple=multiple, latent_dim=latent_dim, h=ae_h, w=ae_w)
+    elif netG == 'autoencoderfc2':
+        net = AutoEncoderWithFC2(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout,
+        multiple=multiple, latent_dim=latent_dim, h=ae_h, w=ae_w)
+    elif netG == 'vae':
+        net = VAE(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout,
+        multiple=multiple, latent_dim=latent_dim, h=ae_h, w=ae_w)
+    elif netG == 'classifier':
+        net = Classifier(input_nc, output_nc, ngf, num_downs=nnG, norm_layer=norm_layer, use_dropout=use_dropout, h=ae_h, w=ae_w)
+    elif netG == 'regressor':
+        net = Regressor(input_nc, ngf, norm_layer=norm_layer, arch=nnG)
+    elif netG == 'resnet_9blocks':#default for cyclegan
+        net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=9)
+    elif netG == 'resnet_6blocks':
+        net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=6)
+    elif netG == 'resnet_nblocks':
+        net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=nnG)
+    elif netG == 'resnet_style2_9blocks':
+        net = ResnetStyle2Generator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=9, model0_res=0, extra_channel=extra_channel)
+    elif netG == 'resnet_style2_6blocks':
+        net = ResnetStyle2Generator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=6, model0_res=0, extra_channel=extra_channel)
+    elif netG == 'resnet_style2_nblocks':
+        net = ResnetStyle2Generator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=nnG, model0_res=0, extra_channel=extra_channel)
+    elif netG == 'unet_128':
+        net = UnetGenerator(input_nc, output_nc, 7, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
+    elif netG == 'unet_256':#default for pix2pix
+        net = UnetGenerator(input_nc, output_nc, 8, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
+    elif netG == 'unet_512':
+        net = UnetGenerator(input_nc, output_nc, 9, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
+    elif netG == 'unet_ndown':
+        net = UnetGenerator(input_nc, output_nc, nnG, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
+    elif netG == 'unetres_ndown':
+        net = UnetResGenerator(input_nc, output_nc, nnG, ngf, norm_layer=norm_layer, use_dropout=use_dropout, nres=nres)
+    elif netG == 'partunet':
+        net = PartUnet(input_nc, output_nc, nnG, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
+    elif netG == 'partunet2':
+        net = PartUnet2(input_nc, output_nc, nnG, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
+    elif netG == 'partunetres':
+        net = PartUnetRes(input_nc, output_nc, nnG, ngf, norm_layer=norm_layer, use_dropout=use_dropout,nres=nres)
+    elif netG == 'partunet2res':
+        net = PartUnet2Res(input_nc, output_nc, nnG, ngf, norm_layer=norm_layer, use_dropout=use_dropout,nres=nres)
+    elif netG == 'partunet2style':
+        net = PartUnet2Style(input_nc, output_nc, nnG, ngf, extra_channel=extra_channel, norm_layer=norm_layer, use_dropout=use_dropout)
+    elif netG == 'partunet2resstyle':
+        net = PartUnet2ResStyle(input_nc, output_nc, nnG, ngf, extra_channel=extra_channel, norm_layer=norm_layer, use_dropout=use_dropout,nres=nres)
+    elif netG == 'combiner':
+         net = Combiner(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=2)
+    elif netG == 'combiner2':
+        net = Combiner2(input_nc, output_nc, nnG, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
+    else:
+        raise NotImplementedError('Generator model name [%s] is not recognized' % netG)
+    return init_net(net, init_type, init_gain, gpu_ids)
+
+
+def define_D(input_nc, ndf, netD,
+             n_layers_D=3, norm='batch', use_sigmoid=False, init_type='normal', init_gain=0.02, gpu_ids=[]):
+    net = None
+    norm_layer = get_norm_layer(norm_type=norm)
+
+    if netD == 'basic':
+        net = NLayerDiscriminator(input_nc, ndf, n_layers=3, norm_layer=norm_layer, use_sigmoid=use_sigmoid)
+    elif netD == 'n_layers':
+        net = NLayerDiscriminator(input_nc, ndf, n_layers_D, norm_layer=norm_layer, use_sigmoid=use_sigmoid)
+    elif netD == 'pixel':
+        net = PixelDiscriminator(input_nc, ndf, norm_layer=norm_layer, use_sigmoid=use_sigmoid)
+    else:
+        raise NotImplementedError('Discriminator model name [%s] is not recognized' % net)
+    return init_net(net, init_type, init_gain, gpu_ids)
+
+
+##############################################################################
+# Classes
+##############################################################################
+
+
+# Defines the GAN loss which uses either LSGAN or the regular GAN.
+# When LSGAN is used, it is basically same as MSELoss,
+# but it abstracts away the need to create the target label tensor
+# that has the same size as the input
+class GANLoss(nn.Module):
+    def __init__(self, use_lsgan=True, target_real_label=1.0, target_fake_label=0.0):
+        super(GANLoss, self).__init__()
+        self.register_buffer('real_label', torch.tensor(target_real_label))
+        self.register_buffer('fake_label', torch.tensor(target_fake_label))
+        if use_lsgan:
+            self.loss = nn.MSELoss()
+        else:#no_lsgan
+            self.loss = nn.BCELoss()
+
+    def get_target_tensor(self, input, target_is_real):
+        if target_is_real:
+            target_tensor = self.real_label
+        else:
+            target_tensor = self.fake_label
+        return target_tensor.expand_as(input)
+
+    def __call__(self, input, target_is_real):
+        target_tensor = self.get_target_tensor(input, target_is_real)
+        return self.loss(input, target_tensor)
+
+
+class AutoEncoderMNIST(nn.Module):
+    def __init__(self):
+        super(AutoEncoderMNIST, self).__init__()
+        self.encoder = nn.Sequential(
+            nn.Conv2d(1, 16, 3, stride=3, padding=1),  # b, 16, 10, 10
+            nn.ReLU(True),
+            nn.MaxPool2d(2, stride=2),  # b, 16, 5, 5
+            nn.Conv2d(16, 8, 3, stride=2, padding=1),  # b, 8, 3, 3
+            nn.ReLU(True),
+            nn.MaxPool2d(2, stride=1)  # b, 8, 2, 2
+        )
+        self.decoder = nn.Sequential(
+            nn.ConvTranspose2d(8, 16, 3, stride=2),  # b, 16, 5, 5
+            nn.ReLU(True),
+            nn.ConvTranspose2d(16, 8, 5, stride=3, padding=1),  # b, 8, 15, 15
+            nn.ReLU(True),
+            nn.ConvTranspose2d(8, 1, 2, stride=2, padding=1),  # b, 1, 28, 28
+            nn.Tanh()
+        )
+
+    def forward(self, x):
+        x = self.encoder(x)
+        x = self.decoder(x)
+        return x
+
+class AutoEncoder(nn.Module):
+    def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, padding_type='reflect'):
+        super(AutoEncoder, self).__init__()
+        self.input_nc = input_nc
+        self.output_nc = output_nc
+        self.ngf = ngf
+        if type(norm_layer) == functools.partial:
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+        
+        model = [nn.Conv2d(input_nc, ngf, kernel_size=4, stride=2, padding=1, bias=use_bias)]
+        n_downsampling = 3
+        for i in range(n_downsampling):
+            mult = 2**i
+            model += [nn.LeakyReLU(0.2),
+                      nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=4,
+                                stride=2, padding=1, bias=use_bias),
+                      norm_layer(ngf * mult * 2)]
+        self.encoder = nn.Sequential(*model)
+
+        model2 = []
+        for i in range(n_downsampling):
+            mult = 2**(n_downsampling - i)
+            model2 += [nn.ReLU(),
+                       nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2),
+                                        kernel_size=4, stride=2,
+                                        padding=1, bias=use_bias),
+                       norm_layer(int(ngf * mult / 2))]
+        model2 += [nn.ReLU()]
+        model2 += [nn.ConvTranspose2d(ngf, output_nc, kernel_size=4, stride=2, padding=1, bias=use_bias)]
+        model2 += [nn.Tanh()]
+        self.decoder = nn.Sequential(*model2)
+
+    def forward(self, x):
+        ax = self.encoder(x) # b, 512, 6, 6
+        y = self.decoder(ax)
+        return y, ax
+
+class AutoEncoderWithFC(nn.Module):
+    def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, multiple=2,latent_dim=1024,  h=96, w=96):
+        super(AutoEncoderWithFC, self).__init__()
+        self.input_nc = input_nc
+        self.output_nc = output_nc
+        self.ngf = ngf
+        if type(norm_layer) == functools.partial:
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+        
+        model = [nn.Conv2d(input_nc, ngf, kernel_size=4, stride=2, padding=1, bias=use_bias)]
+        n_downsampling = 3
+        #multiple = 2
+        for i in range(n_downsampling):
+            mult = multiple**i
+            model += [nn.LeakyReLU(0.2),
+                      nn.Conv2d(int(ngf * mult), int(ngf * mult * multiple), kernel_size=4,
+                                stride=2, padding=1, bias=use_bias),
+                      norm_layer(int(ngf * mult * multiple))]
+        self.encoder = nn.Sequential(*model)
+        self.fc1 = nn.Linear(int(ngf*(multiple**n_downsampling)*h/16*w/16),latent_dim)
+        self.relu = nn.ReLU(latent_dim)
+        self.fc2 = nn.Linear(latent_dim,int(ngf*(multiple**n_downsampling)*h/16*w/16))
+        self.rh = int(h/16)
+        self.rw = int(w/16)
+        model2 = []
+        for i in range(n_downsampling):
+            mult = multiple**(n_downsampling - i)
+            model2 += [nn.ReLU(),
+                       nn.ConvTranspose2d(int(ngf * mult), int(ngf * mult / multiple),
+                                        kernel_size=4, stride=2,
+                                        padding=1, bias=use_bias),
+                       norm_layer(int(ngf * mult / multiple))]
+        model2 += [nn.ReLU()]
+        model2 += [nn.ConvTranspose2d(ngf, output_nc, kernel_size=4, stride=2, padding=1, bias=use_bias)]
+        model2 += [nn.Tanh()]
+        self.decoder = nn.Sequential(*model2)
+
+    def forward(self, x):
+        ax = self.encoder(x) # b, 512, 6, 6
+        ax = ax.view(ax.size(0), -1) # view -- reshape
+        ax = self.relu(self.fc1(ax))
+        ax = self.fc2(ax)
+        ax = ax.view(ax.size(0),-1,self.rh,self.rw)
+        y = self.decoder(ax)
+        return y, ax
+    
+class AutoEncoderWithFC2(nn.Module):
+    def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, multiple=2,latent_dim=1024,  h=96, w=96):
+        super(AutoEncoderWithFC2, self).__init__()
+        self.input_nc = input_nc
+        self.output_nc = output_nc
+        self.ngf = ngf
+        if type(norm_layer) == functools.partial:
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+        
+        model = [nn.Conv2d(input_nc, ngf, kernel_size=4, stride=2, padding=1, bias=use_bias)]
+        n_downsampling = 2
+        #multiple = 2
+        for i in range(n_downsampling):
+            mult = multiple**i
+            model += [nn.LeakyReLU(0.2),
+                      nn.Conv2d(int(ngf * mult), int(ngf * mult * multiple), kernel_size=4,
+                                stride=2, padding=1, bias=use_bias),
+                      norm_layer(int(ngf * mult * multiple))]
+        self.encoder = nn.Sequential(*model)
+        self.fc1 = nn.Linear(int(ngf*(multiple**n_downsampling)*h/8*w/8),latent_dim)
+        self.relu = nn.ReLU(latent_dim)
+        self.fc2 = nn.Linear(latent_dim,int(ngf*(multiple**n_downsampling)*h/8*w/8))
+        self.rh = h/8
+        self.rw = w/8
+        model2 = []
+        for i in range(n_downsampling):
+            mult = multiple**(n_downsampling - i)
+            model2 += [nn.ReLU(),
+                       nn.ConvTranspose2d(int(ngf * mult), int(ngf * mult / multiple),
+                                        kernel_size=4, stride=2,
+                                        padding=1, bias=use_bias),
+                       norm_layer(int(ngf * mult / multiple))]
+        model2 += [nn.ReLU()]
+        model2 += [nn.ConvTranspose2d(ngf, output_nc, kernel_size=4, stride=2, padding=1, bias=use_bias)]
+        model2 += [nn.Tanh()]
+        self.decoder = nn.Sequential(*model2)
+
+    def forward(self, x):
+        ax = self.encoder(x) # b, 256, 12, 12
+        ax = ax.view(ax.size(0), -1) # view -- reshape
+        ax = self.relu(self.fc1(ax))
+        ax = self.fc2(ax)
+        ax = ax.view(ax.size(0),-1,self.rh,self.rw)
+        y = self.decoder(ax)
+        return y, ax
+
+class VAE(nn.Module):
+    def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, multiple=2,latent_dim=1024,  h=96, w=96):
+        super(VAE, self).__init__()
+        self.input_nc = input_nc
+        self.output_nc = output_nc
+        self.ngf = ngf
+        if type(norm_layer) == functools.partial:
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+        
+        model = [nn.Conv2d(input_nc, ngf, kernel_size=4, stride=2, padding=1, bias=use_bias)]
+        n_downsampling = 3
+        for i in range(n_downsampling):
+            mult = multiple**i
+            model += [nn.LeakyReLU(0.2),
+                      nn.Conv2d(int(ngf * mult), int(ngf * mult * multiple), kernel_size=4,
+                                stride=2, padding=1, bias=use_bias),
+                      norm_layer(int(ngf * mult * multiple))]
+        self.encoder_cnn = nn.Sequential(*model)
+
+        self.c_dim = int(ngf*(multiple**n_downsampling)*h/16*w/16)
+        self.rh = h/16
+        self.rw = w/16
+        self.fc1 = nn.Linear(self.c_dim,latent_dim)
+        self.fc2 = nn.Linear(self.c_dim,latent_dim)
+        self.fc3 = nn.Linear(latent_dim,self.c_dim)
+        self.relu = nn.ReLU()
+
+        model2 = []
+        for i in range(n_downsampling):
+            mult = multiple**(n_downsampling - i)
+            model2 += [nn.ReLU(),
+                       nn.ConvTranspose2d(int(ngf * mult), int(ngf * mult / multiple),
+                                        kernel_size=4, stride=2,
+                                        padding=1, bias=use_bias),
+                       norm_layer(int(ngf * mult / multiple))]
+        model2 += [nn.ReLU()]
+        model2 += [nn.ConvTranspose2d(ngf, output_nc, kernel_size=4, stride=2, padding=1, bias=use_bias)]
+        model2 += [nn.Tanh()]#[-1,1]
+        self.decoder_cnn = nn.Sequential(*model2)
+    
+    def encode(self, x):
+        h1 = self.encoder_cnn(x)
+        r1 = h1.view(h1.size(0), -1)
+        return self.fc1(r1), self.fc2(r1)
+    
+    def reparameterize(self, mu, logvar):# not deterministic for test mode
+        std = torch.exp(0.5*logvar)
+        eps = torch.randn_like(std)# torch.rand_like returns a tensor with the same size as input,
+                                   # that is filled with random numbers from a normal distribution N(0,1)
+        return eps.mul(std).add_(mu)
+    
+    def decode(self, z):
+        h4 = self.relu(self.fc3(z))
+        r3 = h4.view(h4.size(0),-1,self.rh,self.rw)
+        return self.decoder_cnn(r3)
+
+    def forward(self, x):
+        mu, logvar = self.encode(x)
+        z = self.reparameterize(mu, logvar)
+        reconx = self.decode(z)
+        return reconx, mu, logvar
+
+class Classifier(nn.Module):
+    def __init__(self, input_nc, classes, ngf=64, num_downs=3, norm_layer=nn.BatchNorm2d, use_dropout=False,
+        h=96, w=96):
+        super(Classifier, self).__init__()
+        self.input_nc = input_nc
+        self.ngf = ngf
+        if type(norm_layer) == functools.partial:
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+        
+        model = [nn.Conv2d(input_nc, ngf, kernel_size=4, stride=2, padding=1, bias=use_bias)]
+        multiple = 2
+        for i in range(num_downs):
+            mult = multiple**i
+            model += [nn.LeakyReLU(0.2),
+                      nn.Conv2d(int(ngf * mult), int(ngf * mult * multiple), kernel_size=4,
+                                stride=2, padding=1, bias=use_bias),
+                      norm_layer(int(ngf * mult * multiple))]
+        self.encoder = nn.Sequential(*model)
+        strides = 2**(num_downs+1)
+        self.fc1 = nn.Linear(int(ngf*h*w/(strides*2)), classes)
+
+    def forward(self, x):
+        ax = self.encoder(x) # b, 512, 6, 6
+        ax = ax.view(ax.size(0), -1) # view -- reshape
+        return self.fc1(ax)
+
+class Regressor(nn.Module):
+    def __init__(self, input_nc, ngf=64, norm_layer=nn.BatchNorm2d, arch=1):
+        super(Regressor, self).__init__()
+        # if use BatchNorm2d, 
+        # no need to use bias as BatchNorm2d has affine parameters
+
+        self.arch = arch
+        
+        if arch == 1:
+            use_bias = True
+            sequence = [
+                nn.Conv2d(input_nc, ngf, kernel_size=3, stride=2, padding=0, bias=use_bias),#11->5
+                nn.LeakyReLU(0.2, True),
+                nn.Conv2d(ngf, 1, kernel_size=5, stride=1, padding=0, bias=use_bias),#5->1
+            ]
+        elif arch == 2:
+            if type(norm_layer) == functools.partial:
+                use_bias = norm_layer.func == nn.InstanceNorm2d
+            else:
+                use_bias = norm_layer == nn.InstanceNorm2d
+            sequence = [
+                nn.Conv2d(input_nc, ngf, kernel_size=3, stride=1, padding=0, bias=use_bias),#11->9
+                nn.LeakyReLU(0.2, True),
+                nn.Conv2d(ngf, ngf*2, kernel_size=3, stride=1, padding=0, bias=use_bias),#9->7
+                norm_layer(ngf*2),
+                nn.LeakyReLU(0.2, True),
+                nn.Conv2d(ngf*2, ngf*4, kernel_size=3, stride=1, padding=0, bias=use_bias),#7->5
+                norm_layer(ngf*4),
+                nn.LeakyReLU(0.2, True),
+                nn.Conv2d(ngf*4, 1, kernel_size=5, stride=1, padding=0, bias=use_bias),#5->1
+            ]
+        elif arch == 3:
+            use_bias = True
+            sequence = [
+                nn.Conv2d(input_nc, ngf, kernel_size=3, stride=1, padding=1, bias=use_bias),#11->11
+                nn.LeakyReLU(0.2, True),
+                nn.Conv2d(ngf, 1, kernel_size=11, stride=1, padding=0, bias=use_bias),#11->1
+            ]
+        elif arch == 4:
+            use_bias = True
+            sequence = [
+                nn.Conv2d(input_nc, ngf, kernel_size=3, stride=1, padding=1, bias=use_bias),#11->11
+                nn.LeakyReLU(0.2, True),
+                nn.Conv2d(ngf, ngf*2, kernel_size=3, stride=1, padding=1, bias=use_bias),#11->11
+                nn.LeakyReLU(0.2, True),
+                nn.Conv2d(ngf*2, ngf*4, kernel_size=3, stride=1, padding=1, bias=use_bias),#11->11
+                nn.LeakyReLU(0.2, True),
+                nn.Conv2d(ngf*4, 1, kernel_size=11, stride=1, padding=0, bias=use_bias),#11->1
+            ]
+        elif arch == 5:
+            use_bias = True
+            sequence = [
+                nn.Conv2d(input_nc, ngf, kernel_size=3, stride=1, padding=1, bias=use_bias),#11->11
+                nn.LeakyReLU(0.2, True),
+                nn.Conv2d(ngf, ngf*2, kernel_size=3, stride=1, padding=1, bias=use_bias),#11->11
+                nn.LeakyReLU(0.2, True),
+                nn.Conv2d(ngf*2, ngf*4, kernel_size=3, stride=1, padding=1, bias=use_bias),#11->11
+                nn.LeakyReLU(0.2, True),
+            ]
+            fc = [
+                nn.Linear(ngf*4*11*11, 4096),
+                nn.ReLU(True),
+                nn.Dropout(),
+                nn.Linear(4096, 1),
+            ]
+            self.fc = nn.Sequential(*fc)
+
+        self.model = nn.Sequential(*sequence)
+    
+    def forward(self, x):
+        if self.arch <= 4:
+            return self.model(x)
+        else:
+            x = self.model(x)
+            x = x.view(x.size(0), -1)
+            x = self.fc(x)
+            return x
+
+
+# Defines the generator that consists of Resnet blocks between a few
+# downsampling/upsampling operations.
+# Code and idea originally from Justin Johnson's architecture.
+# https://github.com/jcjohnson/fast-neural-style/
+class ResnetGenerator(nn.Module):
+    def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, n_blocks=6, padding_type='reflect'):
+        assert(n_blocks >= 0)
+        super(ResnetGenerator, self).__init__()
+        self.input_nc = input_nc
+        self.output_nc = output_nc
+        self.ngf = ngf
+        if type(norm_layer) == functools.partial:
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+
+        model = [nn.ReflectionPad2d(3),
+                 nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0,
+                           bias=use_bias),
+                 norm_layer(ngf),
+                 nn.ReLU(True)]
+
+        n_downsampling = 2
+        for i in range(n_downsampling):
+            mult = 2**i
+            model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3,
+                                stride=2, padding=1, bias=use_bias),
+                      norm_layer(ngf * mult * 2),
+                      nn.ReLU(True)]
+
+        mult = 2**n_downsampling
+        for i in range(n_blocks):
+            model += [ResnetBlock(ngf * mult, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout, use_bias=use_bias)]
+
+        for i in range(n_downsampling):
+            mult = 2**(n_downsampling - i)
+            model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2),
+                                        kernel_size=3, stride=2,
+                                        padding=1, output_padding=1,
+                                        bias=use_bias),
+                    norm_layer(int(ngf * mult / 2)),
+                    nn.ReLU(True)]
+        model += [nn.ReflectionPad2d(3)]
+        model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)]
+        model += [nn.Tanh()]
+
+        self.model = nn.Sequential(*model)
+
+    def forward(self, input):
+        return self.model(input)
+
+class ResnetStyle2Generator(nn.Module):
+    """Resnet-based generator that consists of Resnet blocks between a few downsampling/upsampling operations.
+    We adapt Torch code and idea from Justin Johnson's neural style transfer project(https://github.com/jcjohnson/fast-neural-style)
+    """
+
+    def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, n_blocks=6, padding_type='reflect', extra_channel=3, model0_res=0):
+        """Construct a Resnet-based generator
+
+        Parameters:
+            input_nc (int)      -- the number of channels in input images
+            output_nc (int)     -- the number of channels in output images
+            ngf (int)           -- the number of filters in the last conv layer
+            norm_layer          -- normalization layer
+            use_dropout (bool)  -- if use dropout layers
+            n_blocks (int)      -- the number of ResNet blocks
+            padding_type (str)  -- the name of padding layer in conv layers: reflect | replicate | zero
+        """
+        assert(n_blocks >= 0)
+        super(ResnetStyle2Generator, self).__init__()
+        if type(norm_layer) == functools.partial:
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+
+        model0 = [nn.ReflectionPad2d(3),
+                 nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=use_bias),
+                 norm_layer(ngf),
+                 nn.ReLU(True)]
+
+        n_downsampling = 2
+        for i in range(n_downsampling):  # add downsampling layers
+            mult = 2 ** i
+            model0 += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1, bias=use_bias),
+                      norm_layer(ngf * mult * 2),
+                      nn.ReLU(True)]
+        
+        mult = 2 ** n_downsampling
+        for i in range(model0_res):       # add ResNet blocks
+            model0 += [ResnetBlock(ngf * mult, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout, use_bias=use_bias)]
+
+        model = []
+        model += [nn.Conv2d(ngf * mult + extra_channel, ngf * mult, kernel_size=3, stride=1, padding=1, bias=use_bias),
+                      norm_layer(ngf * mult),
+                      nn.ReLU(True)]
+
+        for i in range(n_blocks-model0_res):       # add ResNet blocks
+            model += [ResnetBlock(ngf * mult, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout, use_bias=use_bias)]
+
+        for i in range(n_downsampling):  # add upsampling layers
+            mult = 2 ** (n_downsampling - i)
+            model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2),
+                                         kernel_size=3, stride=2,
+                                         padding=1, output_padding=1,
+                                         bias=use_bias),
+                      norm_layer(int(ngf * mult / 2)),
+                      nn.ReLU(True)]
+        model += [nn.ReflectionPad2d(3)]
+        model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)]
+        model += [nn.Tanh()]
+
+        self.model0 = nn.Sequential(*model0)
+        self.model = nn.Sequential(*model)
+        print(list(self.modules()))
+
+    def forward(self, input1, input2): # input2 [bs,c]
+        """Standard forward"""
+        f1 = self.model0(input1)
+        [bs,c,h,w] = f1.shape
+        input2 = input2.repeat(h,w,1,1).permute([2,3,0,1])
+        y1 = torch.cat([f1, input2], 1)
+        return self.model(y1)
+
+class Combiner(nn.Module):
+    def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, n_blocks=6, padding_type='reflect'):
+        assert(n_blocks >= 0)
+        super(Combiner, self).__init__()
+        self.input_nc = input_nc
+        self.output_nc = output_nc
+        self.ngf = ngf
+        if type(norm_layer) == functools.partial:
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+
+        model = [nn.ReflectionPad2d(3),
+                 nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0,
+                           bias=use_bias),
+                 norm_layer(ngf),
+                 nn.ReLU(True)]
+
+        for i in range(n_blocks):
+            model += [ResnetBlock(ngf, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout, use_bias=use_bias)]
+
+        model += [nn.ReflectionPad2d(3)]
+        model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)]
+        model += [nn.Tanh()]
+
+        self.model = nn.Sequential(*model)
+
+    def forward(self, input):
+        return self.model(input)
+
+class Combiner2(nn.Module):
+    def __init__(self, input_nc, output_nc, num_downs, ngf=64,
+                 norm_layer=nn.BatchNorm2d, use_dropout=False):
+        super(Combiner2, self).__init__()
+
+        # construct unet structure
+        unet_block = UnetSkipConnectionBlock(ngf, ngf * 2, input_nc=None, submodule=None, norm_layer=norm_layer, innermost=True)
+        unet_block = UnetSkipConnectionBlock(output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, norm_layer=norm_layer)
+
+        self.model = unet_block
+
+    def forward(self, input):
+        return self.model(input)
+
+
+# Define a resnet block
+class ResnetBlock(nn.Module):
+    def __init__(self, dim, padding_type, norm_layer, use_dropout, use_bias):
+        super(ResnetBlock, self).__init__()
+        self.conv_block = self.build_conv_block(dim, padding_type, norm_layer, use_dropout, use_bias)
+
+    def build_conv_block(self, dim, padding_type, norm_layer, use_dropout, use_bias):
+        conv_block = []
+        p = 0
+        if padding_type == 'reflect':
+            conv_block += [nn.ReflectionPad2d(1)]
+        elif padding_type == 'replicate':
+            conv_block += [nn.ReplicationPad2d(1)]
+        elif padding_type == 'zero':
+            p = 1
+        else:
+            raise NotImplementedError('padding [%s] is not implemented' % padding_type)
+
+        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias),
+                       norm_layer(dim),
+                       nn.ReLU(True)]
+        if use_dropout:
+            conv_block += [nn.Dropout(0.5)]
+
+        p = 0
+        if padding_type == 'reflect':
+            conv_block += [nn.ReflectionPad2d(1)]
+        elif padding_type == 'replicate':
+            conv_block += [nn.ReplicationPad2d(1)]
+        elif padding_type == 'zero':
+            p = 1
+        else:
+            raise NotImplementedError('padding [%s] is not implemented' % padding_type)
+        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias),
+                       norm_layer(dim)]
+
+        return nn.Sequential(*conv_block)
+
+    def forward(self, x):
+        out = x + self.conv_block(x)
+        return out
+
+
+# Defines the Unet generator.
+# |num_downs|: number of downsamplings in UNet. For example,
+# if |num_downs| == 7, image of size 128x128 will become of size 1x1
+# at the bottleneck
+class UnetGenerator(nn.Module):
+    def __init__(self, input_nc, output_nc, num_downs, ngf=64,
+                 norm_layer=nn.BatchNorm2d, use_dropout=False):
+        super(UnetGenerator, self).__init__()
+
+        # construct unet structure
+        unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=None, norm_layer=norm_layer, innermost=True)
+        for i in range(num_downs - 5):
+            unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer, use_dropout=use_dropout)
+        unet_block = UnetSkipConnectionBlock(ngf * 4, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+        unet_block = UnetSkipConnectionBlock(ngf * 2, ngf * 4, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+        unet_block = UnetSkipConnectionBlock(ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+        unet_block = UnetSkipConnectionBlock(output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, norm_layer=norm_layer)
+
+        self.model = unet_block
+
+    def forward(self, input):
+        return self.model(input)
+
+class UnetResGenerator(nn.Module):
+    def __init__(self, input_nc, output_nc, num_downs, ngf=64,
+                 norm_layer=nn.BatchNorm2d, use_dropout=False, nres=1):
+        super(UnetResGenerator, self).__init__()
+
+        # construct unet structure
+        unet_block = UnetSkipConnectionResBlock(ngf * 8, ngf * 8, input_nc=None, submodule=None, norm_layer=norm_layer, innermost=True, nres=nres)
+        for i in range(num_downs - 5):
+            unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer, use_dropout=use_dropout)
+        unet_block = UnetSkipConnectionBlock(ngf * 4, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+        unet_block = UnetSkipConnectionBlock(ngf * 2, ngf * 4, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+        unet_block = UnetSkipConnectionBlock(ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+        unet_block = UnetSkipConnectionBlock(output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, norm_layer=norm_layer)
+
+        self.model = unet_block
+
+    def forward(self, input):
+        return self.model(input)
+
+class PartUnet(nn.Module):
+    def __init__(self, input_nc, output_nc, num_downs, ngf=64,
+                 norm_layer=nn.BatchNorm2d, use_dropout=False):
+        super(PartUnet, self).__init__()
+
+        # construct unet structure
+        # 3 downs 
+        unet_block = UnetSkipConnectionBlock(ngf * 2, ngf * 4, input_nc=None, submodule=None, norm_layer=norm_layer, innermost=True)
+        unet_block = UnetSkipConnectionBlock(ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+        unet_block = UnetSkipConnectionBlock(output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, norm_layer=norm_layer)
+
+        self.model = unet_block
+
+    def forward(self, input):
+        return self.model(input)
+
+class PartUnetRes(nn.Module):
+    def __init__(self, input_nc, output_nc, num_downs, ngf=64,
+                 norm_layer=nn.BatchNorm2d, use_dropout=False, nres=1):
+        super(PartUnetRes, self).__init__()
+
+        # construct unet structure
+        # 3 downs 
+        unet_block = UnetSkipConnectionResBlock(ngf * 2, ngf * 4, input_nc=None, submodule=None, norm_layer=norm_layer, innermost=True, nres=nres)
+        unet_block = UnetSkipConnectionBlock(ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+        unet_block = UnetSkipConnectionBlock(output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, norm_layer=norm_layer)
+
+        self.model = unet_block
+
+    def forward(self, input):
+        return self.model(input)
+
+class PartUnet2(nn.Module):
+    def __init__(self, input_nc, output_nc, num_downs, ngf=64,
+                 norm_layer=nn.BatchNorm2d, use_dropout=False):
+        super(PartUnet2, self).__init__()
+
+        # construct unet structure
+        unet_block = UnetSkipConnectionBlock(ngf * 2, ngf * 2, input_nc=None, submodule=None, norm_layer=norm_layer, innermost=True)
+        for i in range(num_downs - 3):
+            unet_block = UnetSkipConnectionBlock(ngf * 2, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer, use_dropout=use_dropout)
+        unet_block = UnetSkipConnectionBlock(ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+        unet_block = UnetSkipConnectionBlock(output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, norm_layer=norm_layer)
+
+        self.model = unet_block
+
+    def forward(self, input):
+        return self.model(input)
+
+class PartUnet2Res(nn.Module):
+    def __init__(self, input_nc, output_nc, num_downs, ngf=64,
+                 norm_layer=nn.BatchNorm2d, use_dropout=False, nres=1):
+        super(PartUnet2Res, self).__init__()
+
+        # construct unet structure
+        unet_block = UnetSkipConnectionResBlock(ngf * 2, ngf * 2, input_nc=None, submodule=None, norm_layer=norm_layer, innermost=True, nres=nres)
+        for i in range(num_downs - 3):
+            unet_block = UnetSkipConnectionBlock(ngf * 2, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer, use_dropout=use_dropout)
+        unet_block = UnetSkipConnectionBlock(ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+        unet_block = UnetSkipConnectionBlock(output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, norm_layer=norm_layer)
+
+        self.model = unet_block
+
+    def forward(self, input):
+        return self.model(input)
+
+class PartUnet2Style(nn.Module):
+    def __init__(self, input_nc, output_nc, num_downs, ngf=64, extra_channel=2,
+                 norm_layer=nn.BatchNorm2d, use_dropout=False):
+        super(PartUnet2Style, self).__init__()
+        # construct unet structure
+        unet_block = UnetSkipConnectionStyleBlock(ngf * 2, ngf * 2, input_nc=None, submodule=None, norm_layer=norm_layer, innermost=True, extra_channel=extra_channel)
+        for i in range(num_downs - 3):
+            unet_block = UnetSkipConnectionStyleBlock(ngf * 2, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer, use_dropout=use_dropout, extra_channel=extra_channel)
+        unet_block = UnetSkipConnectionStyleBlock(ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer, extra_channel=extra_channel)
+        unet_block = UnetSkipConnectionStyleBlock(output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, norm_layer=norm_layer, extra_channel=extra_channel)
+
+        self.model = unet_block
+
+    def forward(self, input, cate):
+        return self.model(input, cate)
+
+class PartUnet2ResStyle(nn.Module):
+    def __init__(self, input_nc, output_nc, num_downs, ngf=64, extra_channel=2,
+                 norm_layer=nn.BatchNorm2d, use_dropout=False, nres=1):
+        super(PartUnet2ResStyle, self).__init__()
+        # construct unet structure
+        unet_block = UnetSkipConnectionResStyleBlock(ngf * 2, ngf * 2, input_nc=None, submodule=None, norm_layer=norm_layer, innermost=True, extra_channel=extra_channel, nres=nres)
+        for i in range(num_downs - 3):
+            unet_block = UnetSkipConnectionStyleBlock(ngf * 2, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer, use_dropout=use_dropout, extra_channel=extra_channel)
+        unet_block = UnetSkipConnectionStyleBlock(ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer, extra_channel=extra_channel)
+        unet_block = UnetSkipConnectionStyleBlock(output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, norm_layer=norm_layer, extra_channel=extra_channel)
+
+        self.model = unet_block
+
+    def forward(self, input, cate):
+        return self.model(input, cate)
+
+
+# Defines the submodule with skip connection.
+# X -------------------identity---------------------- X
+#   |-- downsampling -- |submodule| -- upsampling --|
+class UnetSkipConnectionBlock(nn.Module):
+    def __init__(self, outer_nc, inner_nc, input_nc=None,
+                 submodule=None, outermost=False, innermost=False, norm_layer=nn.BatchNorm2d, use_dropout=False):
+        super(UnetSkipConnectionBlock, self).__init__()
+        self.outermost = outermost
+        if type(norm_layer) == functools.partial:
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+        if input_nc is None:
+            input_nc = outer_nc
+        downconv = nn.Conv2d(input_nc, inner_nc, kernel_size=4,
+                             stride=2, padding=1, bias=use_bias)
+        downrelu = nn.LeakyReLU(0.2, True)
+        downnorm = norm_layer(inner_nc)
+        uprelu = nn.ReLU(True)
+        upnorm = norm_layer(outer_nc)
+
+        if outermost:
+            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
+                                        kernel_size=4, stride=2,
+                                        padding=1)
+            down = [downconv]
+            up = [uprelu, upconv, nn.Tanh()]
+            model = down + [submodule] + up
+        elif innermost:
+            upconv = nn.ConvTranspose2d(inner_nc, outer_nc,
+                                        kernel_size=4, stride=2,
+                                        padding=1, bias=use_bias)
+            down = [downrelu, downconv]
+            up = [uprelu, upconv, upnorm]
+            model = down + up
+        else:
+            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
+                                        kernel_size=4, stride=2,
+                                        padding=1, bias=use_bias)
+            down = [downrelu, downconv, downnorm]
+            up = [uprelu, upconv, upnorm]
+
+            if use_dropout:
+                model = down + [submodule] + up + [nn.Dropout(0.5)]
+            else:
+                model = down + [submodule] + up
+
+        self.model = nn.Sequential(*model)
+
+    def forward(self, x):
+        if self.outermost:
+            return self.model(x)
+        else:
+            return torch.cat([x, self.model(x)], 1)
+
+class UnetSkipConnectionResBlock(nn.Module):
+    def __init__(self, outer_nc, inner_nc, input_nc=None,
+                 submodule=None, outermost=False, innermost=False, norm_layer=nn.BatchNorm2d, use_dropout=False, nres=1):
+        super(UnetSkipConnectionResBlock, self).__init__()
+        self.outermost = outermost
+        if type(norm_layer) == functools.partial:
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+        if input_nc is None:
+            input_nc = outer_nc
+        downconv = nn.Conv2d(input_nc, inner_nc, kernel_size=4,
+                             stride=2, padding=1, bias=use_bias)
+        downrelu = nn.LeakyReLU(0.2, True)
+        downnorm = norm_layer(inner_nc)
+        uprelu = nn.ReLU(True)
+        upnorm = norm_layer(outer_nc)
+
+        if outermost:
+            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
+                                        kernel_size=4, stride=2,
+                                        padding=1)
+            down = [downconv]
+            up = [uprelu, upconv, nn.Tanh()]
+            model = down + [submodule] + up
+        elif innermost:
+            upconv = nn.ConvTranspose2d(inner_nc, outer_nc,
+                                        kernel_size=4, stride=2,
+                                        padding=1, bias=use_bias)
+            down = [downrelu, downconv, downrelu]
+            up = [upconv, upnorm]
+            model = down
+            # resblock: conv norm relu conv norm +
+            for i in range(nres):
+                model += [ResnetBlock(inner_nc, padding_type='reflect', norm_layer=norm_layer, use_dropout=use_dropout, use_bias=use_bias)]
+            model += up
+            #model = down + [submodule] + up
+            print('UnetSkipConnectionResBlock','nres',nres,'inner_nc',inner_nc)
+        else:
+            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
+                                        kernel_size=4, stride=2,
+                                        padding=1, bias=use_bias)
+            down = [downrelu, downconv, downnorm]
+            up = [uprelu, upconv, upnorm]
+
+            if use_dropout:
+                model = down + [submodule] + up + [nn.Dropout(0.5)]
+            else:
+                model = down + [submodule] + up
+
+        self.model = nn.Sequential(*model)
+
+    def forward(self, x):
+        if self.outermost:
+            return self.model(x)
+        else:
+            return torch.cat([x, self.model(x)], 1)
+
+class UnetSkipConnectionStyleBlock(nn.Module):
+    def __init__(self, outer_nc, inner_nc, input_nc=None,
+                 submodule=None, outermost=False, innermost=False,
+                 extra_channel=2, norm_layer=nn.BatchNorm2d, use_dropout=False):
+        super(UnetSkipConnectionStyleBlock, self).__init__()
+        self.outermost = outermost
+        self.innermost = innermost
+        self.extra_channel = extra_channel
+        if type(norm_layer) == functools.partial:
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+        if input_nc is None:
+            input_nc = outer_nc
+        downconv = nn.Conv2d(input_nc, inner_nc, kernel_size=4,
+                             stride=2, padding=1, bias=use_bias)
+        downrelu = nn.LeakyReLU(0.2, True)
+        downnorm = norm_layer(inner_nc)
+        uprelu = nn.ReLU(True)
+        upnorm = norm_layer(outer_nc)
+
+        if outermost:
+            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
+                                        kernel_size=4, stride=2,
+                                        padding=1)
+            down = [downconv]
+            up = [uprelu, upconv, nn.Tanh()]
+            model = down + [submodule] + up
+        elif innermost:
+            upconv = nn.ConvTranspose2d(inner_nc+extra_channel, outer_nc,
+                                        kernel_size=4, stride=2,
+                                        padding=1, bias=use_bias)
+            down = [downrelu, downconv]
+            up = [uprelu, upconv, upnorm]
+            model = down + up
+        else:
+            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
+                                        kernel_size=4, stride=2,
+                                        padding=1, bias=use_bias)
+            down = [downrelu, downconv, downnorm]
+            up = [uprelu, upconv, upnorm]
+
+            if use_dropout:
+                up = up + [nn.Dropout(0.5)]
+            model = down + [submodule] + up
+
+        self.model = nn.Sequential(*model)
+
+        self.downmodel = nn.Sequential(*down)
+        self.upmodel = nn.Sequential(*up)
+        self.submodule = submodule
+
+    def forward(self, x, cate):# cate [bs,c]
+        if self.innermost:
+            y1 = self.downmodel(x)
+            [bs,c,h,w] = y1.shape
+            map = cate.repeat(h,w,1,1).permute([2,3,0,1])
+            y2 = torch.cat([y1,map], 1)
+            y3 = self.upmodel(y2)
+            return torch.cat([x, y3], 1)
+        else:
+            y1 = self.downmodel(x)
+            y2 = self.submodule(y1,cate)
+            y3 = self.upmodel(y2)
+            if self.outermost:
+                return y3
+            else:
+                return torch.cat([x, y3], 1)
+
+class UnetSkipConnectionResStyleBlock(nn.Module):
+    def __init__(self, outer_nc, inner_nc, input_nc=None,
+                 submodule=None, outermost=False, innermost=False,
+                 extra_channel=2, norm_layer=nn.BatchNorm2d, use_dropout=False, nres=1):
+        super(UnetSkipConnectionResStyleBlock, self).__init__()
+        self.outermost = outermost
+        self.innermost = innermost
+        self.extra_channel = extra_channel
+        if type(norm_layer) == functools.partial:
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+        if input_nc is None:
+            input_nc = outer_nc
+        downconv = nn.Conv2d(input_nc, inner_nc, kernel_size=4,
+                             stride=2, padding=1, bias=use_bias)
+        downrelu = nn.LeakyReLU(0.2, True)
+        downnorm = norm_layer(inner_nc)
+        uprelu = nn.ReLU(True)
+        upnorm = norm_layer(outer_nc)
+
+        if outermost:
+            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
+                                        kernel_size=4, stride=2,
+                                        padding=1)
+            down = [downconv]
+            up = [uprelu, upconv, nn.Tanh()]
+            model = down + [submodule] + up
+        elif innermost:
+            upconv = nn.ConvTranspose2d(inner_nc, outer_nc,
+                                        kernel_size=4, stride=2,
+                                        padding=1, bias=use_bias)
+            down = [downrelu, downconv, downrelu]
+            up = [nn.Conv2d(inner_nc+extra_channel, inner_nc, kernel_size=3, stride=1, padding=1, bias=use_bias),
+                      norm_layer(inner_nc),
+                      nn.ReLU(True)]
+            for i in range(nres):
+                up += [ResnetBlock(inner_nc, padding_type='reflect', norm_layer=norm_layer, use_dropout=use_dropout, use_bias=use_bias)]
+            up += [ upconv, upnorm]
+            model = down + up
+            print('UnetSkipConnectionResStyleBlock','nres',nres,'inner_nc',inner_nc)
+        else:
+            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
+                                        kernel_size=4, stride=2,
+                                        padding=1, bias=use_bias)
+            down = [downrelu, downconv, downnorm]
+            up = [uprelu, upconv, upnorm]
+
+            if use_dropout:
+                up = up + [nn.Dropout(0.5)]
+            model = down + [submodule] + up
+
+        self.model = nn.Sequential(*model)
+
+        self.downmodel = nn.Sequential(*down)
+        self.upmodel = nn.Sequential(*up)
+        self.submodule = submodule
+
+    def forward(self, x, cate):# cate [bs,c]
+        # concate in the innermost block
+        if self.innermost:
+            y1 = self.downmodel(x)
+            [bs,c,h,w] = y1.shape
+            map = cate.repeat(h,w,1,1).permute([2,3,0,1])
+            y2 = torch.cat([y1,map], 1)
+            y3 = self.upmodel(y2)
+            return torch.cat([x, y3], 1)
+        else:
+            y1 = self.downmodel(x)
+            y2 = self.submodule(y1,cate)
+            y3 = self.upmodel(y2)
+            if self.outermost:
+                return y3
+            else:
+                return torch.cat([x, y3], 1)
+
+# Defines the PatchGAN discriminator with the specified arguments.
+class NLayerDiscriminator(nn.Module):
+    def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d, use_sigmoid=False):
+        super(NLayerDiscriminator, self).__init__()
+        if type(norm_layer) == functools.partial:
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+
+        kw = 4
+        padw = 1
+        sequence = [
+            nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw),
+            nn.LeakyReLU(0.2, True)
+        ]
+
+        nf_mult = 1
+        nf_mult_prev = 1
+        for n in range(1, n_layers):
+            nf_mult_prev = nf_mult
+            nf_mult = min(2**n, 8)
+            sequence += [
+                nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult,
+                          kernel_size=kw, stride=2, padding=padw, bias=use_bias),
+                norm_layer(ndf * nf_mult),
+                nn.LeakyReLU(0.2, True)
+            ]
+
+        nf_mult_prev = nf_mult
+        nf_mult = min(2**n_layers, 8)
+        sequence += [
+            nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult,
+                      kernel_size=kw, stride=1, padding=padw, bias=use_bias),
+            norm_layer(ndf * nf_mult),
+            nn.LeakyReLU(0.2, True)
+        ]
+
+        sequence += [nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw)]
+
+        if use_sigmoid:#no_lsgan, use sigmoid before calculating bceloss(binary cross entropy)
+            sequence += [nn.Sigmoid()]
+
+        self.model = nn.Sequential(*sequence)
+
+    def forward(self, input):
+        return self.model(input)
+
+
+class PixelDiscriminator(nn.Module):
+    def __init__(self, input_nc, ndf=64, norm_layer=nn.BatchNorm2d, use_sigmoid=False):
+        super(PixelDiscriminator, self).__init__()
+        if type(norm_layer) == functools.partial:
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+
+        self.net = [
+            nn.Conv2d(input_nc, ndf, kernel_size=1, stride=1, padding=0),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(ndf, ndf * 2, kernel_size=1, stride=1, padding=0, bias=use_bias),
+            norm_layer(ndf * 2),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(ndf * 2, 1, kernel_size=1, stride=1, padding=0, bias=use_bias)]
+
+        if use_sigmoid:
+            self.net.append(nn.Sigmoid())
+
+        self.net = nn.Sequential(*self.net)
+
+    def forward(self, input):
+        return self.net(input)

APDrawingGAN2/models/test_model.py

@@ -0,0 +1,214 @@
+from .base_model import BaseModel
+from . import networks
+import torch
+
+
+class TestModel(BaseModel):
+    def name(self):
+        return 'TestModel'
+
+    @staticmethod
+    def modify_commandline_options(parser, is_train=True):
+        assert not is_train, 'TestModel cannot be used in train mode'
+        # uncomment because default CycleGAN did not use dropout ( parser.set_defaults(no_dropout=True) )
+        # parser = CycleGANModel.modify_commandline_options(parser, is_train=False)   
+        parser.set_defaults(pool_size=0, no_lsgan=True, norm='batch')# no_lsgan=True, use_lsgan=False
+        parser.set_defaults(dataset_mode='single')
+        parser.set_defaults(auxiliary_root='auxiliaryeye2o')
+        parser.set_defaults(use_local=True, hair_local=True, bg_local=True)
+        parser.set_defaults(nose_ae=True, others_ae=True, compactmask=True, MOUTH_H=56)
+        parser.set_defaults(soft_border=1)
+        parser.add_argument('--nnG_hairc', type=int, default=6, help='nnG for hair classifier')
+        parser.add_argument('--use_resnet', action='store_true', help='use resnet for generator')
+
+        parser.add_argument('--model_suffix', type=str, default='',
+                            help='In checkpoints_dir, [which_epoch]_net_G[model_suffix].pth will'
+                            ' be loaded as the generator of TestModel')
+
+        return parser
+
+    def initialize(self, opt):
+        assert(not opt.isTrain)
+        BaseModel.initialize(self, opt)
+
+        # specify the training losses you want to print out. The program will call base_model.get_current_losses
+        self.loss_names = []
+        # specify the images you want to save/display. The program will call base_model.get_current_visuals
+        self.visual_names = ['real_A', 'fake_B']
+        # specify the models you want to save to the disk. The program will call base_model.save_networks and base_model.load_networks
+        self.model_names = ['G' + opt.model_suffix]
+        self.auxiliary_model_names = []
+        if self.opt.use_local:
+            self.model_names += ['GLEyel','GLEyer','GLNose','GLMouth','GLHair','GLBG','GCombine']
+            self.auxiliary_model_names += ['CLm','CLh']
+            # auxiliary nets for local output refinement
+            if self.opt.nose_ae:
+                self.auxiliary_model_names += ['AE']
+            if self.opt.others_ae:
+                self.auxiliary_model_names += ['AEel','AEer','AEmowhite','AEmoblack']
+        print('model_names', self.model_names)
+        print('auxiliary_model_names', self.auxiliary_model_names)
+
+        # load/define networks
+        self.netG = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG, opt.norm,
+                                      not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids,
+                                      opt.nnG)
+        print('netG', opt.netG)
+        if self.opt.use_local:
+            netlocal1 = 'partunet' if self.opt.use_resnet == 0 else 'resnet_nblocks'
+            netlocal2 = 'partunet2' if self.opt.use_resnet == 0 else 'resnet_6blocks'
+            netlocal2_style = 'partunet2style' if self.opt.use_resnet == 0 else 'resnet_style2_6blocks'
+            self.netGLEyel = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, netlocal1, opt.norm,
+                                      not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids, nnG=3)
+            self.netGLEyer = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, netlocal1, opt.norm,
+                                      not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids, nnG=3)
+            self.netGLNose = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, netlocal1, opt.norm,
+                                      not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids, nnG=3)
+            self.netGLMouth = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, netlocal1, opt.norm,
+                                        not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids, nnG=3)
+            self.netGLHair = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, netlocal2_style, opt.norm,
+                                      not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids, nnG=4,
+                                      extra_channel=3)
+            self.netGLBG = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, netlocal2, opt.norm,
+                                    not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids, nnG=4)
+            # by default combiner_type is combiner, which uses resnet
+            print('combiner_type', self.opt.combiner_type)
+            self.netGCombine = networks.define_G(2*opt.output_nc, opt.output_nc, opt.ngf, self.opt.combiner_type, opt.norm,
+                                      not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids, 2)
+            # auxiliary classifiers for mouth and hair
+            ratio = self.opt.fineSize / 256
+            self.MOUTH_H = int(self.opt.MOUTH_H * ratio)
+            self.MOUTH_W = int(self.opt.MOUTH_W * ratio)
+            self.netCLm = networks.define_G(opt.input_nc, 2, opt.ngf, 'classifier', opt.norm,
+                                      not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids,
+                                      nnG = 3, ae_h = self.MOUTH_H, ae_w = self.MOUTH_W)
+            self.netCLh = networks.define_G(opt.input_nc, 3, opt.ngf, 'classifier', opt.norm,
+                                      not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids,
+                                      nnG = opt.nnG_hairc, ae_h = opt.fineSize, ae_w = opt.fineSize)
+        # ==================================auxiliary nets (loaded, parameters fixed)=============================
+        if self.opt.use_local and self.opt.nose_ae:
+            ratio = self.opt.fineSize / 256
+            NOSE_H = self.opt.NOSE_H * ratio
+            NOSE_W = self.opt.NOSE_W * ratio
+            self.netAE = networks.define_G(opt.output_nc, opt.output_nc, opt.ngf, self.opt.nose_ae_net, 'batch',
+                                      not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids,
+                                       latent_dim=self.opt.ae_latentno, ae_h=NOSE_H, ae_w=NOSE_W)
+            self.set_requires_grad(self.netAE, False)            
+        if self.opt.use_local and self.opt.others_ae:
+            ratio = self.opt.fineSize / 256
+            EYE_H = self.opt.EYE_H * ratio
+            EYE_W = self.opt.EYE_W * ratio
+            MOUTH_H = self.opt.MOUTH_H * ratio
+            MOUTH_W = self.opt.MOUTH_W * ratio
+            self.netAEel = networks.define_G(opt.output_nc, opt.output_nc, opt.ngf, self.opt.nose_ae_net, 'batch',
+                                      not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids,
+                                      latent_dim=self.opt.ae_latenteye, ae_h=EYE_H, ae_w=EYE_W)
+            self.netAEer = networks.define_G(opt.output_nc, opt.output_nc, opt.ngf, self.opt.nose_ae_net, 'batch',
+                                      not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids,
+                                      latent_dim=self.opt.ae_latenteye, ae_h=EYE_H, ae_w=EYE_W)
+            self.netAEmowhite = networks.define_G(opt.output_nc, opt.output_nc, opt.ngf, self.opt.nose_ae_net, 'batch',
+                                      not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids,
+                                      latent_dim=self.opt.ae_latentmo, ae_h=MOUTH_H, ae_w=MOUTH_W)
+            self.netAEmoblack = networks.define_G(opt.output_nc, opt.output_nc, opt.ngf, self.opt.nose_ae_net, 'batch',
+                                      not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids,
+                                      latent_dim=self.opt.ae_latentmo, ae_h=MOUTH_H, ae_w=MOUTH_W)
+            self.set_requires_grad(self.netAEel, False)
+            self.set_requires_grad(self.netAEer, False)
+            self.set_requires_grad(self.netAEmowhite, False)
+            self.set_requires_grad(self.netAEmoblack, False)
+
+        # assigns the model to self.netG_[suffix] so that it can be loaded
+        # please see BaseModel.load_networks
+        setattr(self, 'netG' + opt.model_suffix, self.netG)
+
+    def set_input(self, input):
+        # we need to use single_dataset mode
+        self.real_A = input['A'].to(self.device)
+        self.image_paths = input['A_paths']
+        self.batch_size = len(self.image_paths)
+        if self.opt.use_local:
+            self.real_A_eyel = input['eyel_A'].to(self.device)
+            self.real_A_eyer = input['eyer_A'].to(self.device)
+            self.real_A_nose = input['nose_A'].to(self.device)
+            self.real_A_mouth = input['mouth_A'].to(self.device)
+            self.center = input['center']
+            if self.opt.soft_border:
+                self.softel = input['soft_eyel_mask'].to(self.device)
+                self.softer = input['soft_eyer_mask'].to(self.device)
+                self.softno = input['soft_nose_mask'].to(self.device)
+                self.softmo = input['soft_mouth_mask'].to(self.device)
+            if self.opt.compactmask:
+                self.cmask = input['cmask'].to(self.device)
+                self.cmask1 = self.cmask*2-1#[0,1]->[-1,1]
+                self.cmaskel = input['cmaskel'].to(self.device)
+                self.cmask1el = self.cmaskel*2-1
+                self.cmasker = input['cmasker'].to(self.device)
+                self.cmask1er = self.cmasker*2-1
+                self.cmaskmo = input['cmaskmo'].to(self.device)
+                self.cmask1mo = self.cmaskmo*2-1
+            self.real_A_hair = input['hair_A'].to(self.device)
+            self.mask = input['mask'].to(self.device) # mask for non-eyes,nose,mouth
+            self.mask2 = input['mask2'].to(self.device) # mask for non-bg
+            self.real_A_bg = input['bg_A'].to(self.device)
+
+    def getonehot(self,outputs,classes):
+        [maxv,index] = torch.max(outputs,1)
+        y = torch.unsqueeze(index,1)
+        onehot = torch.FloatTensor(self.batch_size,classes).to(self.device)
+        onehot.zero_()
+        onehot.scatter_(1,y,1)
+        return onehot  
+
+    def forward(self):
+        if not self.opt.use_local:
+            self.fake_B = self.netG(self.real_A)
+        else:
+            self.fake_B0 = self.netG(self.real_A)
+            # EYES, MOUTH
+            outputs1 = self.netCLm(self.real_A_mouth)
+            onehot1 = self.getonehot(outputs1,2)
+
+            if not self.opt.others_ae:
+                fake_B_eyel = self.netGLEyel(self.real_A_eyel)
+                fake_B_eyer = self.netGLEyer(self.real_A_eyer)
+                fake_B_mouth = self.netGLMouth(self.real_A_mouth)
+            else: # use AE that only constains compact region, need cmask!
+                self.fake_B_eyel1 = self.netGLEyel(self.real_A_eyel)
+                self.fake_B_eyer1 = self.netGLEyer(self.real_A_eyer)
+                self.fake_B_mouth1 = self.netGLMouth(self.real_A_mouth)
+                self.fake_B_eyel2,_ = self.netAEel(self.fake_B_eyel1)
+                self.fake_B_eyer2,_ = self.netAEer(self.fake_B_eyer1)
+                # USE 2 AEs
+                self.fake_B_mouth2 = torch.FloatTensor(self.batch_size,self.opt.output_nc,self.MOUTH_H,self.MOUTH_W).to(self.device)
+                for i in range(self.batch_size):
+                    if onehot1[i][0] == 1:
+                        self.fake_B_mouth2[i],_ = self.netAEmowhite(self.fake_B_mouth1[i].unsqueeze(0))
+                        #print('AEmowhite')
+                    elif onehot1[i][1] == 1:
+                        self.fake_B_mouth2[i],_ = self.netAEmoblack(self.fake_B_mouth1[i].unsqueeze(0))
+                        #print('AEmoblack')
+                fake_B_eyel = self.add_with_mask(self.fake_B_eyel2,self.fake_B_eyel1,self.cmaskel)
+                fake_B_eyer = self.add_with_mask(self.fake_B_eyer2,self.fake_B_eyer1,self.cmasker)
+                fake_B_mouth = self.add_with_mask(self.fake_B_mouth2,self.fake_B_mouth1,self.cmaskmo)
+            # NOSE
+            if not self.opt.nose_ae:
+                fake_B_nose = self.netGLNose(self.real_A_nose)
+            else: # use AE that only constains compact region, need cmask!
+                self.fake_B_nose1 = self.netGLNose(self.real_A_nose)
+                self.fake_B_nose2,_ = self.netAE(self.fake_B_nose1)
+                fake_B_nose = self.add_with_mask(self.fake_B_nose2,self.fake_B_nose1,self.cmask)
+            
+            # HAIR, BG AND PARTCOMBINE
+            outputs2 = self.netCLh(self.real_A_hair)
+            onehot2 = self.getonehot(outputs2,3)
+
+            fake_B_hair = self.netGLHair(self.real_A_hair,onehot2)
+            fake_B_bg = self.netGLBG(self.real_A_bg)
+            self.fake_B_hair = self.masked(fake_B_hair,self.mask*self.mask2)
+            self.fake_B_bg = self.masked(fake_B_bg,self.inverse_mask(self.mask2))
+            if not self.opt.compactmask:
+                self.fake_B1 = self.partCombiner2_bg(fake_B_eyel,fake_B_eyer,fake_B_nose,fake_B_mouth,fake_B_hair,fake_B_bg,self.mask*self.mask2,self.inverse_mask(self.mask2),self.opt.comb_op)
+            else:
+                self.fake_B1 = self.partCombiner2_bg(fake_B_eyel,fake_B_eyer,fake_B_nose,fake_B_mouth,fake_B_hair,fake_B_bg,self.mask*self.mask2,self.inverse_mask(self.mask2),self.opt.comb_op,self.opt.region_enm,self.cmaskel,self.cmasker,self.cmask,self.cmaskmo)
+            
+            self.fake_B = self.netGCombine(torch.cat([self.fake_B0,self.fake_B1],1))

APDrawingGAN2/options/base_options.py

@@ -0,0 +1,192 @@
+import argparse
+import os
+from util import util
+import torch
+import models
+import data
+
+
+class BaseOptions():
+    def __init__(self):
+        self.initialized = False
+
+    def initialize(self, parser):
+        parser.add_argument('--dataroot', type=str, default='', help='path to images (should have subfolders trainA, trainB, valA, valB, etc)')
+        parser.add_argument('--batch_size', type=int, default=1, help='input batch size')
+        parser.add_argument('--loadSize', type=int, default=512, help='scale images to this size')
+        parser.add_argument('--fineSize', type=int, default=512, help='then crop to this size')
+        parser.add_argument('--input_nc', type=int, default=3, help='# of input image channels')
+        parser.add_argument('--output_nc', type=int, default=1, help='# of output image channels')
+        parser.add_argument('--ngf', type=int, default=64, help='# of gen filters in first conv layer')
+        parser.add_argument('--ndf', type=int, default=64, help='# of discrim filters in first conv layer')
+        parser.add_argument('--netD', type=str, default='basic', help='selects model to use for netD')
+        parser.add_argument('--netG', type=str, default='unet_256', help='selects model to use for netG')
+        parser.add_argument('--nnG', type=int, default=9, help='specify nblock for resnet_nblocks, ndown for unet for unet_ndown')
+        parser.add_argument('--n_layers_D', type=int, default=3, help='only used if netD==n_layers')
+        parser.add_argument('--gpu_ids', type=str, default='0', help='gpu ids: e.g. 0  0,1,2, 0,2. use -1 for CPU')
+        parser.add_argument('--gpu_ids_p', type=str, default='0', help='gpu ids for pretrained auxiliary models: e.g. 0  0,1,2, 0,2. use -1 for CPU')
+        parser.add_argument('--name', type=str, default='experiment_name', help='name of the experiment. It decides where to store samples and models')
+        parser.add_argument('--dataset_mode', type=str, default='aligned', help='chooses how datasets are loaded. [unaligned | aligned | single]')
+        parser.add_argument('--model', type=str, default='apdrawing',
+                            help='chooses which model to use. cycle_gan, pix2pix, test, autoencoder')
+        parser.add_argument('--use_local', action='store_true', help='use local part network')
+        parser.add_argument('--lm_dir', type=str, default='dataset/landmark/', help='path to facial landmarks')
+        parser.add_argument('--nose_ae', action='store_true', help='use nose autoencoder')
+        parser.add_argument('--others_ae', action='store_true', help='use autoencoder for eyes and mouth too')
+        parser.add_argument('--nose_ae_net', type=str, default='autoencoderfc', help='net for nose autoencoder [autoencoder | autoencoderfc]')
+        parser.add_argument('--comb_op', type=int, default=1, help='use min-pooling(1) or max-pooling(0) for overlapping regions')
+        parser.add_argument('--hair_local', action='store_true', help='add hair part')
+        parser.add_argument('--bg_local', action='store_true', help='use background mask to seperate background')
+        parser.add_argument('--bg_dir', default='dataset/mask/bg/', type=str, help='choose bg_dir')
+        parser.add_argument('--region_enm', type=int, default=0, help='region type for eyes nose mouth: 0 for rectangle, 1 for campact mask in rectangle, 2 for mask no rectangle (1,2 must have compactmask, 0 use compactmask for AE)')
+        parser.add_argument('--soft_border', type=int, default=0, help='use mask with soft border')
+        parser.add_argument('--EYE_H', type=int, default=40, help='EYE_H')
+        parser.add_argument('--EYE_W', type=int, default=56, help='EYE_W')
+        parser.add_argument('--NOSE_H', type=int, default=48, help='NOSE_H')
+        parser.add_argument('--NOSE_W', type=int, default=48, help='NOSE_W')
+        parser.add_argument('--MOUTH_H', type=int, default=40, help='MOUTH_H')
+        parser.add_argument('--MOUTH_W', type=int, default=64, help='MOUTH_W')
+        parser.add_argument('--average_pos', action='store_true', help='use avg pos in partCombiner')
+        parser.add_argument('--combiner_type', type=str, default='combiner', help='choose combiner type')
+        parser.add_argument('--which_direction', type=str, default='AtoB', help='AtoB or BtoA')
+        parser.add_argument('--num_threads', default=4, type=int, help='# threads for loading data')
+        parser.add_argument('--checkpoints_dir', type=str, default='./checkpoints', help='models are saved here')
+        parser.add_argument('--auxiliary_root', type=str, default='auxiliary', help='auxiliary model folder')
+        parser.add_argument('--norm', type=str, default='instance', help='instance normalization or batch normalization')
+        parser.add_argument('--serial_batches', action='store_true', help='if true, takes images in order to make batches, otherwise takes them randomly')
+        parser.add_argument('--display_winsize', type=int, default=256, help='display window size')
+        parser.add_argument('--display_id', type=int, default=1, help='window id of the web display')
+        parser.add_argument('--display_server', type=str, default="http://localhost", help='visdom server of the web display')
+        parser.add_argument('--display_env', type=str, default='main', help='visdom display environment name (default is "main")')
+        parser.add_argument('--display_port', type=int, default=8097, help='visdom port of the web display')
+        parser.add_argument('--no_dropout', action='store_true', help='no dropout for the generator')
+        parser.add_argument('--max_dataset_size', type=int, default=float("inf"), help='Maximum number of samples allowed per dataset. If the dataset directory contains more than max_dataset_size, only a subset is loaded.')
+        parser.add_argument('--resize_or_crop', type=str, default='resize_and_crop', help='scaling and cropping of images at load time [resize_and_crop|crop|scale_width|scale_width_and_crop]')
+        parser.add_argument('--no_flip', action='store_true', help='if specified, do not flip the images for data augmentation')
+        parser.add_argument('--init_type', type=str, default='normal', help='network initialization [normal|xavier|kaiming|orthogonal]')
+        parser.add_argument('--init_gain', type=float, default=0.02, help='scaling factor for normal, xavier and orthogonal.')
+        parser.add_argument('--verbose', action='store_true', help='if specified, print more debugging information')
+        parser.add_argument('--suffix', default='', type=str, help='customized suffix: opt.name = opt.name + suffix: e.g., {model}_{netG}_size{loadSize}')
+        # compact mask
+        parser.add_argument('--compactmask', action='store_true', help='use compact mask as input and apply to loss')# "when you calculate the (ae) loss, you should also restrict to nose pixels"
+        parser.add_argument('--cmask_dir', type=str, default='dataset/mask/', help='compact mask directory')
+        parser.add_argument('--ae_latentno', type=int, default=1024 ,help='latent space dim for pretrained NOSE AEwithfc')
+        parser.add_argument('--ae_latentmo', type=int, default=1024 ,help='latent space dim for pretrained MOUTH AEwithfc')
+        parser.add_argument('--ae_latenteye', type=int, default=1024 ,help='latent space dim for pretrained EYEL/EYER AEwithfc')
+        parser.add_argument('--ae_small', type=int, default=0 ,help='use latent dim smaller than default 1024 in 4 AEs')
+        # below for autoencoder
+        parser.add_argument('--ae_latent', type=int, default=1024 ,help='latent space dim for autoencoderfc')
+        parser.add_argument('--ae_multiple', type=float, default=2 ,help='filter number change in ae encoder')
+        parser.add_argument('--ae_h', type=int, default=96 ,help='ae input h')
+        parser.add_argument('--ae_w', type=int, default=96 ,help='ae input w')
+        parser.add_argument('--ae_region', type=str, default='nose' ,help='autoencoder for which region')
+        parser.add_argument('--no_ae', action='store_true', help='no ae')        
+        self.initialized = True
+        return parser
+
+    def gather_options(self):
+        # initialize parser with basic options
+        if not self.initialized:
+            parser = argparse.ArgumentParser(
+                formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+            parser = self.initialize(parser)
+
+        # get the basic options
+        opt, _ = parser.parse_known_args()
+
+        # modify model-related parser options
+        model_name = opt.model
+        model_option_setter = models.get_option_setter(model_name)
+        parser = model_option_setter(parser, self.isTrain)
+        opt, _ = parser.parse_known_args()  # parse again with the new defaults
+
+        # modify dataset-related parser options
+        dataset_name = opt.dataset_mode
+        dataset_option_setter = data.get_option_setter(dataset_name)
+        parser = dataset_option_setter(parser, self.isTrain)
+
+        self.parser = parser
+
+        return parser.parse_args()
+
+    def print_options(self, opt):
+        message = ''
+        message += '----------------- Options ---------------\n'
+        for k, v in sorted(vars(opt).items()):
+            comment = ''
+            default = self.parser.get_default(k)
+            if v != default:
+                comment = '\t[default: %s]' % str(default)
+            message += '{:>25}: {:<30}{}\n'.format(str(k), str(v), comment)
+        message += '----------------- End -------------------'
+        print(message)
+
+        # save to the disk
+        expr_dir = os.path.join(opt.checkpoints_dir, opt.name)
+        util.mkdirs(expr_dir)
+        file_name = os.path.join(expr_dir, '{}_opt.txt'.format(opt.phase))
+        with open(file_name, 'wt') as opt_file:
+            opt_file.write(message)
+            opt_file.write('\n')
+
+    def parse(self, print=True):
+
+        opt = self.gather_options()
+        if opt.use_local:
+            opt.loadSize = opt.fineSize
+        if opt.region_enm in [1,2]:
+            opt.compactmask = True
+        if opt.nose_ae or opt.others_ae:
+            opt.compactmask = True
+        if opt.ae_latentno < 1024 and opt.ae_latentmo < 1024 and opt.ae_latenteye < 1024:
+            opt.ae_small = 1
+        opt.isTrain = self.isTrain   # train or test
+
+        # process opt.suffix
+        if opt.suffix:
+            suffix = ('_' + opt.suffix.format(**vars(opt))) if opt.suffix != '' else ''
+            opt.name = opt.name + suffix
+        
+        if self.isTrain and opt.pretrain:
+            opt.nose_ae = False
+            opt.others_ae = False
+            opt.compactmask = False
+            opt.chamfer_loss = False
+        if not self.isTrain and opt.pretrain:
+            opt.nose_ae = False
+            opt.others_ae = False
+            opt.compactmask = False
+        if opt.no_ae:
+            opt.nose_ae = False
+            opt.others_ae = False
+            opt.compactmask = False
+        if self.isTrain and opt.no_dtremap:
+            opt.dt_nonlinear = ''
+            opt.lambda_chamfer = 0.1
+            opt.lambda_chamfer2 = 0.1
+        if self.isTrain and opt.no_dt:
+            opt.chamfer_loss = False
+
+        if print:
+            self.print_options(opt)
+
+        # set gpu ids
+        str_ids = opt.gpu_ids.split(',')
+        opt.gpu_ids = []
+        for str_id in str_ids:
+            id = int(str_id)
+            if id >= 0:
+                opt.gpu_ids.append(id)
+        if len(opt.gpu_ids) > 0:
+            torch.cuda.set_device(opt.gpu_ids[0])
+        
+        # set gpu ids
+        str_ids = opt.gpu_ids_p.split(',')
+        opt.gpu_ids_p = []
+        for str_id in str_ids:
+            id = int(str_id)
+            if id >= 0:
+                opt.gpu_ids_p.append(id)
+
+        self.opt = opt
+        return self.opt

APDrawingGAN2/options/test_options.py

@@ -0,0 +1,23 @@
+from .base_options import BaseOptions
+
+
+class TestOptions(BaseOptions):
+    def initialize(self, parser):
+        parser = BaseOptions.initialize(self, parser)
+        parser.add_argument('--ntest', type=int, default=float("inf"), help='# of test examples.')
+        parser.add_argument('--results_dir', type=str, default='./results/', help='saves results here.')
+        parser.add_argument('--aspect_ratio', type=float, default=1.0, help='aspect ratio of result images')
+        parser.add_argument('--phase', type=str, default='test', help='train, val, test, etc')
+        parser.add_argument('--which_epoch', type=str, default='latest', help='which epoch to load? set to latest to use latest cached model')
+        parser.add_argument('--how_many', type=int, default=50, help='how many test images to run')
+        parser.add_argument('--test_continuity_loss', action='store_true', help='get continuity value in test')
+        parser.add_argument('--netG_line', type=str, default='unet_512', help='selects model to use for netG_line')
+        parser.add_argument('--save2', action='store_true', help='only save real_A and fake_B')
+        parser.add_argument('--imagefolder', type=str, default='images', help='subfolder to save images')
+        parser.add_argument('--pretrain', action='store_true', help='pretrain stage, no dt loss, no ae')        
+
+        parser.set_defaults(model='test')
+        # To avoid cropping, the loadSize should be the same as fineSize
+        parser.set_defaults(loadSize=parser.get_default('fineSize'))
+        self.isTrain = False
+        return parser

APDrawingGAN2/options/train_options.py

@@ -0,0 +1,62 @@
+from .base_options import BaseOptions
+
+
+class TrainOptions(BaseOptions):
+    def initialize(self, parser):
+        parser = BaseOptions.initialize(self, parser)
+        parser.add_argument('--display_freq', type=int, default=400, help='frequency of showing training results on screen')
+        parser.add_argument('--display_ncols', type=int, default=4, help='if positive, display all images in a single visdom web panel with certain number of images per row.')
+        parser.add_argument('--update_html_freq', type=int, default=1000, help='frequency of saving training results to html')
+        parser.add_argument('--print_freq', type=int, default=100, help='frequency of showing training results on console')
+        parser.add_argument('--save_latest_freq', type=int, default=5000, help='frequency of saving the latest results')
+        parser.add_argument('--save_epoch_freq', type=int, default=5, help='frequency of saving checkpoints at the end of epochs')
+        parser.add_argument('--continue_train', action='store_true', help='continue training: load the latest model')
+        parser.add_argument('--epoch_count', type=int, default=1, help='the starting epoch count, we save the model by <epoch_count>, <epoch_count>+<save_latest_freq>, ...')
+        parser.add_argument('--phase', type=str, default='train', help='train, val, test, etc')
+        parser.add_argument('--which_epoch', type=str, default='latest', help='which epoch to load? set to latest to use latest cached model')
+        parser.add_argument('--niter', type=int, default=100, help='# of iter at starting learning rate')
+        parser.add_argument('--niter_decay', type=int, default=100, help='# of iter to linearly decay learning rate to zero')
+        parser.add_argument('--beta1', type=float, default=0.5, help='momentum term of adam')
+        parser.add_argument('--lr', type=float, default=0.0002, help='initial learning rate for adam')
+        parser.add_argument('--no_lsgan', action='store_true', help='do *not* use least square GAN, if false, use vanilla GAN')
+        parser.add_argument('--pool_size', type=int, default=50, help='the size of image buffer that stores previously generated images')
+        parser.add_argument('--no_html', action='store_true', help='do not save intermediate training results to [opt.checkpoints_dir]/[opt.name]/web/')
+        parser.add_argument('--lr_policy', type=str, default='lambda', help='learning rate policy: lambda|step|plateau|cosine')
+        parser.add_argument('--lr_decay_iters', type=int, default=50, help='multiply by a gamma every lr_decay_iters iterations')
+        # ============================================loss=========================================================
+        # chamfer loss
+        parser.add_argument('--chamfer_loss', action='store_true', help='use chamfer loss')
+        parser.add_argument('--chamfer_2way', action='store_true', help='use chamfer loss 2 way')
+        parser.add_argument('--chamfer_only_line', action='store_true', help='use chamfer only on lines')
+        parser.add_argument('--lambda_chamfer', type=float, default=0.1, help='weight for chamfer loss')
+        parser.add_argument('--lambda_chamfer2', type=float, default=0.1, help='weight for chamfer loss2')
+        parser.add_argument('--dt_nonlinear', type=str, default='', help='nonlinear remap on dt [atan | sigmoid | tanh]')
+        parser.add_argument('--dt_xmax', type=float, default=10, help='first mutiply dt to range [0,xmax], then use atan/sigmoid/tanh etc, to have more nonlinearity (not much nonlinearity in range [0,1])')
+        # line continuity loss
+        parser.add_argument('--continuity_loss', action='store_true', help='use line continuity loss')
+        parser.add_argument('--lambda_continuity', type=float, default=10.0, help='weight for continuity loss')
+        parser.add_argument('--emphasis_conti_face', action='store_true', help='constrain conti loss to pixels in original lines (avoid apply to background etc)')
+        parser.add_argument('--facemask_dir', type=str, default='dataset/mask/face/', help='mask folder to constrain conti loss to pixels in original lines')
+        # =====================================auxilary net structure===============================================
+        # dt & line net structure
+        parser.add_argument('--netG_dt', type=str, default='unet_512', help='selects model to use for netG_dt, for chamfer loss')
+        parser.add_argument('--netG_line', type=str, default='unet_512', help='selects model to use for netG_line, for chamfer loss')
+        # multiple discriminators
+        parser.add_argument('--discriminator_local', action='store_true', help='use six diffent local discriminator for 6 local regions')
+        parser.add_argument('--gan_loss_strategy', type=int, default=2, help='specify how to calculate gan loss for g, 1: average global and local discriminators; 2: not change global discriminator weight, 0.25 for local')
+        parser.add_argument('--addw_eye', type=float, default=1.0, help='additional weight for eye region')
+        parser.add_argument('--addw_nose', type=float, default=1.0, help='additional weight for nose region')
+        parser.add_argument('--addw_mouth', type=float, default=1.0, help='additional weight for mouth region')
+        parser.add_argument('--addw_hair', type=float, default=1.0, help='additional weight for hair region')
+        parser.add_argument('--addw_bg', type=float, default=1.0, help='additional weight for bg region')
+        # ==========================================ablation========================================================
+        parser.add_argument('--no_l1_loss', action='store_true', help='no l1 loss')
+        parser.add_argument('--no_G_local_loss', action='store_true', help='not using local transfer loss for local generator output')
+        parser.add_argument('--no_dtremap', action='store_true', help='no dt remap')        
+        parser.add_argument('--no_dt', action='store_true', help='no dt')     
+
+        parser.add_argument('--pretrain', action='store_true', help='pretrain stage, no dt loss, no ae')        
+
+
+        self.isTrain = True
+        return parser

APDrawingGAN2/preprocess/combine_A_and_B.py

@@ -0,0 +1,48 @@
+import os
+import numpy as np
+import cv2
+import argparse
+
+parser = argparse.ArgumentParser('create image pairs')
+parser.add_argument('--fold_A', dest='fold_A', help='input directory for image A', type=str, default='../dataset/50kshoes_edges')
+parser.add_argument('--fold_B', dest='fold_B', help='input directory for image B', type=str, default='../dataset/50kshoes_jpg')
+parser.add_argument('--fold_AB', dest='fold_AB', help='output directory', type=str, default='../dataset/test_AB')
+parser.add_argument('--num_imgs', dest='num_imgs', help='number of images',type=int, default=1000000)
+parser.add_argument('--use_AB', dest='use_AB', help='if true: (0001_A, 0001_B) to (0001_AB)',action='store_true')
+args = parser.parse_args()
+
+for arg in vars(args):
+    print('[%s] = ' % arg,  getattr(args, arg))
+
+splits = os.listdir(args.fold_A)
+
+for sp in splits:
+    img_fold_A = os.path.join(args.fold_A, sp)
+    img_fold_B = os.path.join(args.fold_B, sp)
+    img_list = os.listdir(img_fold_A)
+    if args.use_AB:
+        img_list = [img_path for img_path in img_list if '_A.' in img_path]
+
+    num_imgs = min(args.num_imgs, len(img_list))
+    print('split = %s, use %d/%d images' % (sp, num_imgs, len(img_list)))
+    img_fold_AB = os.path.join(args.fold_AB, sp)
+    if not os.path.isdir(img_fold_AB):
+        os.makedirs(img_fold_AB)
+    print('split = %s, number of images = %d' % (sp, num_imgs))
+    for n in range(num_imgs):
+        name_A = img_list[n]
+        path_A = os.path.join(img_fold_A, name_A)
+        if args.use_AB:
+            name_B = name_A.replace('_A.', '_B.')
+        else:
+            name_B = name_A
+        path_B = os.path.join(img_fold_B, name_B)
+        if os.path.isfile(path_A) and os.path.isfile(path_B):
+            name_AB = name_A
+            if args.use_AB:
+                name_AB = name_AB.replace('_A.', '.') # remove _A
+            path_AB = os.path.join(img_fold_AB, name_AB)
+            im_A = cv2.imread(path_A, cv2.IMREAD_COLOR)
+            im_B = cv2.imread(path_B, cv2.IMREAD_COLOR)
+            im_AB = np.concatenate([im_A, im_B], 1)
+            cv2.imwrite(path_AB, im_AB)

APDrawingGAN2/preprocess/example/img_1701_aligned.txt

@@ -0,0 +1,5 @@
+194 248
+314 249
+261 312
+209 368
+302 371

APDrawingGAN2/preprocess/example/img_1701_aligned_68lm.txt

@@ -0,0 +1,68 @@
+120 261
+124 294
+129 326
+133 358
+142 388
+162 412
+190 430
+220 445
+253 449
+287 447
+317 432
+344 411
+362 385
+370 354
+375 322
+382 291
+385 258
+142 225
+161 209
+188 204
+215 208
+242 218
+269 218
+296 208
+324 206
+351 213
+369 231
+256 244
+256 264
+256 284
+256 305
+232 324
+244 328
+256 332
+267 329
+277 325
+172 252
+186 243
+203 243
+218 253
+203 257
+186 257
+290 254
+305 244
+322 246
+336 255
+322 260
+305 259
+210 368
+229 358
+245 352
+256 354
+267 352
+283 358
+300 368
+284 382
+268 388
+255 389
+244 388
+228 381
+220 368
+245 363
+256 364
+267 364
+290 368
+267 370
+255 372
+244 371

APDrawingGAN2/preprocess/face_align_512.m

@@ -0,0 +1,55 @@
+function [trans_img,trans_facial5point]=face_align_512(impath,facial5point,savedir)
+% align the faces by similarity transformation.
+% using 5 facial landmarks: 2 eyes, nose, 2 mouth corners.
+%   impath: path to image
+%   facial5point: 5x2 size, 5 facial landmark positions, detected by MTCNN
+%   savedir: savedir for cropped image and transformed facial landmarks
+
+%% alignment settings
+imgSize = [512,512];
+coord5point = [180,230;
+    300,230;
+    240,301;
+    186,365.6;
+    294,365.6];%480x480
+coord5point = (coord5point-240)/560 * 512 + 256;
+
+%% face alignment
+
+% load and align, resize image to imgSize
+img      = imread(impath);
+facial5point = double(facial5point);
+transf   = cp2tform(facial5point, coord5point, 'similarity');
+trans_img  = imtransform(img, transf, 'XData', [1 imgSize(2)],...
+                                    'YData', [1 imgSize(1)],...
+                                    'Size', imgSize,...
+                                    'FillValues', [255;255;255]);
+trans_facial5point = round(tformfwd(transf,facial5point));
+
+
+%% save results
+if ~exist(savedir,'dir')
+    mkdir(savedir)
+end
+[~,name,~] = fileparts(impath);
+% save trans_img
+imwrite(trans_img, fullfile(savedir,[name,'_aligned.png']));
+fprintf('write aligned image to %s\n',fullfile(savedir,[name,'_aligned.png']));
+% save trans_facial5point
+write_5pt(fullfile(savedir, [name, '_aligned.txt']), trans_facial5point);
+fprintf('write transformed facial landmark to %s\n',fullfile(savedir,[name,'_aligned.txt']));
+
+%% show results
+imshow(trans_img); hold on;
+plot(trans_facial5point(:,1),trans_facial5point(:,2),'b');
+plot(trans_facial5point(:,1),trans_facial5point(:,2),'r+');
+
+end
+
+function [] = write_5pt(fn, trans_pt)
+fid = fopen(fn, 'w');
+for i = 1:5
+    fprintf(fid, '%d %d\n', trans_pt(i,1), trans_pt(i,2));%will be read as np.int32
+end
+fclose(fid);
+end

APDrawingGAN2/preprocess/get_partmask.py

@@ -0,0 +1,152 @@
+import cv2
+import os, glob, csv, shutil
+import numpy as np
+import dlib
+import math
+from shapely.geometry import Point
+from shapely.geometry import Polygon
+import sys
+
+detector = dlib.get_frontal_face_detector()
+predictor = dlib.shape_predictor('../checkpoints/shape_predictor_68_face_landmarks.dat')
+
+def getfeats(featpath):
+	trans_points = np.empty([68,2],dtype=np.int64) 
+	with open(featpath, 'r') as csvfile:
+		reader = csv.reader(csvfile, delimiter=' ')
+		for ind,row in enumerate(reader):
+			trans_points[ind,:] = row
+	return trans_points
+
+def getinternal(lm1,lm2):
+	lminternal = []
+	if abs(lm1[1]-lm2[1]) > abs(lm1[0]-lm2[0]):
+		if lm1[1] > lm2[1]:
+			tmp = lm1
+			lm1 = lm2
+			lm2 = tmp
+		for y in range(lm1[1]+1,lm2[1]):
+			x = int(round(float(y-lm1[1])/(lm2[1]-lm1[1])*(lm2[0]-lm1[0])+lm1[0]))
+			lminternal.append((x,y))
+	else:
+		if lm1[0] > lm2[0]:
+			tmp = lm1
+			lm1 = lm2
+			lm2 = tmp
+		for x in range(lm1[0]+1,lm2[0]):
+			y = int(round(float(x-lm1[0])/(lm2[0]-lm1[0])*(lm2[1]-lm1[1])+lm1[1]))
+			lminternal.append((x,y))
+	return lminternal
+
+def mulcross(p,x_1,x):#p-x_1,x-x_1
+	vp = [p[0]-x_1[0],p[1]-x_1[1]]
+	vq = [x[0]-x_1[0],x[1]-x_1[1]]
+	return vp[0]*vq[1]-vp[1]*vq[0]
+
+def shape_to_np(shape, dtype="int"):
+	# initialize the list of (x, y)-coordinates
+	coords = np.zeros((shape.num_parts, 2), dtype=dtype)
+	# loop over all facial landmarks and convert them
+	# to a 2-tuple of (x, y)-coordinates
+	for i in range(0, shape.num_parts):
+		coords[i] = (shape.part(i).x, shape.part(i).y)
+	# return the list of (x, y)-coordinates
+	return coords
+
+def get_68lm(imgfile,savepath):
+	image = cv2.imread(imgfile)
+	rgbImg = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+	rects = detector(rgbImg, 1)
+	for (i, rect) in enumerate(rects):
+		landmarks = predictor(rgbImg, rect)
+		landmarks = shape_to_np(landmarks)
+		f = open(savepath,'w')
+		for i in range(len(landmarks)):
+			lm = landmarks[i]
+			print(lm[0], lm[1], file=f)
+		f.close()
+
+def get_partmask(imgfile,part,lmpath,savefile):
+	img = cv2.imread(imgfile)
+	mask = np.zeros(img.shape, np.uint8)
+	lms = getfeats(lmpath)
+
+	if os.path.exists(savefile):
+		return
+	
+	if part == 'nose':
+		# 27,31....,35 -> up, left, right, lower5 -- eight points
+		up = [int(round(1.2*lms[27][0]-0.2*lms[33][0])),int(round(1.2*lms[27][1]-0.2*lms[33][1]))]
+		lower5 = [[0,0]]*5
+		for i in range(31,36):
+			lower5[i-31] = [int(round(1.1*lms[i][0]-0.1*lms[27][0])),int(round(1.1*lms[i][1]-0.1*lms[27][1]))]
+		ratio = 2.5
+		left = [int(round(ratio*lower5[0][0]-(ratio-1)*lower5[1][0])),int(round(ratio*lower5[0][1]-(ratio-1)*lower5[1][1]))]
+		right = [int(round(ratio*lower5[4][0]-(ratio-1)*lower5[3][0])),int(round(ratio*lower5[4][1]-(ratio-1)*lower5[3][1]))]
+		loop = [up,left,lower5[0],lower5[1],lower5[2],lower5[3],lower5[4],right]
+	elif part == 'eyel':
+		height = max(lms[41][1]-lms[37][1],lms[40][1]-lms[38][1])
+		width = lms[39][0]-lms[36][0]
+		ratio = 0.1
+		gap = int(math.ceil(width*ratio))
+		ratio2 = 0.6
+		gaph = int(math.ceil(height*ratio2))
+		ratio3 = 1.5
+		gaph2 = int(math.ceil(height*ratio3))
+		upper = [[lms[17][0]-2*gap,lms[17][1]],[lms[17][0]-2*gap,lms[17][1]-gaph],[lms[18][0],lms[18][1]-gaph],[lms[19][0],lms[19][1]-gaph],[lms[20][0],lms[20][1]-gaph],[lms[21][0]+gap*2,lms[21][1]-gaph]]
+		lower = [[lms[39][0]+gap,lms[40][1]+gaph2],[lms[40][0],lms[40][1]+gaph2],[lms[41][0],lms[41][1]+gaph2],[lms[36][0]-2*gap,lms[41][1]+gaph2]]
+		loop = upper + lower
+		loop.reverse()
+	elif part == 'eyer':
+		height = max(lms[47][1]-lms[43][1],lms[46][1]-lms[44][1])
+		width = lms[45][0]-lms[42][0]
+		ratio = 0.1
+		gap = int(math.ceil(width*ratio))
+		ratio2 = 0.6
+		gaph = int(math.ceil(height*ratio2))
+		ratio3 = 1.5
+		gaph2 = int(math.ceil(height*ratio3))
+		upper = [[lms[22][0]-2*gap,lms[22][1]],[lms[22][0]-2*gap,lms[22][1]-gaph],[lms[23][0],lms[23][1]-gaph],[lms[24][0],lms[24][1]-gaph],[lms[25][0],lms[25][1]-gaph],[lms[26][0]+gap*2,lms[26][1]-gaph]]
+		lower = [[lms[45][0]+2*gap,lms[46][1]+gaph2],[lms[46][0],lms[46][1]+gaph2],[lms[47][0],lms[47][1]+gaph2],[lms[42][0]-gap,lms[42][1]+gaph2]]
+		loop = upper + lower
+		loop.reverse()
+	elif part == 'mouth':
+		height = lms[62][1]-lms[51][1]
+		width = lms[54][0]-lms[48][0]
+		ratio = 1
+		ratio2 = 0.2#0.1
+		gaph = int(math.ceil(ratio*height))
+		gapw = int(math.ceil(ratio2*width))
+		left = [(lms[48][0]-gapw,lms[48][1])]
+		upper = [(lms[i][0], lms[i][1]-gaph) for i in range(48,55)]
+		right = [(lms[54][0]+gapw,lms[54][1])]
+		lower = [(lms[i][0], lms[i][1]+gaph) for i in list(range(54,60))+[48]]
+		loop = left + upper + right + lower
+		loop.reverse()
+		pl = Polygon(loop)
+
+	for i in range(mask.shape[0]):
+		for j in range(mask.shape[1]):
+			if part != 'mouth' and part != 'jaw':
+				p = [j,i]
+				flag = 1
+				for k in range(len(loop)):
+					if mulcross(p,loop[k],loop[(k+1)%len(loop)]) < 0:#y downside... >0 represents counter-clockwise, <0 clockwise
+						flag = 0
+						break
+			else:
+				p = Point(j,i)
+				flag = pl.contains(p)
+			if flag:
+				mask[i,j] = [255,255,255]
+	if not os.path.exists(os.path.dirname(savefile)):
+		os.mkdir(os.path.dirname(savefile))
+	cv2.imwrite(savefile,mask)
+
+if __name__ == '__main__':
+	imgfile = 'example/img_1701_aligned.png'
+	lmfile = 'example/img_1701_aligned_68lm.txt'
+	get_68lm(imgfile,lmfile)
+	for part in ['eyel','eyer','nose','mouth']:
+		savepath = 'example/img_1701_aligned_'+part+'mask.png'
+		get_partmask(imgfile,part,lmfile,savepath)

APDrawingGAN2/preprocess/readme.md

@@ -0,0 +1,71 @@
+## Preprocessing steps
+
+Both training and testing images need: 
+
+- align to 512x512
+- facial landmarks
+- mask for eyes,nose,mouth,background
+
+Training images additionally need:
+
+- mask for face region
+
+
+### 1. Align, resize, crop images to 512x512, and get facial landmarks
+
+All training and testing images in our model are aligned using facial landmarks. And landmarks after alignment are needed in our code.
+
+- First, 5 facial landmark for a face photo need to be detected (we detect using [MTCNN](https://github.com/kpzhang93/MTCNN_face_detection_alignment)(MTCNNv1)).
+
+- Then, we provide a matlab function in `face_align_512.m` to align, resize and crop face photos (and corresponding drawings) to 512x512.Call this function in MATLAB to align the image to 512x512.
+For example, for `img_1701.jpg` in `example` dir, 5 detected facial landmark is saved in `example/img_1701_facial5point.mat`. Call following in MATLAB:
+```bash
+load('example/img_1701_facial5point.mat');
+[trans_img,trans_facial5point]=face_align_512('example/img_1701.jpg',facial5point,'example');
+```
+
+This will align the image, and output aligned image + transformed facial landmark (in txt format) in `example` folder.
+See `face_align_512.m` for more instructions.
+
+The saved transformed facial landmark need to be copied to `dataset/landmark/`, and has the **same filename** with aligned face photos (e.g. `dataset/data/test_single/31.png` should have landmark file `dataset/landmark/31.txt`).
+
+### 2. Prepare background masks
+
+In our work, background mask is segmented by method in
+"Automatic Portrait Segmentation for Image Stylization"
+Xiaoyong Shen, Aaron Hertzmann, Jiaya Jia, Sylvain Paris, Brian Price, Eli Shechtman, Ian Sachs. Computer Graphics Forum, 35(2)(Proc. Eurographics), 2016.
+
+We use code in http://xiaoyongshen.me/webpage_portrait/index.html to detect background masks for aligned face photos.
+An example background mask is shown in `example/img_1701_aligned_bgmask.png`.
+
+The background masks need to be copied to `dataset/mask/bg/`, and has the **same filename** with aligned face photos (e.g. `dataset/data/test_single/31.png` should have background mask `dataset/mask/bg/31.png`)  
+
+### 3. Prepare eyes/nose/mouth masks
+
+We use dlib to extract 68 landmarks for aligned face photos, and use these landmarks to get masks for local regions.
+See an example in `get_partmask.py`, the eyes, nose, mouth masks for `example/img_1701_aligned.png` are `example/img_1701_aligned_[part]mask.png`, where part is in [eyel,eyer,nose,mouth].
+
+The part masks need to be copied to `dataset/mask/[part]/`, and has the **same filename** with aligned face photos.
+
+### 4. (For training) Prepare face masks
+
+We use the face parsing net in https://github.com/cientgu/Mask_Guided_Portrait_Editing to detect face region.
+The face parsing net will label each face into 11 classes, the 0 is for background, 10 is for hair, and the 1~9 are face regions.
+An example face mask is shown in `example/img_1701_aligned_facemask.png`.
+
+The face masks need to be copied to `dataset/mask/face/`, and has the **same filename** with aligned face photos.
+
+### 5. (For training) Combine A and B
+
+We provide a python script to generate training data in the form of pairs of images {A,B}, i.e. pairs {face photo, drawing}. This script will concatenate each pair of images horizontally into one single image. Then we can learn to translate A to B:
+
+Create folder `/path/to/data` with subfolders `A` and `B`. `A` and `B` should each have their own subfolders `train`, `test`, etc. In `/path/to/data/A/train`, put training face photos. In `/path/to/data/B/train`, put the corresponding artist drawings. Repeat same for `test`.
+
+Corresponding images in a pair {A,B} must both be images after aligning and of size 512x512, and have the same filename, e.g., `/path/to/data/A/train/1.png` is considered to correspond to `/path/to/data/B/train/1.png`.
+
+Once the data is formatted this way, call:
+```bash
+python datasets/combine_A_and_B.py --fold_A /path/to/data/A --fold_B /path/to/data/B --fold_AB /path/to/data
+```
+
+This will combine each pair of images (A,B) into a single image file, ready for training.

APDrawingGAN2/readme.md

@@ -0,0 +1,105 @@
+
+# APDrawingGAN++
+
+We provide PyTorch implementations for our TPAMI paper "Line Drawings for Face Portraits from Photos using Global and Local Structure based GANs".
+It is a journal extension of our previous CVPR 2019 work [APDrawingGAN](https://github.com/yiranran/APDrawingGAN).
+
+This project generates artistic portrait drawings from face photos using a GAN-based model.
+You may find useful information in [preprocessing steps](preprocess/readme.md) and [training/testing tips](docs/tips.md).
+
+[[Jittor implementation]](https://github.com/yiranran/APDrawingGAN2-Jittor)
+
+## Our Proposed Framework
+ 
+<img src = 'imgs/architecture-pami.jpg'>
+
+## Sample Results
+Up: input, Down: output
+<p>
+    <img src='imgs/sample/140_large-img_1696_real_A.png' width="16%"/>
+    <img src='imgs/sample/140_large-img_1615_real_A.png' width="16%"/>
+    <img src='imgs/sample/140_large-img_1684_real_A.png' width="16%"/>
+    <img src='imgs/sample/140_large-img_1616_real_A.png' width="16%"/>
+    <img src='imgs/sample/140_large-img_1673_real_A.png' width="16%"/>
+    <img src='imgs/sample/140_large-img_1701_real_A.png' width="16%"/>
+</p>
+<p>
+    <img src='imgs/sample/140_large-img_1696_fake_B.png' width="16%"/>
+    <img src='imgs/sample/140_large-img_1615_fake_B.png' width="16%"/>
+    <img src='imgs/sample/140_large-img_1684_fake_B.png' width="16%"/>
+    <img src='imgs/sample/140_large-img_1616_fake_B.png' width="16%"/>
+    <img src='imgs/sample/140_large-img_1673_fake_B.png' width="16%"/>
+    <img src='imgs/sample/140_large-img_1701_fake_B.png' width="16%"/>
+</p>
+
+## Citation
+If you use this code for your research, please cite our paper.
+```
+@inproceedings{YiXLLR20,
+  title     = {Line Drawings for Face Portraits from Photos using Global and Local Structure based {GAN}s},
+  author    = {Yi, Ran and Xia, Mengfei and Liu, Yong-Jin and Lai, Yu-Kun and Rosin, Paul L},
+  booktitle = {{IEEE} Transactions on Pattern Analysis and Machine Intelligence (TPAMI)},
+  doi       = {10.1109/TPAMI.2020.2987931},
+  year      = {2020}
+}
+```
+
+## Prerequisites
+- Linux or macOS
+- Python 2 or 3
+- CPU or NVIDIA GPU + CUDA CuDNN
+
+
+## Getting Started
+### 1.Installation
+```bash
+pip install -r requirements.txt
+```
+
+### 2.Quick Start (Apply a Pre-trained Model)
+- Download APDrawing dataset from [BaiduYun](https://pan.baidu.com/s/1cN5gEYJ2tnE9WboLA79Z5g)(extract code:0zuv) or [YandexDrive](https://yadi.sk/d/4vWhi8-ZQj_nRw), and extract to `dataset`.
+
+- Download pre-trained models and auxiliary nets from [BaiduYun](https://pan.baidu.com/s/1nrtCHQmgcwbSGxWuAVzWhA)(extract code:imqp) or [YandexDrive](https://yadi.sk/d/DS4271lbEPhGVQ), and extract to `checkpoints`.
+
+- Generate artistic portrait drawings for example photos in `dataset/test_single` using
+``` bash
+python test.py --dataroot dataset/test_single --name apdrawinggan++_author --model test --use_resnet --netG resnet_9blocks --which_epoch 150 --how_many 1000 --gpu_ids 0 --gpu_ids_p 0 --imagefolder images-single
+```
+The test results will be saved to a html file here: `./results/apdrawinggan++_author/test_150/index-single.html`.
+
+- If you want to test on your own data, please first align your pictures and prepare your data's facial landmarks and masks according to tutorial in [preprocessing steps](preprocess/readme.md), then change the --dataroot flag above to your directory of aligned photos.
+
+### 3.Train
+- Run `python -m visdom.server`
+- Train a model (with pre-training as initialization):
+first copy "pre2" models into checkpoints dir of current experiment, e.g. `checkpoints/apdrawinggan++_1`.
+```bash
+mkdir checkpoints/apdrawinggan++_1/
+cp checkpoints/pre2/*.pt checkpoints/apdrawinggan++_1/
+python train.py --dataroot dataset/AB_140_aug3_H_hm2 --name apdrawinggan++_1 --model apdrawingpp_style --use_resnet --netG resnet_9blocks --continue_train --continuity_loss --lambda_continuity 40.0 --gpu_ids 0 --gpu_ids_p 1 --display_env apdrawinggan++_1 --niter 200 --niter_decay 0 --lr 0.0001 --batch_size 1 --emphasis_conti_face --auxiliary_root auxiliaryeye2o
+```
+- To view training results and loss plots, click the URL http://localhost:8097. To see more intermediate results, check out  `./checkpoints/apdrawinggan++_1/web/index.html`
+
+### 4.Test
+- To test the model on test set:
+```bash
+python test.py --dataroot dataset/AB_140_aug3_H_hm2 --name apdrawinggan++_author --model apdrawingpp_style --use_resnet --netG resnet_9blocks --which_epoch 150 --how_many 1000 --gpu_ids 0 --gpu_ids_p 0 --imagefolder images-apd70
+```
+The test results will be saved to a html file: `./results/apdrawinggan++_author/test_150/index-apd70.html`.
+
+- To test the model on images without paired ground truth, same as 2. Apply a pre-trained model.
+
+You can find these scripts at `scripts` directory.
+
+
+## [Preprocessing Steps](preprocess/readme.md)
+Preprocessing steps for your own data (either for testing or training).
+
+
+## [Training/Test Tips](docs/tips.md)
+Best practice for training and testing your models.
+
+You can contact email yr16@mails.tsinghua.edu.cn for any questions.
+
+## Acknowledgments
+Our code is inspired by [pytorch-CycleGAN-and-pix2pix](https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix).

APDrawingGAN2/requirements.txt

@@ -0,0 +1,10 @@
+torch==1.1.0
+torchvision==0.4.0
+dominate==2.4.0
+visdom==0.1.8.9
+scipy==1.1.0
+numpy==1.16.4
+Pillow==4.3.0
+opencv-python==4.1.0.25
+dlib==19.18.0
+shapely==1.7.0

APDrawingGAN2/script/test.sh

@@ -0,0 +1,2 @@
+set -ex
+python test.py --dataroot dataset/AB_140_aug3_H_hm2 --name apdrawinggan++_author --model apdrawingpp_style --use_resnet --netG resnet_9blocks --which_epoch 150 --how_many 1000 --gpu_ids 0 --gpu_ids_p 0 --imagefolder images-apd70

APDrawingGAN2/script/test_single.sh

@@ -0,0 +1,2 @@
+set -ex
+python test.py --dataroot dataset/test_single --name apdrawinggan++_author --model test --use_resnet --netG resnet_9blocks --which_epoch 150 --how_many 1000 --gpu_ids 0 --gpu_ids_p 0 --imagefolder images-single

APDrawingGAN2/script/train.sh

@@ -0,0 +1,3 @@
+set -ex
+python train.py --dataroot dataset/AB_140_aug3_H_hm2 --name apdrawinggan++_1 --model apdrawingpp_style --use_resnet --netG resnet_9blocks --continue_train --continuity_loss --lambda_continuity 40.0 --gpu_ids 0 --gpu_ids_p 1 --display_env apdrawinggan++_1 --niter 200 --niter_decay 0 --lr 0.0001 --batch_size 1 --emphasis_conti_face --auxiliary_root auxiliaryeye2o
+

APDrawingGAN2/test.py

@@ -0,0 +1,69 @@
+import os
+from options.test_options import TestOptions
+from data import CreateDataLoader
+from models import create_model
+from util.visualizer import save_images
+from util import html
+
+
+if __name__ == '__main__':
+    opt = TestOptions().parse()
+    opt.num_threads = 1   # test code only supports num_threads = 1
+    opt.batch_size = 1  # test code only supports batch_size = 1
+    opt.serial_batches = True  # no shuffle
+    opt.no_flip = True  # no flip
+    opt.display_id = -1  # no visdom display
+    data_loader = CreateDataLoader(opt)
+    dataset = data_loader.load_data()
+    model = create_model(opt)
+    model.setup(opt)
+    # create website
+    web_dir = os.path.join(opt.results_dir, opt.name, '%s_%s' % (opt.phase, opt.which_epoch))
+    #webpage = html.HTML(web_dir, 'Experiment = %s, Phase = %s, Epoch = %s' % (opt.name, opt.phase, opt.which_epoch))
+    webpage = html.HTML(web_dir, 'Experiment = %s, Phase = %s, Epoch = %s' % (opt.name, opt.phase, opt.which_epoch),reflesh=0, folder=opt.imagefolder)
+    if opt.test_continuity_loss:
+        file_name = os.path.join(opt.results_dir, opt.name, '%s_%s' % (opt.phase, opt.which_epoch), 'continuity.txt')
+        file_name1 = os.path.join(opt.results_dir, opt.name, '%s_%s' % (opt.phase, opt.which_epoch), 'continuity-r.txt')
+        if os.path.exists(file_name):
+            os.remove(file_name)
+        if os.path.exists(file_name1):
+            os.remove(file_name1)
+    # test
+    #model.eval()
+    for i, data in enumerate(dataset):
+        if i >= opt.how_many:#test code only supports batch_size = 1, how_many means how many test images to run
+            break
+        model.set_input(data)
+        model.test()
+        visuals = model.get_current_visuals()#in test the loadSize is set to the same as fineSize
+        img_path = model.get_image_paths()
+        #if i % 5 == 0:
+        #    print('processing (%04d)-th image... %s' % (i, img_path))
+        save_images(webpage, visuals, img_path, aspect_ratio=opt.aspect_ratio, width=opt.display_winsize)
+
+    webpage.save()
+    if opt.model == 'regressor':
+        print(model.cnt)
+        print(model.value/model.cnt)
+        print(model.minval)
+        print(model.avg/model.cnt)
+        print(model.max)
+        html = os.path.join(web_dir,'cindex'+opt.imagefolder[6:]+'.html')
+        f=open(html,'w')
+        print('<table border="1" style=\"text-align:center;\">',file=f,end='')
+        print('<tr>',file=f,end='')
+        print('<td>image name</td>',file=f,end='')
+        print('<td>realA</td>',file=f,end='')
+        print('<td>realB</td>',file=f,end='')
+        print('<td>fakeB</td>',file=f,end='')
+        print('</tr>',file=f,end='')
+        for info in model.info:
+            basen = os.path.basename(info[0])[:-4]
+            print('<tr>',file=f,end='')
+            print('<td>%s</td>'%basen,file=f,end='')
+            print('<td><img src=\"%s/%s_real_A.png\" style=\"width:44px\"></td>'%(opt.imagefolder,basen),file=f,end='')
+            print('<td>%.4f</td>'%info[1],file=f,end='')
+            print('<td>%.4f</td>'%info[2],file=f,end='')
+            print('</tr>',file=f,end='')
+        print('</table>',file=f,end='')
+        f.close()

APDrawingGAN2/train.py

@@ -0,0 +1,67 @@
+import time
+from options.train_options import TrainOptions
+from data import CreateDataLoader
+from models import create_model
+from util.visualizer import Visualizer
+
+if __name__ == '__main__':
+    start = time.time()
+    opt = TrainOptions().parse()
+    data_loader = CreateDataLoader(opt)
+    dataset = data_loader.load_data()
+    dataset_size = len(data_loader)
+    print('#training images = %d' % dataset_size)
+
+    model = create_model(opt)
+    model.setup(opt)
+    visualizer = Visualizer(opt)
+    total_steps = 0
+    model.save_networks2(opt.which_epoch)
+
+    for epoch in range(opt.epoch_count, opt.niter + opt.niter_decay + 1):
+        epoch_start_time = time.time()
+        iter_data_time = time.time()
+        epoch_iter = 0
+
+        for i, data in enumerate(dataset):
+            iter_start_time = time.time()
+            if total_steps % opt.print_freq == 0:
+                t_data = iter_start_time - iter_data_time
+            visualizer.reset()
+            total_steps += opt.batch_size
+            epoch_iter += opt.batch_size
+            model.set_input(data)
+            model.optimize_parameters()
+
+            if total_steps % opt.display_freq == 0:
+                save_result = total_steps % opt.update_html_freq == 0
+                visualizer.display_current_results(model.get_current_visuals(), epoch, save_result)
+                #print('display',total_steps)
+
+            if total_steps % opt.print_freq == 0:#print freq 100
+                losses = model.get_current_losses()
+                t = (time.time() - iter_start_time) / opt.batch_size
+                visualizer.print_current_losses(epoch, epoch_iter, losses, t, t_data)
+                if opt.display_id > 0:
+                    visualizer.plot_current_losses(epoch, float(epoch_iter) / dataset_size, opt, losses)
+            
+            if total_steps % opt.save_latest_freq == 0:
+                print('saving the latest model (epoch %d, total_steps %d)' %
+                      (epoch, total_steps))
+                #model.save_networks('latest')
+                model.save_networks2('latest')
+
+            iter_data_time = time.time()
+        if epoch % opt.save_epoch_freq == 0:
+            print('saving the model at the end of epoch %d, iters %d' %
+                  (epoch, total_steps))
+            #model.save_networks('latest')
+            #model.save_networks(epoch)
+            model.save_networks2('latest')
+            model.save_networks2(epoch)
+
+        print('End of epoch %d / %d \t Time Taken: %d sec' %
+              (epoch, opt.niter + opt.niter_decay, time.time() - epoch_start_time))
+        model.update_learning_rate()
+
+    print('Total Time Taken: %d sec' % (time.time() - start))

APDrawingGAN2/util/get_data.py

@@ -0,0 +1,115 @@
+from __future__ import print_function
+import os
+import tarfile
+import requests
+from warnings import warn
+from zipfile import ZipFile
+from bs4 import BeautifulSoup
+from os.path import abspath, isdir, join, basename
+
+
+class GetData(object):
+    """
+
+    Download CycleGAN or Pix2Pix Data.
+
+    Args:
+        technique : str
+            One of: 'cyclegan' or 'pix2pix'.
+        verbose : bool
+            If True, print additional information.
+
+    Examples:
+        >>> from util.get_data import GetData
+        >>> gd = GetData(technique='cyclegan')
+        >>> new_data_path = gd.get(save_path='./datasets')  # options will be displayed.
+
+    """
+
+    def __init__(self, technique='cyclegan', verbose=True):
+        url_dict = {
+            'pix2pix': 'https://people.eecs.berkeley.edu/~tinghuiz/projects/pix2pix/datasets',
+            'cyclegan': 'https://people.eecs.berkeley.edu/~taesung_park/CycleGAN/datasets'
+        }
+        self.url = url_dict.get(technique.lower())
+        self._verbose = verbose
+
+    def _print(self, text):
+        if self._verbose:
+            print(text)
+
+    @staticmethod
+    def _get_options(r):
+        soup = BeautifulSoup(r.text, 'lxml')
+        options = [h.text for h in soup.find_all('a', href=True)
+                   if h.text.endswith(('.zip', 'tar.gz'))]
+        return options
+
+    def _present_options(self):
+        r = requests.get(self.url)
+        options = self._get_options(r)
+        print('Options:\n')
+        for i, o in enumerate(options):
+            print("{0}: {1}".format(i, o))
+        choice = input("\nPlease enter the number of the "
+                       "dataset above you wish to download:")
+        return options[int(choice)]
+
+    def _download_data(self, dataset_url, save_path):
+        if not isdir(save_path):
+            os.makedirs(save_path)
+
+        base = basename(dataset_url)
+        temp_save_path = join(save_path, base)
+
+        with open(temp_save_path, "wb") as f:
+            r = requests.get(dataset_url)
+            f.write(r.content)
+
+        if base.endswith('.tar.gz'):
+            obj = tarfile.open(temp_save_path)
+        elif base.endswith('.zip'):
+            obj = ZipFile(temp_save_path, 'r')
+        else:
+            raise ValueError("Unknown File Type: {0}.".format(base))
+
+        self._print("Unpacking Data...")
+        obj.extractall(save_path)
+        obj.close()
+        os.remove(temp_save_path)
+
+    def get(self, save_path, dataset=None):
+        """
+
+        Download a dataset.
+
+        Args:
+            save_path : str
+                A directory to save the data to.
+            dataset : str, optional
+                A specific dataset to download.
+                Note: this must include the file extension.
+                If None, options will be presented for you
+                to choose from.
+
+        Returns:
+            save_path_full : str
+                The absolute path to the downloaded data.
+
+        """
+        if dataset is None:
+            selected_dataset = self._present_options()
+        else:
+            selected_dataset = dataset
+
+        save_path_full = join(save_path, selected_dataset.split('.')[0])
+
+        if isdir(save_path_full):
+            warn("\n'{0}' already exists. Voiding Download.".format(
+                save_path_full))
+        else:
+            self._print('Downloading Data...')
+            url = "{0}/{1}".format(self.url, selected_dataset)
+            self._download_data(url, save_path=save_path)
+
+        return abspath(save_path_full)

APDrawingGAN2/util/html.py

@@ -0,0 +1,68 @@
+import dominate
+from dominate.tags import *
+import os
+
+
+class HTML:
+    def __init__(self, web_dir, title, reflesh=0, folder='images'):
+        self.title = title
+        self.web_dir = web_dir
+        #self.img_dir = os.path.join(self.web_dir, 'images')
+        self.img_dir = os.path.join(self.web_dir, folder)
+        self.folder = folder
+        if not os.path.exists(self.web_dir):
+            os.makedirs(self.web_dir)
+        if not os.path.exists(self.img_dir):
+            os.makedirs(self.img_dir)
+        # print(self.img_dir)
+
+        self.doc = dominate.document(title=title)
+        if reflesh > 0:
+            with self.doc.head:
+                meta(http_equiv="reflesh", content=str(reflesh))
+
+    def get_image_dir(self):
+        return self.img_dir
+
+    def add_header(self, str):
+        with self.doc:
+            h3(str)
+
+    def add_table(self, border=1):
+        self.t = table(border=border, style="table-layout: fixed;")
+        self.doc.add(self.t)
+
+    def add_images(self, ims, txts, links, width=400):
+        self.add_table()
+        with self.t:
+            with tr():
+                for im, txt, link in zip(ims, txts, links):
+                    with td(style="word-wrap: break-word;", halign="center", valign="top"):
+                        with p():
+                            with a(href=os.path.join('images', link)):
+                                #img(style="width:%dpx" % width, src=os.path.join('images', im))
+                                img(style="width:%dpx" % width, src=os.path.join(self.folder, im))
+                            br()
+                            p(txt)
+
+    def save(self):
+        #html_file = '%s/index.html' % self.web_dir
+        html_file = '%s/index%s.html' % (self.web_dir, self.folder[6:])
+        f = open(html_file, 'wt')
+        f.write(self.doc.render())
+        f.close()
+
+
+if __name__ == '__main__':
+    html = HTML('web/', 'test_html')
+    html.add_header('hello world')
+
+    ims = []
+    txts = []
+    links = []
+    for n in range(4):
+        ims.append('image_%d.png' % n)
+        txts.append('text_%d' % n)
+        links.append('image_%d.png' % n)
+    html.add_images(ims, txts, links)
+    html.save()

APDrawingGAN2/util/image_pool.py

@@ -0,0 +1,32 @@
+import random
+import torch
+
+
+class ImagePool():
+    def __init__(self, pool_size):
+        self.pool_size = pool_size
+        if self.pool_size > 0:
+            self.num_imgs = 0
+            self.images = []
+
+    def query(self, images):
+        if self.pool_size == 0:
+            return images
+        return_images = []
+        for image in images:
+            image = torch.unsqueeze(image.data, 0)
+            if self.num_imgs < self.pool_size:
+                self.num_imgs = self.num_imgs + 1
+                self.images.append(image)
+                return_images.append(image)
+            else:
+                p = random.uniform(0, 1)
+                if p > 0.5:
+                    random_id = random.randint(0, self.pool_size - 1)  # randint is inclusive
+                    tmp = self.images[random_id].clone()
+                    self.images[random_id] = image
+                    return_images.append(tmp)
+                else:
+                    return_images.append(image)
+        return_images = torch.cat(return_images, 0)
+        return return_images

APDrawingGAN2/util/util.py

@@ -0,0 +1,60 @@
+from __future__ import print_function
+import torch
+import numpy as np
+from PIL import Image
+import os
+
+
+# Converts a Tensor into an image array (numpy)
+# |imtype|: the desired type of the converted numpy array
+def tensor2im(input_image, imtype=np.uint8):
+    if isinstance(input_image, torch.Tensor):
+        image_tensor = input_image.data
+    else:
+        return input_image
+    image_numpy = image_tensor[0].cpu().float().numpy()
+    if image_numpy.shape[0] == 1:
+        image_numpy = np.tile(image_numpy, (3, 1, 1))
+    image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0
+    return image_numpy.astype(imtype)
+
+
+def diagnose_network(net, name='network'):
+    mean = 0.0
+    count = 0
+    for param in net.parameters():
+        if param.grad is not None:
+            mean += torch.mean(torch.abs(param.grad.data))
+            count += 1
+    if count > 0:
+        mean = mean / count
+    print(name)
+    print(mean)
+
+
+def save_image(image_numpy, image_path):
+    image_pil = Image.fromarray(image_numpy)
+    image_pil.save(image_path)
+
+
+def print_numpy(x, val=True, shp=False):
+    x = x.astype(np.float64)
+    if shp:
+        print('shape,', x.shape)
+    if val:
+        x = x.flatten()
+        print('mean = %3.3f, min = %3.3f, max = %3.3f, median = %3.3f, std=%3.3f' % (
+            np.mean(x), np.min(x), np.max(x), np.median(x), np.std(x)))
+
+
+def mkdirs(paths):
+    if isinstance(paths, list) and not isinstance(paths, str):
+        for path in paths:
+            mkdir(path)
+    else:
+        mkdir(paths)
+
+
+def mkdir(path):
+    if not os.path.exists(path):
+        os.makedirs(path)

APDrawingGAN2/util/visualizer.py

@@ -0,0 +1,171 @@
+import numpy as np
+import os
+import ntpath
+import time
+from . import util
+from . import html
+from scipy.misc import imresize
+
+
+# save image to the disk
+def save_images(webpage, visuals, image_path, aspect_ratio=1.0, width=256):
+    image_dir = webpage.get_image_dir()
+    short_path = ntpath.basename(image_path[0])
+    name = os.path.splitext(short_path)[0]
+
+    webpage.add_header(name)
+    ims, txts, links = [], [], []
+
+    for label, im_data in visuals.items():
+        im = util.tensor2im(im_data)#tensor to numpy array [-1,1]->[0,1]->[0,255]
+        image_name = '%s_%s.png' % (name, label)
+        save_path = os.path.join(image_dir, image_name)
+        h, w, _ = im.shape
+        if aspect_ratio > 1.0:
+            im = imresize(im, (h, int(w * aspect_ratio)), interp='bicubic')
+        if aspect_ratio < 1.0:
+            im = imresize(im, (int(h / aspect_ratio), w), interp='bicubic')
+        util.save_image(im, save_path)
+
+        ims.append(image_name)
+        txts.append(label)
+        links.append(image_name)
+    webpage.add_images(ims, txts, links, width=width)
+
+
+class Visualizer():
+    def __init__(self, opt):
+        self.display_id = opt.display_id
+        self.use_html = opt.isTrain and not opt.no_html
+        self.win_size = opt.display_winsize
+        self.name = opt.name
+        self.opt = opt
+        self.saved = False
+        if self.display_id > 0:
+            import visdom
+            self.ncols = opt.display_ncols
+            self.vis = visdom.Visdom(server=opt.display_server, port=opt.display_port, env=opt.display_env, raise_exceptions=True)
+
+        if self.use_html:
+            self.web_dir = os.path.join(opt.checkpoints_dir, opt.name, 'web')
+            self.img_dir = os.path.join(self.web_dir, 'images')
+            print('create web directory %s...' % self.web_dir)
+            util.mkdirs([self.web_dir, self.img_dir])
+        self.log_name = os.path.join(opt.checkpoints_dir, opt.name, 'loss_log.txt')
+        with open(self.log_name, "a") as log_file:
+            now = time.strftime("%c")
+            log_file.write('================ Training Loss (%s) ================\n' % now)
+
+    def reset(self):
+        self.saved = False
+
+    def throw_visdom_connection_error(self):
+        print('\n\nCould not connect to Visdom server (https://github.com/facebookresearch/visdom) for displaying training progress.\nYou can suppress connection to Visdom using the option --display_id -1. To install visdom, run \n$ pip install visdom\n, and start the server by \n$ python -m visdom.server.\n\n')
+        exit(1)
+
+    # |visuals|: dictionary of images to display or save
+    def display_current_results(self, visuals, epoch, save_result):
+        if self.display_id > 0:  # show images in the browser
+            ncols = self.ncols
+            if ncols > 0:
+                ncols = min(ncols, len(visuals))
+                h, w = next(iter(visuals.values())).shape[:2]
+                table_css = """<style>
+                        table {border-collapse: separate; border-spacing:4px; white-space:nowrap; text-align:center}
+                        table td {width: %dpx; height: %dpx; padding: 4px; outline: 4px solid black}
+                        </style>""" % (w, h)
+                title = self.name
+                label_html = ''
+                label_html_row = ''
+                images = []
+                idx = 0
+                for label, image in visuals.items():
+                    image_numpy = util.tensor2im(image)
+                    label_html_row += '<td>%s</td>' % label
+                    images.append(image_numpy.transpose([2, 0, 1]))
+                    idx += 1
+                    if idx % ncols == 0:
+                        label_html += '<tr>%s</tr>' % label_html_row
+                        label_html_row = ''
+                white_image = np.ones_like(image_numpy.transpose([2, 0, 1])) * 255
+                while idx % ncols != 0:
+                    images.append(white_image)
+                    label_html_row += '<td></td>'
+                    idx += 1
+                if label_html_row != '':
+                    label_html += '<tr>%s</tr>' % label_html_row
+                # pane col = image row
+                try:
+                    self.vis.images(images, nrow=ncols, win=self.display_id + 1,
+                                    padding=2, opts=dict(title=title + ' images'))
+                    label_html = '<table>%s</table>' % label_html
+                    self.vis.text(table_css + label_html, win=self.display_id + 2,
+                                  opts=dict(title=title + ' labels'))
+                except ConnectionError:
+                    self.throw_visdom_connection_error()
+
+            else:
+                idx = 1
+                for label, image in visuals.items():
+                    image_numpy = util.tensor2im(image)
+                    self.vis.image(image_numpy.transpose([2, 0, 1]), opts=dict(title=label),
+                                   win=self.display_id + idx)
+                    idx += 1
+
+        if self.use_html and (save_result or not self.saved):  # save images to a html file
+            self.saved = True
+            for label, image in visuals.items():
+                image_numpy = util.tensor2im(image)
+                img_path = os.path.join(self.img_dir, 'epoch%.3d_%s.png' % (epoch, label))
+                util.save_image(image_numpy, img_path)
+            # update website
+            webpage = html.HTML(self.web_dir, 'Experiment name = %s' % self.name, reflesh=1)
+            for n in range(epoch, 0, -1):
+                webpage.add_header('epoch [%d]' % n)
+                ims, txts, links = [], [], []
+
+                for label, image_numpy in visuals.items():
+                    image_numpy = util.tensor2im(image)
+                    img_path = 'epoch%.3d_%s.png' % (n, label)
+                    ims.append(img_path)
+                    txts.append(label)
+                    links.append(img_path)
+                webpage.add_images(ims, txts, links, width=self.win_size)
+            webpage.save()
+    
+    def save_current_results1(self, visuals, epoch, epoch_iter):
+        if not os.path.exists(self.img_dir+'/detailed'):
+            os.mkdir(self.img_dir+'/detailed')
+        for label, image in visuals.items():
+            image_numpy = util.tensor2im(image)
+            img_path = os.path.join(self.img_dir, 'detailed', 'epoch%.3d_%.3d_%s.png' % (epoch, epoch_iter, label))
+            util.save_image(image_numpy, img_path)
+
+    # losses: dictionary of error labels and values
+    def plot_current_losses(self, epoch, counter_ratio, opt, losses):
+        if not hasattr(self, 'plot_data'):
+            self.plot_data = {'X': [], 'Y': [], 'legend': list(losses.keys())}
+        self.plot_data['X'].append(epoch + counter_ratio)
+        self.plot_data['Y'].append([losses[k] for k in self.plot_data['legend']])
+        try:
+            self.vis.line(
+                X=np.stack([np.array(self.plot_data['X'])] * len(self.plot_data['legend']), 1),
+                Y=np.array(self.plot_data['Y']),
+                opts={
+                    'title': self.name + ' loss over time',
+                    'legend': self.plot_data['legend'],
+                    'xlabel': 'epoch',
+                    'ylabel': 'loss'},
+                win=self.display_id)
+        except ConnectionError:
+            self.throw_visdom_connection_error()
+
+    # losses: same format as |losses| of plot_current_losses
+    def print_current_losses(self, epoch, i, losses, t, t_data):
+        message = '(epoch: %d, iters: %d, time: %.3f, data: %.3f) ' % (epoch, i, t, t_data)
+        for k, v in losses.items():
+            message += '%s: %.6f ' % (k, v)
+
+        print(message)
+        with open(self.log_name, "a") as log_file:
+            log_file.write('%s\n' % message)

README.md

@@ -1,4 +1,5 @@
 ---
+python_version: 3.7
 title: Apdrawing
 emoji: 💻
 colorFrom: indigo

app.py

@@ -0,0 +1,210 @@
+#!/usr/bin/env python
+
+from __future__ import annotations
+import argparse
+import functools
+import os
+import pathlib
+import sys
+from typing import Callable
+import uuid
+
+sys.path.insert(0, 'APDrawingGAN2')
+
+import gradio as gr
+import huggingface_hub
+import numpy as np
+import PIL.Image
+
+from io import BytesIO
+import shutil
+
+from options.test_options import TestOptions
+from data import CreateDataLoader
+from models import create_model
+
+from util import html
+
+import ntpath
+from util import util
+
+
+ORIGINAL_REPO_URL = 'https://github.com/yiranran/APDrawingGAN2'
+TITLE = 'yiranran/APDrawingGAN2'
+DESCRIPTION = f"""This is a demo for {ORIGINAL_REPO_URL}.
+
+"""
+ARTICLE = """
+
+"""
+
+
+MODEL_REPO = 'hylee/apdrawing_model'
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--device', type=str, default='cpu')
+    parser.add_argument('--theme', type=str)
+    parser.add_argument('--live', action='store_true')
+    parser.add_argument('--share', action='store_true')
+    parser.add_argument('--port', type=int)
+    parser.add_argument('--disable-queue',
+                        dest='enable_queue',
+                        action='store_false')
+    parser.add_argument('--allow-flagging', type=str, default='never')
+    parser.add_argument('--allow-screenshot', action='store_true')
+    return parser.parse_args()
+
+
+def load_checkpoint():
+    dir = 'checkpoint'
+    checkpoint_path = huggingface_hub.hf_hub_download(MODEL_REPO,
+                                                'checkpoints.zip',
+                                                force_filename='checkpoints.zip')
+    print(checkpoint_path)
+    shutil.unpack_archive(checkpoint_path, extract_dir=dir)
+
+    print(os.listdir(dir+'/checkpoints'))
+
+    return dir+'/checkpoints'
+
+# save image to the disk
+def save_images2(image_dir, visuals, image_path, aspect_ratio=1.0, width=256):
+    short_path = ntpath.basename(image_path[0])
+    name = os.path.splitext(short_path)[0]
+
+    imgs = []
+
+    for label, im_data in visuals.items():
+        im = util.tensor2im(im_data)#tensor to numpy array [-1,1]->[0,1]->[0,255]
+        image_name = '%s_%s.png' % (name, label)
+        save_path = os.path.join(image_dir, image_name)
+        h, w, _ = im.shape
+        if aspect_ratio > 1.0:
+            im = np.array(PIL.Image.fromarray(arr).resize(im, (h, int(w * aspect_ratio))))
+        if aspect_ratio < 1.0:
+            im = np.array(PIL.Image.fromarray(arr).resize(im, (int(h / aspect_ratio), w)))
+        util.save_image(im, save_path)
+        imgs.append(save_path)
+
+    return imgs
+
+
+SAFEHASH = [x for x in "0123456789-abcdefghijklmnopqrstuvwxyz_ABCDEFGHIJKLMNOPQRSTUVWXYZ"]
+def compress_UUID():
+    '''
+    根据http://www.ietf.org/rfc/rfc1738.txt,由uuid编码扩bai大字符域生成du串
+    包括:[0-9a-zA-Z\-_]共64个
+    长度:(32-2)/3*2=20
+    备注:可在地球上人zhi人都用,使用100年不重复(2^120)
+    :return:String
+    '''
+    row = str(uuid.uuid4()).replace('-', '')
+    safe_code = ''
+    for i in range(10):
+        enbin = "%012d" % int(bin(int(row[i * 3] + row[i * 3 + 1] + row[i * 3 + 2], 16))[2:], 10)
+        safe_code += (SAFEHASH[int(enbin[0:6], 2)] + SAFEHASH[int(enbin[6:12], 2)])
+    safe_code = safe_code.replace('-', '')
+    return safe_code
+
+
+def run(
+    image,
+    model,
+    opt,
+) -> tuple[PIL.Image.Image]:
+
+    dataroot = 'images/'+compress_UUID()
+    opt.dataroot = os.path.join(dataroot, 'src/')
+    os.makedirs(opt.dataroot, exist_ok=True)
+    opt.results_dir = os.path.join(dataroot, 'results/')
+    os.makedirs(opt.results_dir, exist_ok=True)
+
+    shutil.copy(image.name, opt.dataroot)
+
+    data_loader = CreateDataLoader(opt)
+    dataset = data_loader.load_data()
+
+    imgs = [image.name]
+    # test
+    # model.eval()
+    for i, data in enumerate(dataset):
+        if i >= opt.how_many:  # test code only supports batch_size = 1, how_many means how many test images to run
+            break
+        model.set_input(data)
+        model.test()
+        visuals = model.get_current_visuals()  # in test the loadSize is set to the same as fineSize
+        img_path = model.get_image_paths()
+        # if i % 5 == 0:
+        #    print('processing (%04d)-th image... %s' % (i, img_path))
+        imgs = save_images2(opt.results_dir, visuals, img_path, aspect_ratio=opt.aspect_ratio, width=opt.display_winsize)
+
+    print(imgs)
+    return PIL.Image.open(imgs[0])
+
+
+def main():
+    gr.close_all()
+
+    args = parse_args()
+
+    checkpoint_dir = load_checkpoint()
+
+    opt = TestOptions().parse()
+    opt.num_threads = 1  # test code only supports num_threads = 1
+    opt.batch_size = 1  # test code only supports batch_size = 1
+    opt.serial_batches = True  # no shuffle
+    opt.no_flip = True  # no flip
+    opt.display_id = -1  # no visdom display
+
+    '''
+       python test.py --dataroot dataset/test_single --name apdrawinggan++_author --model test --use_resnet --netG resnet_9blocks --which_epoch 150 --how_many 1000 --gpu_ids 0 --gpu_ids_p 0 --imagefolder images-single
+       '''
+    opt.dataroot = 'dataset/test_single'
+    opt.name = 'apdrawinggan++_author'
+    opt.model = 'test'
+    opt.use_resnet = True
+    opt.netG = 'resnet_9blocks'
+    opt.which_epoch = 150
+    opt.how_many = 1000
+    opt.gpu_ids = -1
+    opt.gpu_ids_p = -1
+    opt.imagefolder = 'images-single'
+
+    opt.checkpoints_dir = checkpoint_dir
+
+
+    model = create_model(opt)
+    model.setup(opt)
+
+    func = functools.partial(run, model=model, opt=opt)
+    func = functools.update_wrapper(func, run)
+
+    
+    gr.Interface(
+        func,
+        [
+            gr.inputs.Image(type='file', label='Input Image'),
+        ],
+        [
+            gr.outputs.Image(
+                type='pil',
+                label='Result'),
+        ],
+        #examples=examples,
+        theme=args.theme,
+        title=TITLE,
+        description=DESCRIPTION,
+        article=ARTICLE,
+        allow_screenshot=args.allow_screenshot,
+        allow_flagging=args.allow_flagging,
+        live=args.live,
+    ).launch(
+        enable_queue=args.enable_queue,
+        server_port=args.port,
+        share=args.share,
+    )
+
+
+if __name__ == '__main__':
+    main()

packages.txt

@@ -0,0 +1,2 @@
+
+

requirements.txt

@@ -0,0 +1,8 @@
+torch>=0.4.0
+torchvision>=0.2.1
+dominate>=2.3.1
+visdom>=0.1.8.3
+scipy>=1.1.0
+numpy>=1.14.1
+Pillow>=5.0.0
+opencv-python>=3.4.2