init

APDrawingGAN2/data/aligned_dataset.py

@@ -9,50 +9,55 @@ 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
+    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 = (ts.numpy() / 2 + 0.5) * 255
     img = img.astype('uint8')
-    img = np.transpose(img,(1,2,0))
-    img = img[:,:,::-1]#rgb->bgr
+    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()
+    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()
+
+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))
+        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
+    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):
@@ -71,17 +76,17 @@ class AlignedDataset(BaseDataset):
         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')
+        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:
+        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
+        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)
@@ -90,7 +95,7 @@ class AlignedDataset(BaseDataset):
         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
+        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)
@@ -118,25 +123,25 @@ class AlignedDataset(BaseDataset):
         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]:
+            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)
+                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)
+                        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
@@ -144,32 +149,37 @@ class AlignedDataset(BaseDataset):
                 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]])
+                center = torch.LongTensor(
+                    [[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)]
+                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)
+                        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]
+                        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]
+                    item[regions[i] + '_A'] = A[:, int(center[i, 1] - rhs[i] / 2):int(center[i, 1] + rhs[i] / 2),
+                                              int(center[i, 0] - rws[i] / 2):int(center[i, 0] + rws[i] / 2)]
+                    item[regions[i] + '_B'] = B[:, int(center[i, 1] - rhs[i] / 2):int(center[i, 1] + rhs[i] / 2),
+                                              int(center[i, 0] - rws[i] / 2):int(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]
+                        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')
+                    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')
+                        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:
@@ -180,11 +190,12 @@ class AlignedDataset(BaseDataset):
                     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
+                    if self.opt.region_enm in [0, 1]:
+                        cmask = cmask[:, int(center[i, 1] - rhs[i] / 2):int(center[i, 1] + rhs[i] / 2),
+                                int(center[i, 0] - rws[i] / 2):int(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]
@@ -194,70 +205,75 @@ class AlignedDataset(BaseDataset):
                 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
+                        mask[:, int(center[i, 1] - rhs[i] / 2):int(center[i, 1] + rhs[i] / 2),
+                        int(center[i, 0] - rws[i] / 2):int(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 = [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]
+                        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]
+                            xbi = [int(center[i, 0] - rws[i] / 2), int(center[i, 0] + rws[i] / 2 - 1)]
+                            ybi = [int(center[i, 1] - rhs[i] / 2), int(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]
+                                maskx = bound[:, xbi[j]]
+                                masky = bound[ybi[j], :]
+                                tmp_a = torch.from_numpy(maskx) * xbi[j]
+                                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)
+                                tmp_a = torch.from_numpy(masky) * ybi[j]
+                                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
+                        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
+                        rec[:, int(center[i, 1] - rhs[i] / 2):int(center[i, 1] + rhs[i] / 2),
+                        int(center[i, 0] - rws[i] / 2):int(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
+                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
+                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')
+                    bgpath = os.path.join(bgdir, basen[:-4] + '.png')
                     im_bg = Image.open(bgpath)
-                    mask2 = transforms.ToTensor()(im_bg) # mask out background
+                    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
+                    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)
@@ -270,11 +286,11 @@ class AlignedDataset(BaseDataset):
             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_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]
+            face_mask = transforms.ToTensor()(face_mask)  # [0,1]
             if flipped:
                 face_mask = face_mask.index_select(2, idx)
             item['face_mask'] = face_mask

APDrawingGAN2/data/single_dataset.py

@@ -82,7 +82,7 @@ class SingleDataset(BaseDataset):
             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]])
