init

APDrawingGAN2/data/single_dataset.py

@@ -160,8 +160,10 @@ 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
-                    bg_A = (A/2+0.5) * (torch.ones(mask2.shape)-mask2).repeat(int(input_nc/output_nc),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
+					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
                     item['hair_A'] = hair_A
                     item['bg_A'] = bg_A
                     item['mask'] = mask