add dilation and erosion

README.md

@@ -17,24 +17,47 @@
 ## Description
 Official Implementation of Barbershop. **KEEP UPDATING !**
 
+## Updates
+**27/12/2021** Add dilation and erosion parameters to smooth the boundary.
+#### 24/12/2021 Important Update: Add improved semantic mask inpainting module. Please git pull the newest version.
 
-Mode1: Produce realistic results:
-```
-python main.py --im_path1 90.png --im_path2 15.png --im_path3 117.png --sign realistic
+**18/12/2021** Add a rough version of the project.
+
+**2/6/2021** Add project page.
+
+
+## Installation
+- Clone the repository:
+``` 
+git clone https://github.com/ZPdesu/Barbershop.git
+cd Barbershop
 ```
+- Dependencies:  
+We recommend running this repository using [Anaconda](https://docs.anaconda.com/anaconda/install/). 
+All dependencies for defining the environment are provided in `environment/environment.yaml`.
+
+
+## Download II2S images
+Please download the [II2S](https://drive.google.com/drive/folders/15jsR9yy_pfDHiS9aE3HcYDgwtBbAneId?usp=sharing) 
+and put them in the `input/face` folder.
 
-Mode2: Produce results faithful to the masks:
+
+## Getting Started  
+Preprocess your own images. Please put the raw images in the `unprocessed` folder.
 ```
-python main.py --im_path1 90.png --im_path2 15.png --im_path3 117.png --sign fidelity
+python align_face.py
 ```
 
+Produce realistic results:
+```
+python main.py --im_path1 90.png --im_path2 15.png --im_path3 117.png --sign realistic --smooth 5
+```
 
-## Updates
-#### 24/12/2021 Important Update: Add improved semantic mask inpainting module. Please git pull the newest version.
-
-**18/12/2021** Add a rough version of the project.
+Produce results faithful to the masks:
+```
+python main.py --im_path1 90.png --im_path2 15.png --im_path3 117.png --sign fidelity --smooth 5
+```
 
-**2/6/2021** Add project page.
 
 
 ## Todo List
@@ -44,6 +67,8 @@ python main.py --im_path1 90.png --im_path2 15.png --im_path3 117.png --sign fid
 * add preprocessing step
 * ...
 
+## Acknowledgments
+This code borrows heavily from [II2S](https://github.com/ZPdesu/II2S).
 
 ## BibTeX
 

bash.sh

@@ -1,5 +1,5 @@
 #!/usr/bin/env bash
 
-python main.py --im_path1 90.png --im_path2 15.png --im_path3 117.png --sign realistic
+python main.py --im_path1 16.png --im_path2 15.png --im_path3 117.png --sign realistic --smooth 5
 
 

main.py

@@ -40,9 +40,9 @@ def main(args):
     ii2s.invert_images_in_FS([*im_set])
 
     align = Alignment(args)
-    align.align_images(im_path1, im_path2, sign=args.sign, align_more_region=False)
+    align.align_images(im_path1, im_path2, sign=args.sign, align_more_region=False, smooth=args.smooth)
     if im_path2 != im_path3:
-        align.align_images(im_path1, im_path3, sign=args.sign, align_more_region=False, save_intermediate=False)
+        align.align_images(im_path1, im_path3, sign=args.sign, align_more_region=False, smooth=args.smooth, save_intermediate=False)
 
     blend = Blending(args)
     blend.blend_images(im_path1, im_path2, im_path3, sign=args.sign)
@@ -61,11 +61,11 @@ def main(args):
                         help='The directory of the images to be inverted')
     parser.add_argument('--output_dir', type=str, default='output',
                         help='The directory to save the latent codes and inversion images')
-    parser.add_argument('--im_path1', type=str, default='90.png', help='Identity image')
+    parser.add_argument('--im_path1', type=str, default='16.png', help='Identity image')
     parser.add_argument('--im_path2', type=str, default='15.png', help='Structure image')
     parser.add_argument('--im_path3', type=str, default='117.png', help='Appearance image')
     parser.add_argument('--sign', type=str, default='realistic', help='realistic or fidelity results')
-
+    parser.add_argument('--smooth', type=int, default=5, help='dilation and erosion parameter')
 
     # StyleGAN2 setting
     parser.add_argument('--size', type=int, default=1024)

models/Alignment.py

@@ -17,6 +17,8 @@
 from utils.data_utils import load_FS_latent
 from utils.seg_utils import save_vis_mask
 from utils.model_utils import download_weight
+from utils.data_utils import cuda_unsqueeze
+from utils.image_utils import dilate_erosion_mask_tensor
 
 toPIL = torchvision.transforms.ToPILImage()
 
@@ -210,7 +212,13 @@ def create_down_seg(self, latent_in):
         down_seg, _, _ = self.seg(im)
         return down_seg, gen_im
 
-    def align_images(self, img_path1, img_path2, sign='realistic', align_more_region=True,
+
+    def dilate_erosion(self, free_mask, device, dilate_erosion=5):
+        free_mask = F.interpolate(free_mask.cpu(), size=(256, 256), mode='nearest').squeeze()
+        free_mask_D, free_mask_E = cuda_unsqueeze(dilate_erosion_mask_tensor(free_mask, dilate_erosion=dilate_erosion), device)
+        return free_mask_D, free_mask_E
+
+    def align_images(self, img_path1, img_path2, sign='realistic', align_more_region=False, smooth=5,
                      save_intermediate=True):
 
         ################## img_path1: Identity Image
@@ -310,6 +318,11 @@ def align_images(self, img_path1, img_path2, sign='realistic', align_more_region
         latent_F_out_new = latent_F_out_new.clone().detach()
 
         free_mask = 1 - (1 - hair_mask1.unsqueeze(0)) * (1 - hair_mask_target)
+
+        ##############################
+        free_mask, _ = self.dilate_erosion(free_mask, device, dilate_erosion=smooth)
+        ##############################
+
         free_mask_down_32 = F.interpolate(free_mask.float(), size=(32, 32), mode='bicubic')[0]
         interpolation_low = 1 - free_mask_down_32
 
@@ -319,14 +332,20 @@ def align_images(self, img_path1, img_path2, sign='realistic', align_more_region
 
         if not align_more_region:
             free_mask = hair_mask_target
+            ##########################
+            _, free_mask = self.dilate_erosion(free_mask, device, dilate_erosion=smooth)
+            ##########################
             free_mask_down_32 = F.interpolate(free_mask.float(), size=(32, 32), mode='bicubic')[0]
             interpolation_low = 1 - free_mask_down_32
 
 
         latent_F_mixed = latent_F_out_new + interpolation_low.unsqueeze(0) * (
                 latent_F_mixed - latent_F_out_new)
 
-        free_mask = F.interpolate((hair_mask2.unsqueeze(0) * free_mask).float(), size=(256, 256), mode='nearest').cuda()
+        free_mask = F.interpolate((hair_mask2.unsqueeze(0) * hair_mask_target).float(), size=(256, 256), mode='nearest').cuda()
+        ##########################
+        _, free_mask = self.dilate_erosion(free_mask, device, dilate_erosion=smooth)
+        ##########################
         free_mask_down_32 = F.interpolate(free_mask.float(), size=(32, 32), mode='bicubic')[0]
         interpolation_low = 1 - free_mask_down_32