+            center = torch.LongTensor([[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)]
@@ -95,7 +95,10 @@ class SingleDataset(BaseDataset):
                     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[:,int(center[i,1]-rhs[i]/2):int(center[i,1]+rhs[i]/2),int(center[i,0]-rws[i]/2):int(center[i,0]+rws[i]/2)]
+                item[regions[i]+'_A'] = A[:,(center[i,1]-rhs[i]/2).to(torch.long):
+                                            (center[i,1]+rhs[i]/2).to(torch.long),
+                                        (center[i,0]-rws[i]/2).to(torch.long):
+                                        (center[i,0]+rws[i]/2).to(torch.long)]
                 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:
@@ -111,7 +114,7 @@ class SingleDataset(BaseDataset):
                     cmask0 = (cmask0 >= 0.5).float()
                     cmasks0.append(cmask0)
                     cmask = cmask0.clone()
-                    cmask = cmask[:,int(center[i,1]-rhs[i]/2):int(center[i,1]+rhs[i]/2),int(center[i,0]-rws[i]/2):int(center[i,0]+rws[i]/2)]
+                    cmask = cmask[:,(center[i,1]-rhs[i]/2).to(torch.long):(center[i,1]+rhs[i]/2).to(torch.long),(center[i,0]-rws[i]/2).to(torch.long):(center[i,0]+rws[i]/2).to(torch.long)]
                     cmasks.append(cmask)
                 item['cmaskel'] = cmasks[0]
                 item['cmasker'] = cmasks[1]
@@ -121,7 +124,7 @@ class SingleDataset(BaseDataset):
                 output_nc = self.opt.output_nc
                 mask = torch.ones([output_nc,A.shape[1],A.shape[2]])
                 for i in range(4):
-                    mask[:,int(center[i,1]-rhs[i]/2):int(center[i,1]+rhs[i]/2),int(center[i,0]-rws[i]/2):int(center[i,0]+rws[i]/2)] = 0
+                    mask[:,(center[i,1]-rhs[i]/2).to(torch.long):(center[i,1]+rhs[i]/2).to(torch.long),(center[i,0]-rws[i]/2).to(torch.long):(center[i,0]+rws[i]/2).to(torch.long)] = 0
                 if self.opt.soft_border:
                     imgsize = self.opt.fineSize
                     maskn = mask[0].numpy()
@@ -136,11 +139,11 @@ class SingleDataset(BaseDataset):
                     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]
+                        xbi = [(center[i,0]-rws[i]/2).to(torch.long), (center[i,0]+rws[i]/2-1).to(torch.long)]
+                        ybi = [(center[i,1]-rhs[i]/2).to(torch.long), (center[i,1]+rhs[i]/2-1).to(torch.long)]
                         for j in range(2):
-                            maskx = bound[:,int(xbi[j])]
-                            masky = bound[int(ybi[j]),:]
+                            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]
@@ -160,10 +163,8 @@ class SingleDataset(BaseDataset):
                     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
-                    hair_A = (A/2+0.5) * mask.repeat(3,1,1) * mask2.repeat(3,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
-                    bg_A = (A/2+0.5) * (torch.ones(mask2.shape)-mask2).repeat(3,1,1) * 2 - 1
+                    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

APDrawingGAN2/models/apdrawingpp_style_model.py

@@ -6,13 +6,13 @@ import os
 import math
 
 W = 11
-aa = int(math.floor(512./W))
-res = 512 - W*aa
+aa = int(math.floor(512. / W))
+res = 512 - W * aa
 
 
-def padpart(A,part,centers,opt,device):
+def padpart(A, part, centers, opt, device):
     IMAGE_SIZE = opt.fineSize
-    bs,nc,_,_ = A.shape
+    bs, nc, _, _ = A.shape
     ratio = IMAGE_SIZE / 256
     NOSE_W = opt.NOSE_W * ratio
     NOSE_H = opt.NOSE_H * ratio
@@ -20,37 +20,52 @@ def padpart(A,part,centers,opt,device):
     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
+    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])
+            A_p[i] = torch.nn.ConstantPad2d((
+                                            int(center[2, 0] - NOSE_W / 2), IMAGE_SIZE - int(center[2, 0] + NOSE_W / 2),
+                                            int(center[2, 1] - NOSE_H / 2),
+                                            IMAGE_SIZE - int(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])
+            A_p[i] = torch.nn.ConstantPad2d((int(center[0, 0] - EYE_W / 2), IMAGE_SIZE - int(center[0, 0] + EYE_W / 2),
+                                             int(center[0, 1] - EYE_H / 2), IMAGE_SIZE - int(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])
+            A_p[i] = torch.nn.ConstantPad2d((int(center[1, 0] - EYE_W / 2), IMAGE_SIZE - int(center[1, 0] + EYE_W / 2),
+                                             int(center[1, 1] - EYE_H / 2), IMAGE_SIZE - int(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])
+            A_p[i] = torch.nn.ConstantPad2d((int(center[3, 0] - MOUTH_W / 2),
+                                             IMAGE_SIZE - int(center[3, 0] + MOUTH_W / 2),
+                                             int(center[3, 1] - MOUTH_H / 2),
+                                             IMAGE_SIZE - int(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]
+
+
+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)
+        nldt = torch.atan(dt * xmax) / torch.atan(xmax)
     elif type == 'sigmoid':
-        nldt = (torch.sigmoid(dt*xmax)-0.5) / (torch.sigmoid(xmax)-0.5)
+        nldt = (torch.sigmoid(dt * xmax) - 0.5) / (torch.sigmoid(xmax) - 0.5)
     elif type == 'tanh':
-        nldt = torch.tanh(dt*xmax) / torch.tanh(xmax)
+        nldt = torch.tanh(dt * xmax) / torch.tanh(xmax)
     elif type == 'pow':
-        nldt = torch.pow(dt*xmax,2) / torch.pow(xmax,2)
+        nldt = torch.pow(dt * xmax, 2) / torch.pow(xmax, 2)
     elif type == 'exp':
-        if xmax.item()>1:
+        if xmax.item() > 1:
             xmax = xmax / 3
-        nldt = (torch.exp(dt*xmax)-1) / (torch.exp(xmax)-1)
-    #print("remap dt:", type, xmax.item())
+        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'
@@ -60,7 +75,7 @@ class APDrawingPPStyleModel(BaseModel):
 
         # 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(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)
@@ -107,15 +122,15 @@ class APDrawingPPStyleModel(BaseModel):
             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.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']
+                    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']
+                    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']
+                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']
@@ -129,7 +144,7 @@ class APDrawingPPStyleModel(BaseModel):
         if self.isTrain:
             self.model_names = ['G', 'D']
             if self.opt.discriminator_local:
-                self.model_names += ['DLEyel','DLEyer','DLNose','DLMouth','DLHair','DLBG']
+                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']
@@ -141,13 +156,13 @@ class APDrawingPPStyleModel(BaseModel):
             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']
+            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']
+                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
@@ -159,55 +174,61 @@ class APDrawingPPStyleModel(BaseModel):
         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)
+                                          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)
+                                                   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)
+                                                   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)
+                                                   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)
+                                                    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)
+                                                   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)
-                
-        
+                                                 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)
+                                               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)
+                                               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)
+                                               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)
+                                                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)
+                                               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)
+                                             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)
+            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)
+                                            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)
-        
+                                            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)
@@ -220,34 +241,39 @@ class APDrawingPPStyleModel(BaseModel):
             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))
+                                                    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())
+                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))
-            
+                                                    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())
+                                                    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))
+                                                    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)            
+                                           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
@@ -255,53 +281,51 @@ class APDrawingPPStyleModel(BaseModel):
             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)
+                                             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)
+                                             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)
+                                                  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)
+                                                  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)
+                                                  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)
+                                            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)
+                                            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)
+                                              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)
+                                              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:
+        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)
+                                                  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'
@@ -318,7 +342,7 @@ class APDrawingPPStyleModel(BaseModel):
             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]:
+            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)
@@ -327,17 +351,17 @@ class APDrawingPPStyleModel(BaseModel):
                 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.cmask1 = self.cmask * 2 - 1  # [0,1]->[-1,1]
                 self.cmaskel = input['cmaskel'].to(self.device)
-                self.cmask1el = self.cmaskel*2-1
+                self.cmask1el = self.cmaskel * 2 - 1
                 self.cmasker = input['cmasker'].to(self.device)
-                self.cmask1er = self.cmasker*2-1
+                self.cmask1er = self.cmasker * 2 - 1
                 self.cmaskmo = input['cmaskmo'].to(self.device)
-                self.cmask1mo = self.cmaskmo*2-1
+                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.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):
@@ -345,13 +369,13 @@ class APDrawingPPStyleModel(BaseModel):
             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)
+
+    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)
+        onehot.scatter_(1, y, 1)
         return onehot
 
     def forward(self):
@@ -361,40 +385,41 @@ class APDrawingPPStyleModel(BaseModel):
             self.fake_B0 = self.netG(self.real_A)
             # EYES, MOUTH
             outputs1 = self.netCLm(self.real_A_mouth)
-            onehot1 = self.getonehot(outputs1,2)
+            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!
+            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)
+                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)
+                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')
+                        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)
+                        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!
+            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)
+                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]:
+            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
@@ -405,41 +430,48 @@ class APDrawingPPStyleModel(BaseModel):
                     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)
-            
+            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)
+            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')
+                    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])
+                    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_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))
+            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)
+                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))
+                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.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)
@@ -458,21 +490,20 @@ class APDrawingPPStyleModel(BaseModel):
                     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
+        # 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']
+            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])
@@ -487,7 +518,7 @@ class APDrawingPPStyleModel(BaseModel):
         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']
+            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])
@@ -504,12 +535,12 @@ class APDrawingPPStyleModel(BaseModel):
     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)
+        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
+            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])
@@ -524,31 +555,34 @@ class APDrawingPPStyleModel(BaseModel):
         # 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']
+            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
+                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
+                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
+                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:
@@ -558,22 +592,23 @@ class APDrawingPPStyleModel(BaseModel):
             # 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
+            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())
-            
+                # 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
+            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) 
+                self.loss_G_chamfer = self.loss_G_chamfer.cuda(gpu)
 
-            # d_CM(G(p_i),a_i)
+                # d_CM(G(p_i),a_i)
             if gpu_p != gpu:
                 dt1gt = self.dt1gt.cuda(gpu_p)
                 dt2gt = self.dt2gt.cuda(gpu_p)
@@ -583,13 +618,14 @@ class APDrawingPPStyleModel(BaseModel):
             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
+                # 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
+            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)
 
@@ -599,11 +635,10 @@ class APDrawingPPStyleModel(BaseModel):
             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
+                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_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:
@@ -627,7 +662,7 @@ class APDrawingPPStyleModel(BaseModel):
         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)
@@ -661,32 +696,32 @@ class APDrawingPPStyleModel(BaseModel):
         patches = []
         if self.isTrain and self.opt.emphasis_conti_face:
             weights = []
-            W2 = int(W/2)
-        t = np.random.randint(res,size=2)
+            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:
+                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])
+                    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
-    
+            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)
+        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:  
+                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

@@ -20,8 +20,8 @@ class BaseModel():
         self.gpu_ids = opt.gpu_ids
         self.gpu_ids_p = opt.gpu_ids_p
         self.isTrain = opt.isTrain
-        self.device = torch.device('cpu')
-        self.device_p = torch.device('cpu')
+        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':
@@ -105,7 +105,7 @@ class BaseModel():
                     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))
@@ -120,7 +120,7 @@ class BaseModel():
                     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':
@@ -142,7 +142,7 @@ class BaseModel():
                 if getattr(module, key) is None:
                     state_dict.pop('.'.join(keys))
             if module.__class__.__name__.startswith('InstanceNorm') and \
-               (key == 'num_batches_tracked'):
+                    (key == 'num_batches_tracked'):
                 state_dict.pop('.'.join(keys))
         else:
             self.__patch_instance_norm_state_dict(state_dict, getattr(module, key), keys, i + 1)
@@ -171,7 +171,7 @@ class BaseModel():
                 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))
@@ -184,19 +184,19 @@ class BaseModel():
                 if isinstance(net, torch.nn.DataParallel):
                     net = net.module
                 if name[0] == 'G':
-                    print('loading the model %s from %s' % (name,gen_name))
+                    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))
+                    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:
@@ -214,7 +214,8 @@ class BaseModel():
                 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']:
+                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))
@@ -251,18 +252,19 @@ class BaseModel():
 
     # =============================================================================================================
     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):
+        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
@@ -276,7 +278,7 @@ class BaseModel():
         if comb_op == 0:
             # use max pooling, pad black for eyes etc
             padvalue = -1
-            if region_enm in [1,2]:
+            if region_enm in [1, 2]:
                 eyel = eyel * cmaskel
                 eyer = eyer * cmasker
                 nose = nose * cmaskno
@@ -284,12 +286,12 @@ class BaseModel():
         else:
             # use min pooling, pad white for eyes etc
             padvalue = 1
-            if region_enm in [1,2]:
+            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
+        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
@@ -298,20 +300,32 @@ class BaseModel():
             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)
+            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])
+                    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
@@ -328,13 +342,14 @@ class BaseModel():
             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):
+
+    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]:
+            if region_enm in [1, 2]:
                 eyel = eyel * cmaskel
                 eyer = eyer * cmasker
                 nose = nose * cmaskno
@@ -343,12 +358,12 @@ class BaseModel():
             # use min pooling, pad white for eyes etc
             padvalue = 1
             hair = self.addone_with_mask(hair, mask)
-            if region_enm in [1,2]:
+            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
+        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
@@ -357,17 +372,26 @@ class BaseModel():
             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)
+            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])
+                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
@@ -378,22 +402,23 @@ class BaseModel():
             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)
+            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)
+            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):
+
+    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]:
+            if region_enm in [1, 2]:
                 eyel = eyel * cmaskel
                 eyer = eyer * cmasker
                 nose = nose * cmaskno
@@ -403,12 +428,12 @@ class BaseModel():
             padvalue = 1
             hair = self.addone_with_mask(hair, maskh)
             bg = self.addone_with_mask(bg, maskb)
-            if region_enm in [1,2]:
+            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
+        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
@@ -417,17 +442,29 @@ class BaseModel():
             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)
+            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])
+                center = self.center[i]  # x,y
+                eyel_p[i] = torch.nn.ConstantPad2d((int(center[0, 0] - EYE_W / 2),
+                                                    IMAGE_SIZE - int(center[0, 0] + EYE_W / 2),
+                                                    int(center[0, 1] - EYE_H / 2),
+                                                    IMAGE_SIZE - int(center[0, 1] + EYE_H / 2)), padvalue)(eyel[i])
+                eyer_p[i] = torch.nn.ConstantPad2d((int(center[1, 0] - EYE_W / 2),
+                                                    IMAGE_SIZE - int(center[1, 0] + EYE_W / 2),
+                                                    int(center[1, 1] - EYE_H / 2),
+                                                    IMAGE_SIZE - int(center[1, 1] + EYE_H / 2)), padvalue)(eyer[i])
+                nose_p[i] = torch.nn.ConstantPad2d((int(center[2, 0] - NOSE_W / 2),
+                                                    IMAGE_SIZE - int(center[2, 0] + NOSE_W / 2),
+                                                    int(center[2, 1] - NOSE_H / 2),
+                                                    IMAGE_SIZE - int(center[2, 1] + NOSE_H / 2)), padvalue)(nose[i])
+                mouth_p[i] = torch.nn.ConstantPad2d((int(center[3, 0] - MOUTH_W / 2),
+                                                     IMAGE_SIZE - int(center[3, 0] + MOUTH_W / 2),
+                                                     int(center[3, 1] - MOUTH_H / 2),
+                                                     IMAGE_SIZE - int(center[3, 1] + MOUTH_H / 2)), padvalue)(mouth[i])
         elif region_enm in [2]:
             eyel_p = eyel
             eyer_p = eyer
@@ -437,17 +474,17 @@ class BaseModel():
             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)
+            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)
+            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):
+
+    def partCombiner3(self, face, hair, maskf, maskh, comb_op=1):
         if comb_op == 0:
             # use max pooling, pad black etc
             padvalue = -1
@@ -459,27 +496,25 @@ class BaseModel():
             face = self.addone_with_mask(face, maskf)
             hair = self.addone_with_mask(hair, maskh)
         if comb_op == 0:
-            result = torch.max(face,hair)
+            result = torch.max(face, hair)
         else:
-            result = torch.min(face,hair)
+            result = torch.min(face, hair)
         return result
 
-
     def tocv2(ts):
-        img = (ts.numpy()/2+0.5)*255
+        img = (ts.numpy() / 2 + 0.5) * 255
         img = img.astype('uint8')
-        img = np.transpose(img,(1,2,0))
-        img = img[:,:,::-1]#rgb->bgr
+        img = np.transpose(img, (1, 2, 0))
+        img = img[:, :, ::-1]  # rgb->bgr
         return img
-    
+
     def totor(img):
-        img = img[:,:,::-1]
+        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):
+
+    def ContinuityForTest(self, real=0):
         # Patch-based
         self.get_patches()
         self.outputs = self.netRegressor(self.fake_B_patches)
@@ -494,16 +529,17 @@ class BaseModel():
             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')
+            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
+
+    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]:
+        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
@@ -511,28 +547,32 @@ class BaseModel():
             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)
+            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]
+                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
+            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):
+
+    def getaddw(self, local_name):
         addw = 1
-        if local_name in ['DLEyel','DLEyer','eyel','eyer','DLFace','face']:
+        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