fix nvidia 3000 series usage

Add Docker environment & web demo

.vscode/launch.json

@@ -0,0 +1,15 @@
+{
+    // Use IntelliSense to learn about possible attributes.
+    // Hover to view descriptions of existing attributes.
+    // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "type": "pwa-chrome",
+            "request": "launch",
+            "name": "Launch Chrome against localhost",
+            "url": "http://localhost:8080",
+            "webRoot": "${workspaceFolder}"
+        }
+    ]
+}

README.md

@@ -16,8 +16,18 @@ Our paper can be downloaded from [[Arxiv]](https://arxiv.org/pdf/2106.06340v1.pd
 ## Attention
 ***This project is for technical and academic use only. Please do not apply it to illegal and unethical scenarios.***
 
+If you find this project useful, please star it. It is the greatest appreciation of our work.
+
 ## Top News <img width=8% src="./docs/img/new.gif"/>
 
+**`2021-08-30`**: Docker has been supported, please refer [here](https://replicate.ai/neuralchen/simswap-image) for details.
+
+**`2021-08-17`**: We have updated the [Preparation](./docs/guidance/preparation.md), The main change is that the gpu version of onnx is now installed by default, Now the time to process a video is greatly reduced.
+
+**`2021-07-19`**: ***Obvious border abruptness has been resolved***. We add the ability to using mask and upgrade the old algorithm for better visual effect, please go to [Inference for image or video face swapping](./docs/guidance/usage.md) for details. Please don’t forget to go to [Preparation](./docs/guidance/preparation.md) to check the latest set up. (Thanks for the help from [@woctezuma](https://github.com/woctezuma) and [@instant-high](https://github.com/instant-high))
+
+**`2021-07-04`**: A new Colab performing **multi specific** face video swapping has been added. You can check it out [here](https://colab.research.google.com/github/neuralchen/SimSwap/blob/main/MultiSpecific.ipynb) 
+
 **`2021-07-03`**: We add the scripts for **multi specific** face swapping, please go to [Inference for image or video face swapping](./docs/guidance/usage.md) for details. 
 
 **`2021-07-02`**: We add the scripts for designating a **specific** person in arbitrary video or image to change face, please go to [Inference for image or video face swapping](./docs/guidance/usage.md) for details.
@@ -44,6 +54,12 @@ Our paper can be downloaded from [[Arxiv]](https://arxiv.org/pdf/2106.06340v1.pd
 
 [Colab demo](https://colab.research.google.com/github/neuralchen/SimSwap/blob/main/SimSwap%20colab.ipynb)
 
+<div style="background: yellow; width:140px; font-weight:bold;font-family: sans-serif;">Stronger feature</div>
+
+[Colab fo switching specific faces in multi-face videos](https://colab.research.google.com/github/neuralchen/SimSwap/blob/main/MultiSpecific.ipynb) 
+
+[Image face swapping demo & Docker image on Replicate](https://replicate.ai/neuralchen/simswap-image)
+
 Training: **coming soon**
 
 
@@ -55,6 +71,10 @@ Training: **coming soon**
 <img width=24% src="./docs/img/zhoujielun.webp"/>
 <img width=24% src="./docs/img/zhuyin.webp"/>
 </div>
+<div>
+<img width=49% src="./docs/img/mama_mask_short.webp"/>
+<img width=49% src="./docs/img/mama_mask_wuyifan_short.webp"/>
+</div>
 
 ## Results
 ![Results1](/docs/img/results1.PNG)
@@ -70,7 +90,7 @@ Training: **coming soon**
 
 **High-quality videos can be found in the link below:**
 
-[[Mama(video) 1080p]](https://drive.google.com/file/d/1JTruy6BTnT1EK1PSaZ4x-F8RhtZU_kT3/view?usp=sharing)
+[[Mama(video) 1080p]](https://drive.google.com/file/d/1mnSlwzz7f4H2O7UwApAHo64mgK4xSNyK/view?usp=sharing)
 
 [[Google Drive link for video 1]](https://drive.google.com/file/d/1hdne7Gw39d34zt3w1NYV3Ln5cT8PfCNm/view?usp=sharing)
 
@@ -139,5 +159,8 @@ Learn about our other projects
 ## Acknowledgements
 
 <!--ts-->
+* [Deepfacelab](https://github.com/iperov/DeepFaceLab)
 * [Insightface](https://github.com/deepinsight/insightface)
+* [Face-parsing.PyTorch](https://github.com/zllrunning/face-parsing.PyTorch)
+* [BiSeNet](https://github.com/CoinCheung/BiSeNet)
 <!--te-->

SimSwap colab.ipynb

@@ -227,11 +227,15 @@
         "from google_drive_downloader import GoogleDriveDownloader\n",
         "\n",
         "### it seems that google drive link may not be permenant, you can find this ID from our open url.\n",
-        "GoogleDriveDownloader.download_file_from_google_drive(file_id='1TLNdIufzwesDbyr_nVTR7Zrx9oRHLM_N',\n",
-        "                                    dest_path='./arcface_model/arcface_checkpoint.tar')\n",
-        "GoogleDriveDownloader.download_file_from_google_drive(file_id='1PXkRiBUYbu1xWpQyDEJvGKeqqUFthJcI',\n",
-        "                                    dest_path='./checkpoints.zip')\n",
-        "!unzip ./checkpoints.zip  -d ./checkpoints"
+        "# GoogleDriveDownloader.download_file_from_google_drive(file_id='1TLNdIufzwesDbyr_nVTR7Zrx9oRHLM_N',\n",
+        "#                                     dest_path='./arcface_model/arcface_checkpoint.tar')\n",
+        "# GoogleDriveDownloader.download_file_from_google_drive(file_id='1PXkRiBUYbu1xWpQyDEJvGKeqqUFthJcI',\n",
+        "#                                     dest_path='./checkpoints.zip')\n",
+        "\n",
+        "!wget -P ./arcface_model https://github.com/neuralchen/SimSwap/releases/download/1.0/arcface_checkpoint.tar\n",
+        "!wget https://github.com/neuralchen/SimSwap/releases/download/1.0/checkpoints.zip\n",
+        "!unzip ./checkpoints.zip  -d ./checkpoints\n",
+        "!wget -P ./parsing_model/checkpoint https://github.com/neuralchen/SimSwap/releases/download/1.0/79999_iter.pth"
       ],
       "execution_count": 5,
       "outputs": [
@@ -392,6 +396,7 @@
         "opt.temp_path = './tmp'\n",
         "opt.Arc_path = './arcface_model/arcface_checkpoint.tar'\n",
         "opt.isTrain = False\n",
+        "opt.use_mask = True  ## new feature up-to-date\n",
         "\n",
         "crop_size = 224\n",
         "\n",
@@ -399,29 +404,29 @@
         "model = create_model(opt)\n",
         "model.eval()\n",
         "\n",
-        "\n",
         "app = Face_detect_crop(name='antelope', root='./insightface_func/models')\n",
         "app.prepare(ctx_id= 0, det_thresh=0.6, det_size=(640,640))\n",
         "\n",
-        "pic_a = opt.pic_a_path\n",
-        "# img_a = Image.open(pic_a).convert('RGB')\n",
-        "img_a_whole = cv2.imread(pic_a)\n",
-        "img_a_align_crop, _ = app.get(img_a_whole,crop_size)\n",
-        "img_a_align_crop_pil = Image.fromarray(cv2.cvtColor(img_a_align_crop[0],cv2.COLOR_BGR2RGB)) \n",
-        "img_a = transformer_Arcface(img_a_align_crop_pil)\n",
-        "img_id = img_a.view(-1, img_a.shape[0], img_a.shape[1], img_a.shape[2])\n",
+        "with torch.no_grad():\n",
+        "    pic_a = opt.pic_a_path\n",
+        "    # img_a = Image.open(pic_a).convert('RGB')\n",
+        "    img_a_whole = cv2.imread(pic_a)\n",
+        "    img_a_align_crop, _ = app.get(img_a_whole,crop_size)\n",
+        "    img_a_align_crop_pil = Image.fromarray(cv2.cvtColor(img_a_align_crop[0],cv2.COLOR_BGR2RGB)) \n",
+        "    img_a = transformer_Arcface(img_a_align_crop_pil)\n",
+        "    img_id = img_a.view(-1, img_a.shape[0], img_a.shape[1], img_a.shape[2])\n",
         "\n",
-        "# convert numpy to tensor\n",
-        "img_id = img_id.cuda()\n",
+        "    # convert numpy to tensor\n",
+        "    img_id = img_id.cuda()\n",
         "\n",
-        "#create latent id\n",
-        "img_id_downsample = F.interpolate(img_id, scale_factor=0.5)\n",
-        "latend_id = model.netArc(img_id_downsample)\n",
-        "latend_id = latend_id.detach().to('cpu')\n",
-        "latend_id = latend_id/np.linalg.norm(latend_id,axis=1,keepdims=True)\n",
-        "latend_id = latend_id.to('cuda')\n",
+        "    #create latent id\n",
+        "    img_id_downsample = F.interpolate(img_id, scale_factor=0.5)\n",
+        "    latend_id = model.netArc(img_id_downsample)\n",
+        "    latend_id = latend_id.detach().to('cpu')\n",
+        "    latend_id = latend_id/np.linalg.norm(latend_id,axis=1,keepdims=True)\n",
+        "    latend_id = latend_id.to('cuda')\n",
         "\n",
-        "video_swap(opt.video_path, latend_id, model, app, opt.output_path,temp_results_dir=opt.temp_path)"
+        "    video_swap(opt.video_path, latend_id, model, app, opt.output_path, temp_results_dir=opt.temp_path, use_mask=opt.use_mask)"
       ],
       "execution_count": 9,
       "outputs": [

cog.yaml

@@ -0,0 +1,20 @@
+build:
+  gpu: true
+  python_version: "3.8"
+  system_packages:
+    - "libgl1-mesa-glx"
+    - "libglib2.0-0"
+  python_packages:
+    - "imageio==2.9.0"
+    - "torch==1.8.0"
+    - "torchvision==0.9.0"
+    - "numpy==1.21.1"
+    - "insightface==0.2.1"
+    - "ipython==7.21.0"
+    - "Pillow==8.3.1"
+    - "opencv-python==4.5.3.56"
+    - "Fraction==1.5.1"
+    - "onnxruntime-gpu==1.8.1"
+    - "moviepy==1.0.3"
+
+predict: "predict.py:Predictor"

docs/css/page.css

@@ -22,6 +22,7 @@
     margin: 1% 1% 1% 1%;
     width: 100%;
     height: auto;
+
 }
 
 .which-image img {
@@ -29,7 +30,20 @@
     width: 100%;
 }
 
-.image-display img {
+.image-display2 img {
     float: right;
-    width: 65%;
+    width: 100%;
+}
+/* .image-display{
+    align-items: center;
+}
+.image-display2{
+    align-items: center;
+} */
+.select-show {
+    border-style: dashed;
+    border-width: 2px;
+    border-color: purple;
+
+    /* padding: 2px; */
 }

docs/guidance/preparation.md

@@ -6,11 +6,14 @@
 ```
 conda create -n simswap python=3.6
 conda activate simswap
-conda install pytorch==1.8.0 torchvision==0.9.0 torchaudio==0.8.0 cudatoolkit=10.2 -c pytorch
+conda install pytorch==1.8.0 torchvision==0.9.0 torchaudio==0.8.0 cudatoolkit=11.1 -c pytorch -c nvidia -c conda-forge
 (option): pip install --ignore-installed imageio
 pip install insightface==0.2.1 onnxruntime moviepy
+(option): pip install onnxruntime-gpu  (If you want to reduce the inference time)(It will be diffcult to install onnxruntime-gpu , the specify version of onnxruntime-gpu may depends on your machine and cuda version.)
 ```
+- ***We have now updated the prepare document. The main change gpu version of onnx is supported now. If you have configured the environment before, now use pip install onnxruntime-gpu ,You can increase the computing speed.***
 - We use the face detection and alignment methods from **[insightface](https://github.com/deepinsight/insightface)** for image preprocessing. Please download the relative files and unzip them to ./insightface_func/models from [this link](https://onedrive.live.com/?authkey=%21ADJ0aAOSsc90neY&cid=4A83B6B633B029CC&id=4A83B6B633B029CC%215837&parId=4A83B6B633B029CC%215834&action=locate).
+- We use the face parsing from **[face-parsing.PyTorch](https://github.com/zllrunning/face-parsing.PyTorch)** for image postprocessing. Please download the relative file and place it in ./parsing_model/checkpoint from [this link](https://drive.google.com/file/d/154JgKpzCPW82qINcVieuPH3fZ2e0P812/view).
 - The pytorch and cuda versions above are most recommanded. They may vary.
 - Using insightface with different versions is not recommanded. Please use this specific version.
 - These settings are tested valid on both Windows and Ununtu.
@@ -25,4 +28,4 @@ There are two archive files in the drive: **checkpoints.zip** and **arcface_chec
 [[Baidu Drive]](https://pan.baidu.com/s/1wFV11RVZMHqd-ky4YpLdcA) Password: ```jd2v```
 
 ### Note
-We expect users to have GPU with at least 8G memory. For those who do not, we provide [[Colab Notebook implementation]](https://colab.research.google.com/github/neuralchen/SimSwap/blob/main/SimSwap%20colab.ipynb).
+We expect users to have GPU with at least 3G memory. For those who do not, we provide [[Colab Notebook implementation]](https://colab.research.google.com/github/neuralchen/SimSwap/blob/main/SimSwap%20colab.ipynb).

docs/guidance/usage.md

@@ -10,31 +10,32 @@
 # Usage
 
 ###### Before running, please make sure you have installed the environment and downloaded requested files according to the [preparation guidance](./preparation.md).
+###### The below example command lines are using mask by default.
 
 ### Simple face swapping for already face-aligned images
 ```
-python test_one_image.py --isTrain false  --name people --Arc_path arcface_model/arcface_checkpoint.tar --pic_a_path crop_224/6.jpg --pic_b_path crop_224/ds.jpg --output_path output/
+python test_one_image.py --isTrain false --name people --Arc_path arcface_model/arcface_checkpoint.tar --pic_a_path crop_224/6.jpg --pic_b_path crop_224/ds.jpg --output_path output/
 ```
 
 ### Face swapping for video
 
 - Swap only one face within the video(the one with highest confidence by face detection).
 ```
-python test_video_swapsingle.py --isTrain false  --name people --Arc_path arcface_model/arcface_checkpoint.tar --pic_a_path ./demo_file/Iron_man.jpg --video_path ./demo_file/multi_people_1080p.mp4 --output_path ./output/multi_test_swapsingle.mp4 --temp_path ./temp_results
+python test_video_swapsingle.py --isTrain false --use_mask --name people --Arc_path arcface_model/arcface_checkpoint.tar --pic_a_path ./demo_file/Iron_man.jpg --video_path ./demo_file/multi_people_1080p.mp4 --output_path ./output/multi_test_swapsingle.mp4 --temp_path ./temp_results
 ```
 - Swap all faces within the video.
 ```
-python test_video_swapmulti.py --isTrain false  --name people --Arc_path arcface_model/arcface_checkpoint.tar --pic_a_path ./demo_file/Iron_man.jpg --video_path ./demo_file/multi_people_1080p.mp4 --output_path ./output/multi_test_swapmulti.mp4 --temp_path ./temp_results
+python test_video_swapmulti.py --isTrain false --use_mask --name people --Arc_path arcface_model/arcface_checkpoint.tar --pic_a_path ./demo_file/Iron_man.jpg --video_path ./demo_file/multi_people_1080p.mp4 --output_path ./output/multi_test_swapmulti.mp4 --temp_path ./temp_results
 ```
 - Swap the ***specific*** face within the video.
 ```
-python test_video_swapspecific.py --pic_specific_path ./demo_file/specific1.png --isTrain false  --name people --Arc_path arcface_model/arcface_checkpoint.tar --pic_a_path ./demo_file/Iron_man.jpg --video_path ./demo_file/multi_people_1080p.mp4 --output_path ./output/multi_test_specific.mp4 --temp_path ./temp_results 
+python test_video_swapspecific.py --use_mask --pic_specific_path ./demo_file/specific1.png --isTrain false  --name people --Arc_path arcface_model/arcface_checkpoint.tar --pic_a_path ./demo_file/Iron_man.jpg --video_path ./demo_file/multi_people_1080p.mp4 --output_path ./output/multi_test_specific.mp4 --temp_path ./temp_results 
 ```
 When changing the specified face, you need to give a picture of the person whose face is to be changed. Then assign the picture path to the argument "***--pic_specific_path***". This picture should be a front face and show the entire head and neck, which can help accurately change the face (if you still don’t know how to choose the picture, you can refer to the specific*.png of [./demo_file/](https://github.com/neuralchen/SimSwap/tree/main/demo_file)). It would be better if this picture was taken from the video to be changed.
 
 - Swap ***multi specific*** face with **multi specific id** within the video.
 ```
-python test_video_swap_multispecific.py --isTrain false  --name people --Arc_path arcface_model/arcface_checkpoint.tar --video_path ./demo_file/multi_people_1080p.mp4 --output_path ./output/multi_test_multispecific.mp4 --temp_path ./temp_results --multisepcific_dir ./demo_file/multispecific
+python test_video_swap_multispecific.py --isTrain false --use_mask  --name people --Arc_path arcface_model/arcface_checkpoint.tar --video_path ./demo_file/multi_people_1080p.mp4 --output_path ./output/multi_test_multispecific.mp4 --temp_path ./temp_results --multisepcific_dir ./demo_file/multispecific
 ```
 The folder you assign to ***"--multisepcific_dir"*** should be looked like:
 ```
@@ -55,34 +56,37 @@ The result is that the face corresponding to SRC_01.jpg (png) in the video will
 
 - Swap only one face within one image(the one with highest confidence by face detection). The result would be saved to ./output/result_whole_swapsingle.jpg
 ```
-python test_wholeimage_swapsingle.py --isTrain false  --name people --Arc_path arcface_model/arcface_checkpoint.tar --pic_a_path ./demo_file/Iron_man.jpg --pic_b_path ./demo_file/multi_people.jpg --output_path ./output/
+python test_wholeimage_swapsingle.py --isTrain false --use_mask  --name people --Arc_path arcface_model/arcface_checkpoint.tar --pic_a_path ./demo_file/Iron_man.jpg --pic_b_path ./demo_file/multi_people.jpg --output_path ./output/
 ```
 - Swap all faces within one image. The result would be saved to ./output/result_whole_swapmulti.jpg
 ```
-python test_wholeimage_swapmulti.py --isTrain false  --name people --Arc_path arcface_model/arcface_checkpoint.tar --pic_a_path ./demo_file/Iron_man.jpg --pic_b_path ./demo_file/multi_people.jpg --output_path ./output/
+python test_wholeimage_swapmulti.py --isTrain false --use_mask  --name people --Arc_path arcface_model/arcface_checkpoint.tar --pic_a_path ./demo_file/Iron_man.jpg --pic_b_path ./demo_file/multi_people.jpg --output_path ./output/
 ```
 - Swap **specific** face within one image. The result would be saved to ./output/result_whole_swapspecific.jpg
 ```
-python test_wholeimage_swapspecific.py --isTrain false  --name people --Arc_path arcface_model/arcface_checkpoint.tar --pic_a_path ./demo_file/Iron_man.jpg --pic_b_path ./demo_file/multi_people.jpg --output_path ./output/ --pic_specific_path ./demo_file/specific2.png
+python test_wholeimage_swapspecific.py --isTrain false --use_mask  --name people --Arc_path arcface_model/arcface_checkpoint.tar --pic_a_path ./demo_file/Iron_man.jpg --pic_b_path ./demo_file/multi_people.jpg --output_path ./output/ --pic_specific_path ./demo_file/specific2.png
 ```
 - Swap **multi specific** face with **multi specific id** within one image. The result would be saved to ./output/result_whole_swap_multispecific.jpg
 ```
-python test_wholeimage_swap_multispecific.py --isTrain false  --name people --Arc_path arcface_model/arcface_checkpoint.tar --pic_b_path ./demo_file/multi_people.jpg --output_path ./output/ --multisepcific_dir ./demo_file/multispecific
+python test_wholeimage_swap_multispecific.py --isTrain false --use_mask  --name people --Arc_path arcface_model/arcface_checkpoint.tar --pic_b_path ./demo_file/multi_people.jpg --output_path ./output/ --multisepcific_dir ./demo_file/multispecific
 ```
 ### About watermark of simswap logo
-The above example command line is to add the simswap logo as the watermark by default. After our discussion, we have added a hyper parameter to control whether to remove watermark.
+The above example command lines are to add the simswap logo as the watermark by default. After our discussion, we have added a hyper parameter to control whether to remove watermark.
 
 The usage of removing the watermark is to add an argument: "***--no_simswaplogo***" to the command line, take the command line of "Swap all faces within one image" as an example, the following command line can get the result without watermark:
 ```
-python test_wholeimage_swapmulti.py --no_simswaplogo --isTrain false  --name people --Arc_path arcface_model/arcface_checkpoint.tar --pic_a_path ./demo_file/Iron_man.jpg --pic_b_path ./demo_file/multi_people.jpg --output_path ./output/
+python test_wholeimage_swapmulti.py --no_simswaplogo --isTrain false --use_mask  --name people --Arc_path arcface_model/arcface_checkpoint.tar --pic_a_path ./demo_file/Iron_man.jpg --pic_b_path ./demo_file/multi_people.jpg --output_path ./output/
 ```
+### About using mask for better result
+We provide two methods to paste the face back to the original image after changing the face: Using mask or using bounding box. At present, the effect of using mask is the best. All the above code examples are using mask. If you want to use the bounding box, you only need to remove the --use_mask in the code example.
+Difference between using mask and not using mask can be found [here](https://imgsli.com/NjE3OTA).
 
-
-
-
-Difference between single face swapping and all face swapping are shown below.
+### Difference between single face swapping and all face swapping are shown below.
 <img src="../img/multi_face_comparison.png"/>
 
+
+
+
 ### Parameters
 |  Parameters   | Function  |
 |  :----  | :----  |
@@ -95,6 +99,7 @@ Difference between single face swapping and all face swapping are shown below.
 | --temp_path  | Path to store intermediate files  |
 | --output_path  | Path of directory to store the face swapping result  |
 | --no_simswaplogo  |The hyper parameter to control whether to remove watermark |
+| --use_mask  |The hyper parameter to control whether to use face parsing for the better visual effects(I recommend to use)|
 
 ### Note
-We expect users to have GPU with at least 6G memory. For those who do not, we will provide Colab Notebook implementation in the future.
+We expect users to have GPU with at least 3G memory.the For those who do not, we will provide Colab Notebook implementation in the future.

docs/index.html

@@ -150,19 +150,19 @@
                 <h3>Single Face Video Swap</h3>
             </div>
         </div>
-        <div class="row">
+        <div class="row select-show">
             <div class="col-md-12">
                 <div class="col-md-3 which-image-container">
                     <div class="which-image" id=anni>
                         <img src="./img/id/anni.jpg" onclick="select_source(0)" />
                     </div>
                     <div class="which-image" id=chenglong>
-                        <img src="./img/id/chenglong.jpg" onclick="select_source(1)" />
+                        <img src="./img/id/zhoujielun.jpg" onclick="select_source(1)" />
                     </div>
                 </div>
                 <div class="col-md-3 which-image-container">
                     <div class="which-image" id=zhoujielun>
-                        <img src="./img/id/zhoujielun.jpg" onclick="select_source(2)" />
+                        <img src="./img/id/chenglong.jpg" onclick="select_source(2)" />
                     </div>
                     <div class="which-image" id=zhuyin>
                         <img src="./img/id/zhuyin.jpg" onclick="select_source(3)" />
@@ -172,11 +172,23 @@
                     <img id=jiroujinlun src="./img/anni.webp" />
                 </div>
             </div>
-
-
-            <!-- <img width=24% src="./img/chenglong.webp" />
-            <img width=24% src="./img/zhoujielun.webp" />
-            <img width=24% src="./img/zhuyin.webp" /> -->
+        </div>
+        <div class="row select-show" style="margin-top: 10px;">
+            <div class="col-md-12">
+                <div class="col-md-3 which-image-container">
+                    <div class="which-image" id=Iron_man>
+                        <img src="./img/id/Iron_man.jpg" onclick="select_source2(0)" />
+                    </div>
+                </div>
+                <div class="col-md-3 which-image-container">
+                    <div class="which-image" id=wuyifan>
+                        <img src="./img/id/wuyifan.png" onclick="select_source2(1)" />
+                    </div>
+                </div>
+                <div class="col-md-6 image-display2">
+                    <img id=mama src="./img/mama_mask_short.webp" />
+                </div>
+            </div>
         </div>
         <div class="row">
             <div class="col-md-*">
@@ -319,6 +331,15 @@
     <script src="./js/bootstrap.min.js"></script>
     <!-- IE10 viewport hack for Surface/desktop Windows 8 bug -->
     <script src="./js/ie10-viewport-bug-workaround.js"></script>
+    <script>
+        document.getElementById('Iron_man').style.borderWidth = '5px';
+        document.getElementById('Iron_man').style.borderColor = 'red';
+        document.getElementById('Iron_man').style.borderStyle = 'outset';
+
+        document.getElementById('anni').style.borderWidth = '5px';
+        document.getElementById('anni').style.borderColor = 'red';
+        document.getElementById('anni').style.borderStyle = 'outset';
+    </script>
 </body>
 
 </html>

docs/js/which-image.js

@@ -2,8 +2,8 @@
  * @FilePath: \SimSwap\docs\js\which-image.js
  * @Author: Ziang Liu
  * @Date: 2021-07-03 16:34:56
- * @LastEditors: Ziang Liu
- * @LastEditTime: 2021-07-03 16:44:10
+ * @LastEditors: AceSix
+ * @LastEditTime: 2021-07-20 00:46:27
  * Copyright (C) 2021 SJTU. All rights reserved.
  */
 
@@ -24,4 +24,27 @@ function select_source(number) {
     }
     document.getElementById('jiroujinlun').src = './img/' + item_id + '.webp';
 
+}
+
+function select_source2(number) {
+    var items = ['Iron_man', 'wuyifan'];
+    var item_id = items[number];
+
+    for (i = 0; i < 2; i++) {
+        if (number == i) {
+            document.getElementById(items[i]).style.borderWidth = '5px';
+            document.getElementById(items[i]).style.borderColor = 'red';
+            document.getElementById(items[i]).style.borderStyle = 'outset';
+        } else {
+            document.getElementById(items[i]).style.border = 'none';
+        }
+    }
+    if (item_id=='Iron_man'){
+        document.getElementById('mama').src = './img/mama_mask_short.webp';
+    }
+    else{
+        document.getElementById('mama').src = './img/mama_mask_wuyifan_short.webp';
+
+    }
+    
 }

download-weights.sh

@@ -0,0 +1,9 @@
+#!/bin/sh
+wget -P ./arcface_model https://github.com/neuralchen/SimSwap/releases/download/1.0/arcface_checkpoint.tar
+wget https://github.com/neuralchen/SimSwap/releases/download/1.0/checkpoints.zip
+unzip ./checkpoints.zip  -d ./checkpoints
+rm checkpoints.zip
+wget --no-check-certificate "https://sh23tw.dm.files.1drv.com/y4mmGiIkNVigkSwOKDcV3nwMJulRGhbtHdkheehR5TArc52UjudUYNXAEvKCii2O5LAmzGCGK6IfleocxuDeoKxDZkNzDRSt4ZUlEt8GlSOpCXAFEkBwaZimtWGDRbpIGpb_pz9Nq5jATBQpezBS6G_UtspWTkgrXHHxhviV2nWy8APPx134zOZrUIbkSF6xnsqzs3uZ_SEX_m9Rey0ykpx9w" -O antelope.zip
+mkdir -p insightface_func/models
+unzip ./antelope.zip -d ./insightface_func/models/
+rm antelope.zip

options/test_options.py

@@ -24,6 +24,7 @@ class TestOptions(BaseOptions):
         self.parser.add_argument("--output_path", type=str, default='./output/', help="results path")
         self.parser.add_argument('--id_thres', type=float, default=0.03, help='how many test images to run')
         self.parser.add_argument('--no_simswaplogo', action='store_true', help='Remove the watermark')
+        self.parser.add_argument('--use_mask', action='store_true', help='Use mask for better result')
         
-
+ 
         self.isTrain = False

parsing_model/model.py

@@ -0,0 +1,283 @@
+#!/usr/bin/python
+# -*- encoding: utf-8 -*-
+
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision
+
+from parsing_model.resnet import Resnet18
+# from modules.bn import InPlaceABNSync as BatchNorm2d
+
+
+class ConvBNReLU(nn.Module):
+    def __init__(self, in_chan, out_chan, ks=3, stride=1, padding=1, *args, **kwargs):
+        super(ConvBNReLU, self).__init__()
+        self.conv = nn.Conv2d(in_chan,
+                out_chan,
+                kernel_size = ks,
+                stride = stride,
+                padding = padding,
+                bias = False)
+        self.bn = nn.BatchNorm2d(out_chan)
+        self.init_weight()
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = F.relu(self.bn(x))
+        return x
+
+    def init_weight(self):
+        for ly in self.children():
+            if isinstance(ly, nn.Conv2d):
+                nn.init.kaiming_normal_(ly.weight, a=1)
+                if not ly.bias is None: nn.init.constant_(ly.bias, 0)
+
+class BiSeNetOutput(nn.Module):
+    def __init__(self, in_chan, mid_chan, n_classes, *args, **kwargs):
+        super(BiSeNetOutput, self).__init__()
+        self.conv = ConvBNReLU(in_chan, mid_chan, ks=3, stride=1, padding=1)
+        self.conv_out = nn.Conv2d(mid_chan, n_classes, kernel_size=1, bias=False)
+        self.init_weight()
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.conv_out(x)
+        return x
+
+    def init_weight(self):
+        for ly in self.children():
+            if isinstance(ly, nn.Conv2d):
+                nn.init.kaiming_normal_(ly.weight, a=1)
+                if not ly.bias is None: nn.init.constant_(ly.bias, 0)
+
+    def get_params(self):
+        wd_params, nowd_params = [], []
+        for name, module in self.named_modules():
+            if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d):
+                wd_params.append(module.weight)
+                if not module.bias is None:
+                    nowd_params.append(module.bias)
+            elif isinstance(module, nn.BatchNorm2d):
+                nowd_params += list(module.parameters())
+        return wd_params, nowd_params
+
+
+class AttentionRefinementModule(nn.Module):
+    def __init__(self, in_chan, out_chan, *args, **kwargs):
+        super(AttentionRefinementModule, self).__init__()
+        self.conv = ConvBNReLU(in_chan, out_chan, ks=3, stride=1, padding=1)
+        self.conv_atten = nn.Conv2d(out_chan, out_chan, kernel_size= 1, bias=False)
+        self.bn_atten = nn.BatchNorm2d(out_chan)
+        self.sigmoid_atten = nn.Sigmoid()
+        self.init_weight()
+
+    def forward(self, x):
+        feat = self.conv(x)
+        atten = F.avg_pool2d(feat, feat.size()[2:])
+        atten = self.conv_atten(atten)
+        atten = self.bn_atten(atten)
+        atten = self.sigmoid_atten(atten)
+        out = torch.mul(feat, atten)
+        return out
+
+    def init_weight(self):
+        for ly in self.children():
+            if isinstance(ly, nn.Conv2d):
+                nn.init.kaiming_normal_(ly.weight, a=1)
+                if not ly.bias is None: nn.init.constant_(ly.bias, 0)
+
+
+class ContextPath(nn.Module):
+    def __init__(self, *args, **kwargs):
+        super(ContextPath, self).__init__()
+        self.resnet = Resnet18()
+        self.arm16 = AttentionRefinementModule(256, 128)
+        self.arm32 = AttentionRefinementModule(512, 128)
+        self.conv_head32 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1)
+        self.conv_head16 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1)
+        self.conv_avg = ConvBNReLU(512, 128, ks=1, stride=1, padding=0)
+
+        self.init_weight()
+
+    def forward(self, x):
+        H0, W0 = x.size()[2:]
+        feat8, feat16, feat32 = self.resnet(x)
+        H8, W8 = feat8.size()[2:]
+        H16, W16 = feat16.size()[2:]
+        H32, W32 = feat32.size()[2:]
+
+        avg = F.avg_pool2d(feat32, feat32.size()[2:])
+        avg = self.conv_avg(avg)
+        avg_up = F.interpolate(avg, (H32, W32), mode='nearest')
+
+        feat32_arm = self.arm32(feat32)
+        feat32_sum = feat32_arm + avg_up
+        feat32_up = F.interpolate(feat32_sum, (H16, W16), mode='nearest')
+        feat32_up = self.conv_head32(feat32_up)
+
+        feat16_arm = self.arm16(feat16)
+        feat16_sum = feat16_arm + feat32_up
+        feat16_up = F.interpolate(feat16_sum, (H8, W8), mode='nearest')
+        feat16_up = self.conv_head16(feat16_up)
+
+        return feat8, feat16_up, feat32_up  # x8, x8, x16
+
+    def init_weight(self):
+        for ly in self.children():
+            if isinstance(ly, nn.Conv2d):
+                nn.init.kaiming_normal_(ly.weight, a=1)
+                if not ly.bias is None: nn.init.constant_(ly.bias, 0)
+
+    def get_params(self):
+        wd_params, nowd_params = [], []
+        for name, module in self.named_modules():
+            if isinstance(module, (nn.Linear, nn.Conv2d)):
+                wd_params.append(module.weight)
+                if not module.bias is None:
+                    nowd_params.append(module.bias)
+            elif isinstance(module, nn.BatchNorm2d):
+                nowd_params += list(module.parameters())
+        return wd_params, nowd_params
+
+
+### This is not used, since I replace this with the resnet feature with the same size
+class SpatialPath(nn.Module):
+    def __init__(self, *args, **kwargs):
+        super(SpatialPath, self).__init__()
+        self.conv1 = ConvBNReLU(3, 64, ks=7, stride=2, padding=3)
+        self.conv2 = ConvBNReLU(64, 64, ks=3, stride=2, padding=1)
+        self.conv3 = ConvBNReLU(64, 64, ks=3, stride=2, padding=1)
+        self.conv_out = ConvBNReLU(64, 128, ks=1, stride=1, padding=0)
+        self.init_weight()
+
+    def forward(self, x):
+        feat = self.conv1(x)
+        feat = self.conv2(feat)
+        feat = self.conv3(feat)
+        feat = self.conv_out(feat)
+        return feat
+
+    def init_weight(self):
+        for ly in self.children():
+            if isinstance(ly, nn.Conv2d):
+                nn.init.kaiming_normal_(ly.weight, a=1)
+                if not ly.bias is None: nn.init.constant_(ly.bias, 0)
+
+    def get_params(self):
+        wd_params, nowd_params = [], []
+        for name, module in self.named_modules():
+            if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d):
+                wd_params.append(module.weight)
+                if not module.bias is None:
+                    nowd_params.append(module.bias)
+            elif isinstance(module, nn.BatchNorm2d):
+                nowd_params += list(module.parameters())
+        return wd_params, nowd_params
+
+
+class FeatureFusionModule(nn.Module):
+    def __init__(self, in_chan, out_chan, *args, **kwargs):
+        super(FeatureFusionModule, self).__init__()
+        self.convblk = ConvBNReLU(in_chan, out_chan, ks=1, stride=1, padding=0)
+        self.conv1 = nn.Conv2d(out_chan,
+                out_chan//4,
+                kernel_size = 1,
+                stride = 1,
+                padding = 0,
+                bias = False)
+        self.conv2 = nn.Conv2d(out_chan//4,
+                out_chan,
+                kernel_size = 1,
+                stride = 1,
+                padding = 0,
+                bias = False)
+        self.relu = nn.ReLU(inplace=True)
+        self.sigmoid = nn.Sigmoid()
+        self.init_weight()
+
+    def forward(self, fsp, fcp):
+        fcat = torch.cat([fsp, fcp], dim=1)
+        feat = self.convblk(fcat)
+        atten = F.avg_pool2d(feat, feat.size()[2:])
+        atten = self.conv1(atten)
+        atten = self.relu(atten)
+        atten = self.conv2(atten)
+        atten = self.sigmoid(atten)
+        feat_atten = torch.mul(feat, atten)
+        feat_out = feat_atten + feat
+        return feat_out
+
+    def init_weight(self):
+        for ly in self.children():
+            if isinstance(ly, nn.Conv2d):
+                nn.init.kaiming_normal_(ly.weight, a=1)
+                if not ly.bias is None: nn.init.constant_(ly.bias, 0)
+
+    def get_params(self):
+        wd_params, nowd_params = [], []
+        for name, module in self.named_modules():
+            if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d):
+                wd_params.append(module.weight)
+                if not module.bias is None:
+                    nowd_params.append(module.bias)
+            elif isinstance(module, nn.BatchNorm2d):
+                nowd_params += list(module.parameters())
+        return wd_params, nowd_params
+
+
+class BiSeNet(nn.Module):
+    def __init__(self, n_classes, *args, **kwargs):
+        super(BiSeNet, self).__init__()
+        self.cp = ContextPath()
+        ## here self.sp is deleted
+        self.ffm = FeatureFusionModule(256, 256)
+        self.conv_out = BiSeNetOutput(256, 256, n_classes)
+        self.conv_out16 = BiSeNetOutput(128, 64, n_classes)
+        self.conv_out32 = BiSeNetOutput(128, 64, n_classes)
+        self.init_weight()
+
+    def forward(self, x):
+        H, W = x.size()[2:]
+        feat_res8, feat_cp8, feat_cp16 = self.cp(x)  # here return res3b1 feature
+        feat_sp = feat_res8  # use res3b1 feature to replace spatial path feature
+        feat_fuse = self.ffm(feat_sp, feat_cp8)
+
+        feat_out = self.conv_out(feat_fuse)
+        feat_out16 = self.conv_out16(feat_cp8)
+        feat_out32 = self.conv_out32(feat_cp16)
+
+        feat_out = F.interpolate(feat_out, (H, W), mode='bilinear', align_corners=True)
+        feat_out16 = F.interpolate(feat_out16, (H, W), mode='bilinear', align_corners=True)
+        feat_out32 = F.interpolate(feat_out32, (H, W), mode='bilinear', align_corners=True)
+        return feat_out, feat_out16, feat_out32
+
+    def init_weight(self):
+        for ly in self.children():
+            if isinstance(ly, nn.Conv2d):
+                nn.init.kaiming_normal_(ly.weight, a=1)
+                if not ly.bias is None: nn.init.constant_(ly.bias, 0)
+
+    def get_params(self):
+        wd_params, nowd_params, lr_mul_wd_params, lr_mul_nowd_params = [], [], [], []
+        for name, child in self.named_children():
+            child_wd_params, child_nowd_params = child.get_params()
+            if isinstance(child, FeatureFusionModule) or isinstance(child, BiSeNetOutput):
+                lr_mul_wd_params += child_wd_params
+                lr_mul_nowd_params += child_nowd_params
+            else:
+                wd_params += child_wd_params
+                nowd_params += child_nowd_params
+        return wd_params, nowd_params, lr_mul_wd_params, lr_mul_nowd_params
+
+
+if __name__ == "__main__":
+    net = BiSeNet(19)
+    net.cuda()
+    net.eval()
+    in_ten = torch.randn(16, 3, 640, 480).cuda()
+    out, out16, out32 = net(in_ten)
+    print(out.shape)
+
+    net.get_params()

parsing_model/resnet.py

@@ -0,0 +1,109 @@
+#!/usr/bin/python
+# -*- encoding: utf-8 -*-
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.model_zoo as modelzoo
+
+# from modules.bn import InPlaceABNSync as BatchNorm2d
+
+resnet18_url = 'https://download.pytorch.org/models/resnet18-5c106cde.pth'
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=1, bias=False)
+
+
+class BasicBlock(nn.Module):
+    def __init__(self, in_chan, out_chan, stride=1):
+        super(BasicBlock, self).__init__()
+        self.conv1 = conv3x3(in_chan, out_chan, stride)
+        self.bn1 = nn.BatchNorm2d(out_chan)
+        self.conv2 = conv3x3(out_chan, out_chan)
+        self.bn2 = nn.BatchNorm2d(out_chan)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = None
+        if in_chan != out_chan or stride != 1:
+            self.downsample = nn.Sequential(
+                nn.Conv2d(in_chan, out_chan,
+                          kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(out_chan),
+                )
+
+    def forward(self, x):
+        residual = self.conv1(x)
+        residual = F.relu(self.bn1(residual))
+        residual = self.conv2(residual)
+        residual = self.bn2(residual)
+
+        shortcut = x
+        if self.downsample is not None:
+            shortcut = self.downsample(x)
+
+        out = shortcut + residual
+        out = self.relu(out)
+        return out
+
+
+def create_layer_basic(in_chan, out_chan, bnum, stride=1):
+    layers = [BasicBlock(in_chan, out_chan, stride=stride)]
+    for i in range(bnum-1):
+        layers.append(BasicBlock(out_chan, out_chan, stride=1))
+    return nn.Sequential(*layers)
+
+
+class Resnet18(nn.Module):
+    def __init__(self):
+        super(Resnet18, self).__init__()
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
+                               bias=False)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = create_layer_basic(64, 64, bnum=2, stride=1)
+        self.layer2 = create_layer_basic(64, 128, bnum=2, stride=2)
+        self.layer3 = create_layer_basic(128, 256, bnum=2, stride=2)
+        self.layer4 = create_layer_basic(256, 512, bnum=2, stride=2)
+        self.init_weight()
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = F.relu(self.bn1(x))
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        feat8 = self.layer2(x) # 1/8
+        feat16 = self.layer3(feat8) # 1/16
+        feat32 = self.layer4(feat16) # 1/32
+        return feat8, feat16, feat32
+
+    def init_weight(self):
+        state_dict = modelzoo.load_url(resnet18_url)
+        self_state_dict = self.state_dict()
+        for k, v in state_dict.items():
+            if 'fc' in k: continue
+            self_state_dict.update({k: v})
+        self.load_state_dict(self_state_dict)
+
+    def get_params(self):
+        wd_params, nowd_params = [], []
+        for name, module in self.named_modules():
+            if isinstance(module, (nn.Linear, nn.Conv2d)):
+                wd_params.append(module.weight)
+                if not module.bias is None:
+                    nowd_params.append(module.bias)
+            elif isinstance(module,  nn.BatchNorm2d):
+                nowd_params += list(module.parameters())
+        return wd_params, nowd_params
+
+
+if __name__ == "__main__":
+    net = Resnet18()
+    x = torch.randn(16, 3, 224, 224)
+    out = net(x)
+    print(out[0].size())
+    print(out[1].size())
+    print(out[2].size())
+    net.get_params()

predict.py

@@ -0,0 +1,101 @@
+import cog
+import tempfile
+from pathlib import Path
+import argparse
+import cv2
+import torch
+from PIL import Image
+import torch.nn.functional as F
+from torchvision import transforms
+from models.models import create_model
+from options.test_options import TestOptions
+from util.reverse2original import reverse2wholeimage
+from util.norm import SpecificNorm
+from test_wholeimage_swapmulti import _totensor
+from insightface_func.face_detect_crop_multi import Face_detect_crop as Face_detect_crop_multi
+from insightface_func.face_detect_crop_single import Face_detect_crop as Face_detect_crop_single
+
+
+class Predictor(cog.Predictor):
+    def setup(self):
+        self.transformer_Arcface = transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+        ])
+
+    @cog.input("source", type=Path, help="source image")
+    @cog.input("target", type=Path, help="target image")
+    @cog.input("mode", type=str, options=['single', 'all'], default='all',
+               help="swap a single face (the one with highest confidence by face detection) or all faces in the target image")
+    def predict(self, source, target, mode='all'):
+
+        app = Face_detect_crop_multi(name='antelope', root='./insightface_func/models')
+
+        if mode == 'single':
+            app = Face_detect_crop_single(name='antelope', root='./insightface_func/models')
+
+        app.prepare(ctx_id=0, det_thresh=0.6, det_size=(640, 640))
+
+        options = TestOptions()
+        options.initialize()
+        opt = options.parser.parse_args(["--Arc_path", 'arcface_model/arcface_checkpoint.tar', "--pic_a_path", str(source),
+                                         "--pic_b_path", str(target), "--isTrain", False, "--no_simswaplogo"])
+
+        str_ids = opt.gpu_ids.split(',')
+        opt.gpu_ids = []
+        for str_id in str_ids:
+            id = int(str_id)
+            if id >= 0:
+                opt.gpu_ids.append(id)
+
+        # set gpu ids
+        if len(opt.gpu_ids) > 0:
+            torch.cuda.set_device(opt.gpu_ids[0])
+
+        torch.nn.Module.dump_patches = True
+        model = create_model(opt)
+        model.eval()
+
+        crop_size = 224
+        spNorm = SpecificNorm()
+
+        with torch.no_grad():
+            pic_a = opt.pic_a_path
+            img_a_whole = cv2.imread(pic_a)
+            img_a_align_crop, _ = app.get(img_a_whole, crop_size)
+            img_a_align_crop_pil = Image.fromarray(cv2.cvtColor(img_a_align_crop[0], cv2.COLOR_BGR2RGB))
+            img_a = self.transformer_Arcface(img_a_align_crop_pil)
+            img_id = img_a.view(-1, img_a.shape[0], img_a.shape[1], img_a.shape[2])
+
+            # convert numpy to tensor
+            img_id = img_id.cuda()
+
+            # create latent id
+            img_id_downsample = F.interpolate(img_id, scale_factor=0.5)
+            latend_id = model.netArc(img_id_downsample)
+            latend_id = F.normalize(latend_id, p=2, dim=1)
+
+            ############## Forward Pass ######################
+
+            pic_b = opt.pic_b_path
+            img_b_whole = cv2.imread(pic_b)
+            img_b_align_crop_list, b_mat_list = app.get(img_b_whole, crop_size)
+
+            swap_result_list = []
+            b_align_crop_tenor_list = []
+
+            for b_align_crop in img_b_align_crop_list:
+                b_align_crop_tenor = _totensor(cv2.cvtColor(b_align_crop, cv2.COLOR_BGR2RGB))[None, ...].cuda()
+
+                swap_result = model(None, b_align_crop_tenor, latend_id, None, True)[0]
+                swap_result_list.append(swap_result)
+                b_align_crop_tenor_list.append(b_align_crop_tenor)
+
+            net = None
+
+            out_path = Path(tempfile.mkdtemp()) / "output.png"
+
+            reverse2wholeimage(b_align_crop_tenor_list, swap_result_list, b_mat_list, crop_size, img_b_whole, None,
+                               str(out_path), opt.no_simswaplogo,
+                               pasring_model=net, use_mask=opt.use_mask, norm=spNorm)
+            return out_path

test_one_image.py

@@ -22,10 +22,10 @@ transformer_Arcface = transforms.Compose([
         transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
     ])
 
-# detransformer = transforms.Compose([
-#         transforms.Normalize([0, 0, 0], [1/0.229, 1/0.224, 1/0.225]),
-#         transforms.Normalize([-0.485, -0.456, -0.406], [1, 1, 1])
-#     ])
+detransformer = transforms.Compose([
+        transforms.Normalize([0, 0, 0], [1/0.229, 1/0.224, 1/0.225]),
+        transforms.Normalize([-0.485, -0.456, -0.406], [1, 1, 1])
+    ])
 if __name__ == '__main__':
     opt = TestOptions().parse()
 
@@ -35,51 +35,52 @@ if __name__ == '__main__':
     model = create_model(opt)
     model.eval()
 
+    with torch.no_grad():
+        
+        pic_a = opt.pic_a_path
+        img_a = Image.open(pic_a).convert('RGB')
+        img_a = transformer_Arcface(img_a)
+        img_id = img_a.view(-1, img_a.shape[0], img_a.shape[1], img_a.shape[2])
 
-    pic_a = opt.pic_a_path
-    img_a = Image.open(pic_a).convert('RGB')
-    img_a = transformer_Arcface(img_a)
-    img_id = img_a.view(-1, img_a.shape[0], img_a.shape[1], img_a.shape[2])
+        pic_b = opt.pic_b_path
 
-    pic_b = opt.pic_b_path
+        img_b = Image.open(pic_b).convert('RGB')
+        img_b = transformer(img_b)
+        img_att = img_b.view(-1, img_b.shape[0], img_b.shape[1], img_b.shape[2])
 
-    img_b = Image.open(pic_b).convert('RGB')
-    img_b = transformer(img_b)
-    img_att = img_b.view(-1, img_b.shape[0], img_b.shape[1], img_b.shape[2])
+        # convert numpy to tensor
+        img_id = img_id.cuda()
+        img_att = img_att.cuda()
 
-    # convert numpy to tensor
-    img_id = img_id.cuda()
-    img_att = img_att.cuda()
-
-    #create latent id
-    img_id_downsample = F.interpolate(img_id, scale_factor=0.5)
-    latend_id = model.netArc(img_id_downsample)
-    latend_id = latend_id.detach().to('cpu')
-    latend_id = latend_id/np.linalg.norm(latend_id,axis=1,keepdims=True)
-    latend_id = latend_id.to('cuda')
+        #create latent id
+        img_id_downsample = F.interpolate(img_id, scale_factor=0.5)
+        latend_id = model.netArc(img_id_downsample)
+        latend_id = latend_id.detach().to('cpu')
+        latend_id = latend_id/np.linalg.norm(latend_id,axis=1,keepdims=True)
+        latend_id = latend_id.to('cuda')
 
 
-    ############## Forward Pass ######################
-    img_fake = model(img_id, img_att, latend_id, latend_id, True)
+        ############## Forward Pass ######################
+        img_fake = model(img_id, img_att, latend_id, latend_id, True)
 
 
-    for i in range(img_id.shape[0]):
-        if i == 0:
-            row1 = img_id[i]
-            row2 = img_att[i]
-            row3 = img_fake[i]
-        else:
-            row1 = torch.cat([row1, img_id[i]], dim=2)
-            row2 = torch.cat([row2, img_att[i]], dim=2)
-            row3 = torch.cat([row3, img_fake[i]], dim=2)
+        for i in range(img_id.shape[0]):
+            if i == 0:
+                row1 = img_id[i]
+                row2 = img_att[i]
+                row3 = img_fake[i]
+            else:
+                row1 = torch.cat([row1, img_id[i]], dim=2)
+                row2 = torch.cat([row2, img_att[i]], dim=2)
+                row3 = torch.cat([row3, img_fake[i]], dim=2)
 
-    #full = torch.cat([row1, row2, row3], dim=1).detach()
-    full = row3.detach()
-    full = full.permute(1, 2, 0)
-    output = full.to('cpu')
-    output = np.array(output)
-    output = output[..., ::-1]
+        #full = torch.cat([row1, row2, row3], dim=1).detach()
+        full = row3.detach()
+        full = full.permute(1, 2, 0)
+        output = full.to('cpu')
+        output = np.array(output)
+        output = output[..., ::-1]
 
-    output = output*255
+        output = output*255
 
-    cv2.imwrite(opt.output_path + 'result.jpg',output)
+        cv2.imwrite(opt.output_path + 'result.jpg',output)

test_video_swap_multispecific.py

@@ -51,44 +51,44 @@ if __name__ == '__main__':
     source_specific_id_nonorm_list = []
     source_path = os.path.join(multisepcific_dir,'SRC_*')
     source_specific_images_path = sorted(glob.glob(source_path))
-
-    for source_specific_image_path in source_specific_images_path:
-        specific_person_whole = cv2.imread(source_specific_image_path)
-        specific_person_align_crop, _ = app.get(specific_person_whole,crop_size)
-        specific_person_align_crop_pil = Image.fromarray(cv2.cvtColor(specific_person_align_crop[0],cv2.COLOR_BGR2RGB)) 
-        specific_person = transformer_Arcface(specific_person_align_crop_pil)
-        specific_person = specific_person.view(-1, specific_person.shape[0], specific_person.shape[1], specific_person.shape[2])
-        # convert numpy to tensor
-        specific_person = specific_person.cuda()
-        #create latent id
-        specific_person_downsample = F.interpolate(specific_person, scale_factor=0.5)
-        specific_person_id_nonorm = model.netArc(specific_person_downsample)
-        source_specific_id_nonorm_list.append(specific_person_id_nonorm.clone())
+    with torch.no_grad():
+        for source_specific_image_path in source_specific_images_path:
+            specific_person_whole = cv2.imread(source_specific_image_path)
+            specific_person_align_crop, _ = app.get(specific_person_whole,crop_size)
+            specific_person_align_crop_pil = Image.fromarray(cv2.cvtColor(specific_person_align_crop[0],cv2.COLOR_BGR2RGB)) 
+            specific_person = transformer_Arcface(specific_person_align_crop_pil)
+            specific_person = specific_person.view(-1, specific_person.shape[0], specific_person.shape[1], specific_person.shape[2])
+            # convert numpy to tensor
+            specific_person = specific_person.cuda()
+            #create latent id
+            specific_person_downsample = F.interpolate(specific_person, scale_factor=0.5)
+            specific_person_id_nonorm = model.netArc(specific_person_downsample)
+            source_specific_id_nonorm_list.append(specific_person_id_nonorm.clone())
 
 
-    # The person who provides id information (list)
-    target_id_norm_list = []
-    target_path = os.path.join(multisepcific_dir,'DST_*')
-    target_images_path = sorted(glob.glob(target_path))
+        # The person who provides id information (list)
+        target_id_norm_list = []
+        target_path = os.path.join(multisepcific_dir,'DST_*')
+        target_images_path = sorted(glob.glob(target_path))
 
-    for target_image_path in target_images_path:
-        img_a_whole = cv2.imread(target_image_path)
-        img_a_align_crop, _ = app.get(img_a_whole,crop_size)
-        img_a_align_crop_pil = Image.fromarray(cv2.cvtColor(img_a_align_crop[0],cv2.COLOR_BGR2RGB)) 
-        img_a = transformer_Arcface(img_a_align_crop_pil)
-        img_id = img_a.view(-1, img_a.shape[0], img_a.shape[1], img_a.shape[2])
-        # convert numpy to tensor
-        img_id = img_id.cuda()
-        #create latent id
-        img_id_downsample = F.interpolate(img_id, scale_factor=0.5)
-        latend_id = model.netArc(img_id_downsample)
-        latend_id = F.normalize(latend_id, p=2, dim=1)
-        target_id_norm_list.append(latend_id.clone())
+        for target_image_path in target_images_path:
+            img_a_whole = cv2.imread(target_image_path)
+            img_a_align_crop, _ = app.get(img_a_whole,crop_size)
+            img_a_align_crop_pil = Image.fromarray(cv2.cvtColor(img_a_align_crop[0],cv2.COLOR_BGR2RGB)) 
+            img_a = transformer_Arcface(img_a_align_crop_pil)
+            img_id = img_a.view(-1, img_a.shape[0], img_a.shape[1], img_a.shape[2])
+            # convert numpy to tensor
+            img_id = img_id.cuda()
+            #create latent id
+            img_id_downsample = F.interpolate(img_id, scale_factor=0.5)
+            latend_id = model.netArc(img_id_downsample)
+            latend_id = F.normalize(latend_id, p=2, dim=1)
+            target_id_norm_list.append(latend_id.clone())
 
-    assert len(target_id_norm_list) == len(source_specific_id_nonorm_list), "The number of images in source and target directory must be same !!!"
+        assert len(target_id_norm_list) == len(source_specific_id_nonorm_list), "The number of images in source and target directory must be same !!!"
 
 
 
-    video_swap(opt.video_path, target_id_norm_list,source_specific_id_nonorm_list, opt.id_thres, \
-        model, app, opt.output_path,temp_results_dir=opt.temp_path,no_simswaplogo=opt.no_simswaplogo)
+        video_swap(opt.video_path, target_id_norm_list,source_specific_id_nonorm_list, opt.id_thres, \
+            model, app, opt.output_path,temp_results_dir=opt.temp_path,no_simswaplogo=opt.no_simswaplogo,use_mask=opt.use_mask)
 

test_video_swapmulti.py

@@ -44,29 +44,31 @@ if __name__ == '__main__':
     app = Face_detect_crop(name='antelope', root='./insightface_func/models')
     app.prepare(ctx_id= 0, det_thresh=0.6, det_size=(640,640))
 
-    pic_a = opt.pic_a_path
-    # img_a = Image.open(pic_a).convert('RGB')
-    img_a_whole = cv2.imread(pic_a)
-    img_a_align_crop, _ = app.get(img_a_whole,crop_size)
-    img_a_align_crop_pil = Image.fromarray(cv2.cvtColor(img_a_align_crop[0],cv2.COLOR_BGR2RGB)) 
-    img_a = transformer_Arcface(img_a_align_crop_pil)
-    img_id = img_a.view(-1, img_a.shape[0], img_a.shape[1], img_a.shape[2])
+    with torch.no_grad():
+        pic_a = opt.pic_a_path
+        # img_a = Image.open(pic_a).convert('RGB')
+        img_a_whole = cv2.imread(pic_a)
+        img_a_align_crop, _ = app.get(img_a_whole,crop_size)
+        img_a_align_crop_pil = Image.fromarray(cv2.cvtColor(img_a_align_crop[0],cv2.COLOR_BGR2RGB)) 
+        img_a = transformer_Arcface(img_a_align_crop_pil)
+        img_id = img_a.view(-1, img_a.shape[0], img_a.shape[1], img_a.shape[2])
 
-    # pic_b = opt.pic_b_path
-    # img_b_whole = cv2.imread(pic_b)
-    # img_b_align_crop, b_mat = app.get(img_b_whole,crop_size)
-    # img_b_align_crop_pil = Image.fromarray(cv2.cvtColor(img_b_align_crop,cv2.COLOR_BGR2RGB)) 
-    # img_b = transformer(img_b_align_crop_pil)
-    # img_att = img_b.view(-1, img_b.shape[0], img_b.shape[1], img_b.shape[2])
+        # pic_b = opt.pic_b_path
+        # img_b_whole = cv2.imread(pic_b)
+        # img_b_align_crop, b_mat = app.get(img_b_whole,crop_size)
+        # img_b_align_crop_pil = Image.fromarray(cv2.cvtColor(img_b_align_crop,cv2.COLOR_BGR2RGB)) 
+        # img_b = transformer(img_b_align_crop_pil)
+        # img_att = img_b.view(-1, img_b.shape[0], img_b.shape[1], img_b.shape[2])
 
-    # convert numpy to tensor
-    img_id = img_id.cuda()
-    # img_att = img_att.cuda()
+        # convert numpy to tensor
+        img_id = img_id.cuda()
+        # img_att = img_att.cuda()
 
-    #create latent id
-    img_id_downsample = F.interpolate(img_id, scale_factor=0.5)
-    latend_id = model.netArc(img_id_downsample)
-    latend_id = F.normalize(latend_id, p=2, dim=1)
+        #create latent id
+        img_id_downsample = F.interpolate(img_id, scale_factor=0.5)
+        latend_id = model.netArc(img_id_downsample)
+        latend_id = F.normalize(latend_id, p=2, dim=1)
 
-    video_swap(opt.video_path, latend_id, model, app, opt.output_path,temp_results_dir=opt.temp_path,no_simswaplogo=opt.no_simswaplogo)
+        video_swap(opt.video_path, latend_id, model, app, opt.output_path,temp_results_dir=opt.temp_path,\
+            no_simswaplogo=opt.no_simswaplogo,use_mask=opt.use_mask)
 

test_video_swapsingle.py

@@ -43,30 +43,31 @@ if __name__ == '__main__':
 
     app = Face_detect_crop(name='antelope', root='./insightface_func/models')
     app.prepare(ctx_id= 0, det_thresh=0.6, det_size=(640,640))
+    with torch.no_grad():
+        pic_a = opt.pic_a_path
+        # img_a = Image.open(pic_a).convert('RGB')
+        img_a_whole = cv2.imread(pic_a)
+        img_a_align_crop, _ = app.get(img_a_whole,crop_size)
+        img_a_align_crop_pil = Image.fromarray(cv2.cvtColor(img_a_align_crop[0],cv2.COLOR_BGR2RGB)) 
+        img_a = transformer_Arcface(img_a_align_crop_pil)
+        img_id = img_a.view(-1, img_a.shape[0], img_a.shape[1], img_a.shape[2])
 
-    pic_a = opt.pic_a_path
-    # img_a = Image.open(pic_a).convert('RGB')
-    img_a_whole = cv2.imread(pic_a)
-    img_a_align_crop, _ = app.get(img_a_whole,crop_size)
-    img_a_align_crop_pil = Image.fromarray(cv2.cvtColor(img_a_align_crop[0],cv2.COLOR_BGR2RGB)) 
-    img_a = transformer_Arcface(img_a_align_crop_pil)
-    img_id = img_a.view(-1, img_a.shape[0], img_a.shape[1], img_a.shape[2])
+        # pic_b = opt.pic_b_path
+        # img_b_whole = cv2.imread(pic_b)
+        # img_b_align_crop, b_mat = app.get(img_b_whole,crop_size)
+        # img_b_align_crop_pil = Image.fromarray(cv2.cvtColor(img_b_align_crop,cv2.COLOR_BGR2RGB)) 
+        # img_b = transformer(img_b_align_crop_pil)
+        # img_att = img_b.view(-1, img_b.shape[0], img_b.shape[1], img_b.shape[2])
 
-    # pic_b = opt.pic_b_path
-    # img_b_whole = cv2.imread(pic_b)
-    # img_b_align_crop, b_mat = app.get(img_b_whole,crop_size)
-    # img_b_align_crop_pil = Image.fromarray(cv2.cvtColor(img_b_align_crop,cv2.COLOR_BGR2RGB)) 
-    # img_b = transformer(img_b_align_crop_pil)
-    # img_att = img_b.view(-1, img_b.shape[0], img_b.shape[1], img_b.shape[2])
+        # convert numpy to tensor
+        img_id = img_id.cuda()
+        # img_att = img_att.cuda()
 
-    # convert numpy to tensor
-    img_id = img_id.cuda()
-    # img_att = img_att.cuda()
+        #create latent id
+        img_id_downsample = F.interpolate(img_id, scale_factor=0.5)
+        latend_id = model.netArc(img_id_downsample)
+        latend_id = F.normalize(latend_id, p=2, dim=1)
 
-    #create latent id
-    img_id_downsample = F.interpolate(img_id, scale_factor=0.5)
-    latend_id = model.netArc(img_id_downsample)
-    latend_id = F.normalize(latend_id, p=2, dim=1)
-
-    video_swap(opt.video_path, latend_id, model, app, opt.output_path,temp_results_dir=opt.temp_path,no_simswaplogo=opt.no_simswaplogo)
+        video_swap(opt.video_path, latend_id, model, app, opt.output_path,temp_results_dir=opt.temp_path,\
+            no_simswaplogo=opt.no_simswaplogo,use_mask=opt.use_mask)
 

test_video_swapspecific.py

@@ -43,42 +43,42 @@ if __name__ == '__main__':
 
     app = Face_detect_crop(name='antelope', root='./insightface_func/models')
     app.prepare(ctx_id= 0, det_thresh=0.6, det_size=(640,640))
+    with torch.no_grad():
+        pic_a = opt.pic_a_path
+        # img_a = Image.open(pic_a).convert('RGB')
+        img_a_whole = cv2.imread(pic_a)
+        img_a_align_crop, _ = app.get(img_a_whole,crop_size)
+        img_a_align_crop_pil = Image.fromarray(cv2.cvtColor(img_a_align_crop[0],cv2.COLOR_BGR2RGB)) 
+        img_a = transformer_Arcface(img_a_align_crop_pil)
+        img_id = img_a.view(-1, img_a.shape[0], img_a.shape[1], img_a.shape[2])
 
-    pic_a = opt.pic_a_path
-    # img_a = Image.open(pic_a).convert('RGB')
-    img_a_whole = cv2.imread(pic_a)
-    img_a_align_crop, _ = app.get(img_a_whole,crop_size)
-    img_a_align_crop_pil = Image.fromarray(cv2.cvtColor(img_a_align_crop[0],cv2.COLOR_BGR2RGB)) 
-    img_a = transformer_Arcface(img_a_align_crop_pil)
-    img_id = img_a.view(-1, img_a.shape[0], img_a.shape[1], img_a.shape[2])
+        # pic_b = opt.pic_b_path
+        # img_b_whole = cv2.imread(pic_b)
+        # img_b_align_crop, b_mat = app.get(img_b_whole,crop_size)
+        # img_b_align_crop_pil = Image.fromarray(cv2.cvtColor(img_b_align_crop,cv2.COLOR_BGR2RGB)) 
+        # img_b = transformer(img_b_align_crop_pil)
+        # img_att = img_b.view(-1, img_b.shape[0], img_b.shape[1], img_b.shape[2])
 
-    # pic_b = opt.pic_b_path
-    # img_b_whole = cv2.imread(pic_b)
-    # img_b_align_crop, b_mat = app.get(img_b_whole,crop_size)
-    # img_b_align_crop_pil = Image.fromarray(cv2.cvtColor(img_b_align_crop,cv2.COLOR_BGR2RGB)) 
-    # img_b = transformer(img_b_align_crop_pil)
-    # img_att = img_b.view(-1, img_b.shape[0], img_b.shape[1], img_b.shape[2])
+        # convert numpy to tensor
+        img_id = img_id.cuda()
+        # img_att = img_att.cuda()
 
-    # convert numpy to tensor
-    img_id = img_id.cuda()
-    # img_att = img_att.cuda()
-
-    #create latent id
-    img_id_downsample = F.interpolate(img_id, scale_factor=0.5)
-    latend_id = model.netArc(img_id_downsample)
-    latend_id = F.normalize(latend_id, p=2, dim=1)
+        #create latent id
+        img_id_downsample = F.interpolate(img_id, scale_factor=0.5)
+        latend_id = model.netArc(img_id_downsample)
+        latend_id = F.normalize(latend_id, p=2, dim=1)
 
 
-    # The specific person to be swapped
-    specific_person_whole = cv2.imread(pic_specific)
-    specific_person_align_crop, _ = app.get(specific_person_whole,crop_size)
-    specific_person_align_crop_pil = Image.fromarray(cv2.cvtColor(specific_person_align_crop[0],cv2.COLOR_BGR2RGB)) 
-    specific_person = transformer_Arcface(specific_person_align_crop_pil)
-    specific_person = specific_person.view(-1, specific_person.shape[0], specific_person.shape[1], specific_person.shape[2])
-    specific_person = specific_person.cuda()
-    specific_person_downsample = F.interpolate(specific_person, scale_factor=0.5)
-    specific_person_id_nonorm = model.netArc(specific_person_downsample)
+        # The specific person to be swapped
+        specific_person_whole = cv2.imread(pic_specific)
+        specific_person_align_crop, _ = app.get(specific_person_whole,crop_size)
+        specific_person_align_crop_pil = Image.fromarray(cv2.cvtColor(specific_person_align_crop[0],cv2.COLOR_BGR2RGB)) 
+        specific_person = transformer_Arcface(specific_person_align_crop_pil)
+        specific_person = specific_person.view(-1, specific_person.shape[0], specific_person.shape[1], specific_person.shape[2])
+        specific_person = specific_person.cuda()
+        specific_person_downsample = F.interpolate(specific_person, scale_factor=0.5)
+        specific_person_id_nonorm = model.netArc(specific_person_downsample)
 
-    video_swap(opt.video_path, latend_id,specific_person_id_nonorm, opt.id_thres, \
-        model, app, opt.output_path,temp_results_dir=opt.temp_path,no_simswaplogo=opt.no_simswaplogo)
+        video_swap(opt.video_path, latend_id,specific_person_id_nonorm, opt.id_thres, \
+            model, app, opt.output_path,temp_results_dir=opt.temp_path,no_simswaplogo=opt.no_simswaplogo,use_mask=opt.use_mask)
 

test_wholeimage_swap_multispecific.py

@@ -15,6 +15,7 @@ from util.add_watermark import watermark_image
 import torch.nn as nn
 from util.norm import SpecificNorm
 import glob
+from parsing_model.model import BiSeNet
 
 def lcm(a, b): return abs(a * b) / fractions.gcd(a, b) if a and b else 0
 
@@ -53,93 +54,107 @@ if __name__ == '__main__':
     app = Face_detect_crop(name='antelope', root='./insightface_func/models')
     app.prepare(ctx_id= 0, det_thresh=0.6, det_size=(640,640))
 
-    # The specific person to be swapped(source)
+    with torch.no_grad():
+        # The specific person to be swapped(source)
 
-    source_specific_id_nonorm_list = []
-    source_path = os.path.join(multisepcific_dir,'SRC_*')
-    source_specific_images_path = sorted(glob.glob(source_path))
+        source_specific_id_nonorm_list = []
+        source_path = os.path.join(multisepcific_dir,'SRC_*')
+        source_specific_images_path = sorted(glob.glob(source_path))
 
-    for source_specific_image_path in source_specific_images_path:
-        specific_person_whole = cv2.imread(source_specific_image_path)
-        specific_person_align_crop, _ = app.get(specific_person_whole,crop_size)
-        specific_person_align_crop_pil = Image.fromarray(cv2.cvtColor(specific_person_align_crop[0],cv2.COLOR_BGR2RGB)) 
-        specific_person = transformer_Arcface(specific_person_align_crop_pil)
-        specific_person = specific_person.view(-1, specific_person.shape[0], specific_person.shape[1], specific_person.shape[2])
-        # convert numpy to tensor
-        specific_person = specific_person.cuda()
-        #create latent id
-        specific_person_downsample = F.interpolate(specific_person, scale_factor=0.5)
-        specific_person_id_nonorm = model.netArc(specific_person_downsample)
-        source_specific_id_nonorm_list.append(specific_person_id_nonorm.clone())
+        for source_specific_image_path in source_specific_images_path:
+            specific_person_whole = cv2.imread(source_specific_image_path)
+            specific_person_align_crop, _ = app.get(specific_person_whole,crop_size)
+            specific_person_align_crop_pil = Image.fromarray(cv2.cvtColor(specific_person_align_crop[0],cv2.COLOR_BGR2RGB)) 
+            specific_person = transformer_Arcface(specific_person_align_crop_pil)
+            specific_person = specific_person.view(-1, specific_person.shape[0], specific_person.shape[1], specific_person.shape[2])
+            # convert numpy to tensor
+            specific_person = specific_person.cuda()
+            #create latent id
+            specific_person_downsample = F.interpolate(specific_person, scale_factor=0.5)
+            specific_person_id_nonorm = model.netArc(specific_person_downsample)
+            source_specific_id_nonorm_list.append(specific_person_id_nonorm.clone())
 
 
-    # The person who provides id information (list)
-    target_id_norm_list = []
-    target_path = os.path.join(multisepcific_dir,'DST_*')
-    target_images_path = sorted(glob.glob(target_path))
+        # The person who provides id information (list)
+        target_id_norm_list = []
+        target_path = os.path.join(multisepcific_dir,'DST_*')
+        target_images_path = sorted(glob.glob(target_path))
 
-    for target_image_path in target_images_path:
-        img_a_whole = cv2.imread(target_image_path)
-        img_a_align_crop, _ = app.get(img_a_whole,crop_size)
-        img_a_align_crop_pil = Image.fromarray(cv2.cvtColor(img_a_align_crop[0],cv2.COLOR_BGR2RGB)) 
-        img_a = transformer_Arcface(img_a_align_crop_pil)
-        img_id = img_a.view(-1, img_a.shape[0], img_a.shape[1], img_a.shape[2])
-        # convert numpy to tensor
-        img_id = img_id.cuda()
-        #create latent id
-        img_id_downsample = F.interpolate(img_id, scale_factor=0.5)
-        latend_id = model.netArc(img_id_downsample)
-        latend_id = F.normalize(latend_id, p=2, dim=1)
-        target_id_norm_list.append(latend_id.clone())
+        for target_image_path in target_images_path:
+            img_a_whole = cv2.imread(target_image_path)
+            img_a_align_crop, _ = app.get(img_a_whole,crop_size)
+            img_a_align_crop_pil = Image.fromarray(cv2.cvtColor(img_a_align_crop[0],cv2.COLOR_BGR2RGB)) 
+            img_a = transformer_Arcface(img_a_align_crop_pil)
+            img_id = img_a.view(-1, img_a.shape[0], img_a.shape[1], img_a.shape[2])
+            # convert numpy to tensor
+            img_id = img_id.cuda()
+            #create latent id
+            img_id_downsample = F.interpolate(img_id, scale_factor=0.5)
+            latend_id = model.netArc(img_id_downsample)
+            latend_id = F.normalize(latend_id, p=2, dim=1)
+            target_id_norm_list.append(latend_id.clone())
 
-    assert len(target_id_norm_list) == len(source_specific_id_nonorm_list), "The number of images in source and target directory must be same !!!"
+        assert len(target_id_norm_list) == len(source_specific_id_nonorm_list), "The number of images in source and target directory must be same !!!"
 
-    ############## Forward Pass ######################
+        ############## Forward Pass ######################
 
-    pic_b = opt.pic_b_path
-    img_b_whole = cv2.imread(pic_b)
+        pic_b = opt.pic_b_path
+        img_b_whole = cv2.imread(pic_b)
 
-    img_b_align_crop_list, b_mat_list = app.get(img_b_whole,crop_size)
-    # detect_results = None
-    swap_result_list = []
+        img_b_align_crop_list, b_mat_list = app.get(img_b_whole,crop_size)
+        # detect_results = None
+        swap_result_list = []
 
-    id_compare_values = [] 
-    b_align_crop_tenor_list = []
-    for b_align_crop in img_b_align_crop_list:
+        id_compare_values = [] 
+        b_align_crop_tenor_list = []
+        for b_align_crop in img_b_align_crop_list:
 
-        b_align_crop_tenor = _totensor(cv2.cvtColor(b_align_crop,cv2.COLOR_BGR2RGB))[None,...].cuda()
+            b_align_crop_tenor = _totensor(cv2.cvtColor(b_align_crop,cv2.COLOR_BGR2RGB))[None,...].cuda()
 
-        b_align_crop_tenor_arcnorm = spNorm(b_align_crop_tenor)
-        b_align_crop_tenor_arcnorm_downsample = F.interpolate(b_align_crop_tenor_arcnorm, scale_factor=0.5)
-        b_align_crop_id_nonorm = model.netArc(b_align_crop_tenor_arcnorm_downsample)
+            b_align_crop_tenor_arcnorm = spNorm(b_align_crop_tenor)
+            b_align_crop_tenor_arcnorm_downsample = F.interpolate(b_align_crop_tenor_arcnorm, scale_factor=0.5)
+            b_align_crop_id_nonorm = model.netArc(b_align_crop_tenor_arcnorm_downsample)
 
-        id_compare_values.append([])
-        for source_specific_id_nonorm_tmp in source_specific_id_nonorm_list:
-            id_compare_values[-1].append(mse(b_align_crop_id_nonorm,source_specific_id_nonorm_tmp).detach().cpu().numpy())
-        b_align_crop_tenor_list.append(b_align_crop_tenor)
+            id_compare_values.append([])
+            for source_specific_id_nonorm_tmp in source_specific_id_nonorm_list:
+                id_compare_values[-1].append(mse(b_align_crop_id_nonorm,source_specific_id_nonorm_tmp).detach().cpu().numpy())
+            b_align_crop_tenor_list.append(b_align_crop_tenor)
 
-    id_compare_values_array = np.array(id_compare_values).transpose(1,0)
-    min_indexs = np.argmin(id_compare_values_array,axis=0)
-    min_value = np.min(id_compare_values_array,axis=0)
+        id_compare_values_array = np.array(id_compare_values).transpose(1,0)
+        min_indexs = np.argmin(id_compare_values_array,axis=0)
+        min_value = np.min(id_compare_values_array,axis=0)
 
-    swap_result_list = [] 
-    swap_result_matrix_list = []
+        swap_result_list = [] 
+        swap_result_matrix_list = []
+        swap_result_ori_pic_list = []
 
-    for tmp_index, min_index in enumerate(min_indexs):
-        if min_value[tmp_index] < opt.id_thres:
-            swap_result = model(None, b_align_crop_tenor_list[tmp_index], target_id_norm_list[min_index], None, True)[0]
-            swap_result_list.append(swap_result)
-            swap_result_matrix_list.append(b_mat_list[tmp_index])
+        for tmp_index, min_index in enumerate(min_indexs):
+            if min_value[tmp_index] < opt.id_thres:
+                swap_result = model(None, b_align_crop_tenor_list[tmp_index], target_id_norm_list[min_index], None, True)[0]
+                swap_result_list.append(swap_result)
+                swap_result_matrix_list.append(b_mat_list[tmp_index])
+                swap_result_ori_pic_list.append(b_align_crop_tenor_list[tmp_index])
+            else:
+                pass
+
+        if len(swap_result_list) !=0:
+
+            if opt.use_mask:
+                n_classes = 19
+                net = BiSeNet(n_classes=n_classes)
+                net.cuda()
+                save_pth = os.path.join('./parsing_model/checkpoint', '79999_iter.pth')
+                net.load_state_dict(torch.load(save_pth))
+                net.eval()
+            else:
+                net =None
+        
+            reverse2wholeimage(swap_result_ori_pic_list, swap_result_list, swap_result_matrix_list, crop_size, img_b_whole, logoclass,\
+                os.path.join(opt.output_path, 'result_whole_swap_multispecific.jpg'), opt.no_simswaplogo,pasring_model =net,use_mask=opt.use_mask, norm = spNorm)
+
+            print(' ')
+
+            print('************ Done ! ************')
+        
         else:
-            pass
-
-    if len(swap_result_list) !=0:
-    
-        reverse2wholeimage(swap_result_list, swap_result_matrix_list, crop_size, img_b_whole, logoclass, os.path.join(opt.output_path, 'result_whole_swap_multispecific.jpg'), opt.no_simswaplogo)
-
-        print(' ')
-
-        print('************ Done ! ************')
-    
-    else:
-        print('The people you specified are not found on the picture: {}'.format(pic_b))
+            print('The people you specified are not found on the picture: {}'.format(pic_b))

test_wholeimage_swapmulti.py

@@ -12,6 +12,8 @@ from insightface_func.face_detect_crop_multi import Face_detect_crop
 from util.reverse2original import reverse2wholeimage
 import os
 from util.add_watermark import watermark_image
+from util.norm import SpecificNorm
+from parsing_model.model import BiSeNet
 
 def lcm(a, b): return abs(a * b) / fractions.gcd(a, b) if a and b else 0
 
@@ -35,46 +37,60 @@ if __name__ == '__main__':
     logoclass = watermark_image('./simswaplogo/simswaplogo.png')
     model = create_model(opt)
     model.eval()
-
+    spNorm =SpecificNorm()
 
     app = Face_detect_crop(name='antelope', root='./insightface_func/models')
     app.prepare(ctx_id= 0, det_thresh=0.6, det_size=(640,640))
 
-    pic_a = opt.pic_a_path
+    with torch.no_grad():
+        pic_a = opt.pic_a_path
 
-    img_a_whole = cv2.imread(pic_a)
-    img_a_align_crop, _ = app.get(img_a_whole,crop_size)
-    img_a_align_crop_pil = Image.fromarray(cv2.cvtColor(img_a_align_crop[0],cv2.COLOR_BGR2RGB)) 
-    img_a = transformer_Arcface(img_a_align_crop_pil)
-    img_id = img_a.view(-1, img_a.shape[0], img_a.shape[1], img_a.shape[2])
+        img_a_whole = cv2.imread(pic_a)
+        img_a_align_crop, _ = app.get(img_a_whole,crop_size)
+        img_a_align_crop_pil = Image.fromarray(cv2.cvtColor(img_a_align_crop[0],cv2.COLOR_BGR2RGB)) 
+        img_a = transformer_Arcface(img_a_align_crop_pil)
+        img_id = img_a.view(-1, img_a.shape[0], img_a.shape[1], img_a.shape[2])
 
-    # convert numpy to tensor
-    img_id = img_id.cuda()
+        # convert numpy to tensor
+        img_id = img_id.cuda()
 
-    #create latent id
-    img_id_downsample = F.interpolate(img_id, scale_factor=0.5)
-    latend_id = model.netArc(img_id_downsample)
-    latend_id = F.normalize(latend_id, p=2, dim=1)
+        #create latent id
+        img_id_downsample = F.interpolate(img_id, scale_factor=0.5)
+        latend_id = model.netArc(img_id_downsample)
+        latend_id = F.normalize(latend_id, p=2, dim=1)
 
 
-    ############## Forward Pass ######################
+        ############## Forward Pass ######################
 
-    pic_b = opt.pic_b_path
-    img_b_whole = cv2.imread(pic_b)
+        pic_b = opt.pic_b_path
+        img_b_whole = cv2.imread(pic_b)
 
-    img_b_align_crop_list, b_mat_list = app.get(img_b_whole,crop_size)
-    # detect_results = None
-    swap_result_list = []
+        img_b_align_crop_list, b_mat_list = app.get(img_b_whole,crop_size)
+        # detect_results = None
+        swap_result_list = []
+        b_align_crop_tenor_list = []
 
-    for b_align_crop in img_b_align_crop_list:
+        for b_align_crop in img_b_align_crop_list:
 
-        b_align_crop_tenor = _totensor(cv2.cvtColor(b_align_crop,cv2.COLOR_BGR2RGB))[None,...].cuda()
+            b_align_crop_tenor = _totensor(cv2.cvtColor(b_align_crop,cv2.COLOR_BGR2RGB))[None,...].cuda()
 
-        swap_result = model(None, b_align_crop_tenor, latend_id, None, True)[0]
-        swap_result_list.append(swap_result)
+            swap_result = model(None, b_align_crop_tenor, latend_id, None, True)[0]
+            swap_result_list.append(swap_result)
+            b_align_crop_tenor_list.append(b_align_crop_tenor)
 
 
-    reverse2wholeimage(swap_result_list, b_mat_list, crop_size, img_b_whole, logoclass, os.path.join(opt.output_path, 'result_whole_swapmulti.jpg'),opt.no_simswaplogo)
-    print(' ')
+        if opt.use_mask:
+            n_classes = 19
+            net = BiSeNet(n_classes=n_classes)
+            net.cuda()
+            save_pth = os.path.join('./parsing_model/checkpoint', '79999_iter.pth')
+            net.load_state_dict(torch.load(save_pth))
+            net.eval()
+        else:
+            net =None
 
-    print('************ Done ! ************')
+        reverse2wholeimage(b_align_crop_tenor_list,swap_result_list, b_mat_list, crop_size, img_b_whole, logoclass, \
+            os.path.join(opt.output_path, 'result_whole_swapmulti.jpg'),opt.no_simswaplogo,pasring_model =net,use_mask=opt.use_mask, norm = spNorm)
+        print(' ')
+
+        print('************ Done ! ************')

test_wholeimage_swapsingle.py

@@ -12,6 +12,8 @@ from insightface_func.face_detect_crop_single import Face_detect_crop
 from util.reverse2original import reverse2wholeimage
 import os
 from util.add_watermark import watermark_image
+from util.norm import SpecificNorm
+from parsing_model.model import BiSeNet
 
 def lcm(a, b): return abs(a * b) / fractions.gcd(a, b) if a and b else 0
 
@@ -35,45 +37,60 @@ if __name__ == '__main__':
     model = create_model(opt)
     model.eval()
 
-
+    spNorm =SpecificNorm()
     app = Face_detect_crop(name='antelope', root='./insightface_func/models')
     app.prepare(ctx_id= 0, det_thresh=0.6, det_size=(640,640))
 
-    pic_a = opt.pic_a_path
+    with torch.no_grad():
+        pic_a = opt.pic_a_path
 
-    img_a_whole = cv2.imread(pic_a)
-    img_a_align_crop, _ = app.get(img_a_whole,crop_size)
-    img_a_align_crop_pil = Image.fromarray(cv2.cvtColor(img_a_align_crop[0],cv2.COLOR_BGR2RGB)) 
-    img_a = transformer_Arcface(img_a_align_crop_pil)
-    img_id = img_a.view(-1, img_a.shape[0], img_a.shape[1], img_a.shape[2])
+        img_a_whole = cv2.imread(pic_a)
+        img_a_align_crop, _ = app.get(img_a_whole,crop_size)
+        img_a_align_crop_pil = Image.fromarray(cv2.cvtColor(img_a_align_crop[0],cv2.COLOR_BGR2RGB)) 
+        img_a = transformer_Arcface(img_a_align_crop_pil)
+        img_id = img_a.view(-1, img_a.shape[0], img_a.shape[1], img_a.shape[2])
 
-    # convert numpy to tensor
-    img_id = img_id.cuda()
+        # convert numpy to tensor
+        img_id = img_id.cuda()
 
-    #create latent id
-    img_id_downsample = F.interpolate(img_id, scale_factor=0.5)
-    latend_id = model.netArc(img_id_downsample)
-    latend_id = F.normalize(latend_id, p=2, dim=1)
+        #create latent id
+        img_id_downsample = F.interpolate(img_id, scale_factor=0.5)
+        latend_id = model.netArc(img_id_downsample)
+        latend_id = F.normalize(latend_id, p=2, dim=1)
 
 
-    ############## Forward Pass ######################
+        ############## Forward Pass ######################
 
-    pic_b = opt.pic_b_path
-    img_b_whole = cv2.imread(pic_b)
+        pic_b = opt.pic_b_path
+        img_b_whole = cv2.imread(pic_b)
 
-    img_b_align_crop_list, b_mat_list = app.get(img_b_whole,crop_size)
-    # detect_results = None
-    swap_result_list = []
+        img_b_align_crop_list, b_mat_list = app.get(img_b_whole,crop_size)
+        # detect_results = None
+        swap_result_list = []
 
-    for b_align_crop in img_b_align_crop_list:
+        b_align_crop_tenor_list = []
 
-        b_align_crop_tenor = _totensor(cv2.cvtColor(b_align_crop,cv2.COLOR_BGR2RGB))[None,...].cuda()
+        for b_align_crop in img_b_align_crop_list:
 
-        swap_result = model(None, b_align_crop_tenor, latend_id, None, True)[0]
-        swap_result_list.append(swap_result)
+            b_align_crop_tenor = _totensor(cv2.cvtColor(b_align_crop,cv2.COLOR_BGR2RGB))[None,...].cuda()
 
-    reverse2wholeimage(swap_result_list, b_mat_list, crop_size, img_b_whole, logoclass, os.path.join(opt.output_path, 'result_whole_swapsingle.jpg'), opt.no_simswaplogo)
+            swap_result = model(None, b_align_crop_tenor, latend_id, None, True)[0]
+            swap_result_list.append(swap_result)
+            b_align_crop_tenor_list.append(b_align_crop_tenor)
 
-    print(' ')
+        if opt.use_mask:
+            n_classes = 19
+            net = BiSeNet(n_classes=n_classes)
+            net.cuda()
+            save_pth = os.path.join('./parsing_model/checkpoint', '79999_iter.pth')
+            net.load_state_dict(torch.load(save_pth))
+            net.eval()
+        else:
+            net =None
 
-    print('************ Done ! ************')
+        reverse2wholeimage(b_align_crop_tenor_list, swap_result_list, b_mat_list, crop_size, img_b_whole, logoclass, \
+            os.path.join(opt.output_path, 'result_whole_swapsingle.jpg'), opt.no_simswaplogo,pasring_model =net,use_mask=opt.use_mask, norm = spNorm)
+
+        print(' ')
+
+        print('************ Done ! ************')

test_wholeimage_swapspecific.py

@@ -14,6 +14,7 @@ import os
 from util.add_watermark import watermark_image
 import torch.nn as nn
 from util.norm import SpecificNorm
+from parsing_model.model import BiSeNet
 
 def lcm(a, b): return abs(a * b) / fractions.gcd(a, b) if a and b else 0
 
@@ -110,11 +111,22 @@ if __name__ == '__main__':
     min_index = np.argmin(id_compare_values_array)
     min_value = id_compare_values_array[min_index]
 
+    if opt.use_mask:
+        n_classes = 19
+        net = BiSeNet(n_classes=n_classes)
+        net.cuda()
+        save_pth = os.path.join('./parsing_model/checkpoint', '79999_iter.pth')
+        net.load_state_dict(torch.load(save_pth))
+        net.eval()
+    else:
+        net =None
+
     if min_value < opt.id_thres:
 
         swap_result = model(None, b_align_crop_tenor_list[min_index], latend_id, None, True)[0]
 
-        reverse2wholeimage([swap_result], [b_mat_list[min_index]], crop_size, img_b_whole, logoclass, os.path.join(opt.output_path, 'result_whole_swapspecific.jpg'), opt.no_simswaplogo)
+        reverse2wholeimage([b_align_crop_tenor_list[min_index]], [swap_result], [b_mat_list[min_index]], crop_size, img_b_whole, logoclass, \
+            os.path.join(opt.output_path, 'result_whole_swapspecific.jpg'), opt.no_simswaplogo,pasring_model =net,use_mask=opt.use_mask, norm = spNorm)
 
         print(' ')
 

util/reverse2original.py

@@ -1,13 +1,91 @@
 import cv2
 import numpy as np
 # import  time
-from util.add_watermark import  watermark_image
+import torch
+from torch.nn import functional as F
+import torch.nn as nn
 
-def reverse2wholeimage(swaped_imgs, mats, crop_size, oriimg, logoclass, save_path = '', no_simswaplogo = False):
+
+def encode_segmentation_rgb(segmentation, no_neck=True):
+    parse = segmentation
+
+    face_part_ids = [1, 2, 3, 4, 5, 6, 10, 12, 13] if no_neck else [1, 2, 3, 4, 5, 6, 7, 8, 10, 12, 13, 14]
+    mouth_id = 11
+    # hair_id = 17
+    face_map = np.zeros([parse.shape[0], parse.shape[1]])
+    mouth_map = np.zeros([parse.shape[0], parse.shape[1]])
+    # hair_map = np.zeros([parse.shape[0], parse.shape[1]])
+
+    for valid_id in face_part_ids:
+        valid_index = np.where(parse==valid_id)
+        face_map[valid_index] = 255
+    valid_index = np.where(parse==mouth_id)
+    mouth_map[valid_index] = 255
+    # valid_index = np.where(parse==hair_id)
+    # hair_map[valid_index] = 255
+    #return np.stack([face_map, mouth_map,hair_map], axis=2)
+    return np.stack([face_map, mouth_map], axis=2)
+
+
+class SoftErosion(nn.Module):
+    def __init__(self, kernel_size=15, threshold=0.6, iterations=1):
+        super(SoftErosion, self).__init__()
+        r = kernel_size // 2
+        self.padding = r
+        self.iterations = iterations
+        self.threshold = threshold
+
+        # Create kernel
+        y_indices, x_indices = torch.meshgrid(torch.arange(0., kernel_size), torch.arange(0., kernel_size))
+        dist = torch.sqrt((x_indices - r) ** 2 + (y_indices - r) ** 2)
+        kernel = dist.max() - dist
+        kernel /= kernel.sum()
+        kernel = kernel.view(1, 1, *kernel.shape)
+        self.register_buffer('weight', kernel)
+
+    def forward(self, x):
+        x = x.float()
+        for i in range(self.iterations - 1):
+            x = torch.min(x, F.conv2d(x, weight=self.weight, groups=x.shape[1], padding=self.padding))
+        x = F.conv2d(x, weight=self.weight, groups=x.shape[1], padding=self.padding)
+
+        mask = x >= self.threshold
+        x[mask] = 1.0
+        x[~mask] /= x[~mask].max()
+
+        return x, mask
+
+
+def postprocess(swapped_face, target, target_mask,smooth_mask):
+    # target_mask = cv2.resize(target_mask, (self.size,  self.size))
+
+    mask_tensor = torch.from_numpy(target_mask.copy().transpose((2, 0, 1))).float().mul_(1/255.0).cuda()
+    face_mask_tensor = mask_tensor[0] + mask_tensor[1]
+    
+    soft_face_mask_tensor, _ = smooth_mask(face_mask_tensor.unsqueeze_(0).unsqueeze_(0))
+    soft_face_mask_tensor.squeeze_()
+
+    soft_face_mask = soft_face_mask_tensor.cpu().numpy()
+    soft_face_mask = soft_face_mask[:, :, np.newaxis]
+
+    result =  swapped_face * soft_face_mask + target * (1 - soft_face_mask)
+    result = result[:,:,::-1]# .astype(np.uint8)
+    return result
+
+def reverse2wholeimage(b_align_crop_tenor_list,swaped_imgs, mats, crop_size, oriimg, logoclass, save_path = '', \
+                    no_simswaplogo = False,pasring_model =None,norm = None, use_mask = False):
 
     target_image_list = []
     img_mask_list = []
-    for swaped_img, mat in zip(swaped_imgs, mats):
+    if use_mask:
+        smooth_mask = SoftErosion(kernel_size=17, threshold=0.9, iterations=7).cuda()
+    else:
+        pass
+
+    # print(len(swaped_imgs))
+    # print(mats)
+    # print(len(b_align_crop_tenor_list))
+    for swaped_img, mat ,source_img in zip(swaped_imgs, mats,b_align_crop_tenor_list):
         swaped_img = swaped_img.cpu().detach().numpy().transpose((1, 2, 0))
         img_white = np.full((crop_size,crop_size), 255, dtype=float)
 
@@ -23,7 +101,28 @@ def reverse2wholeimage(swaped_imgs, mats, crop_size, oriimg, logoclass, save_pat
         mat_rev[1][2] = -(mat[0][2]*mat[1][0]-mat[0][0]*mat[1][2])/div2
 
         orisize = (oriimg.shape[1], oriimg.shape[0])
-        target_image = cv2.warpAffine(swaped_img, mat_rev, orisize)
+        if use_mask:
+            source_img_norm = norm(source_img)
+            source_img_512  = F.interpolate(source_img_norm,size=(512,512))
+            out = pasring_model(source_img_512)[0]
+            parsing = out.squeeze(0).detach().cpu().numpy().argmax(0)
+            vis_parsing_anno = parsing.copy().astype(np.uint8)
+            tgt_mask = encode_segmentation_rgb(vis_parsing_anno)
+            if tgt_mask.sum() >= 5000:
+                # face_mask_tensor = tgt_mask[...,0] + tgt_mask[...,1]
+                target_mask = cv2.resize(tgt_mask, (224,  224))
+                # print(source_img)
+                target_image_parsing = postprocess(swaped_img, source_img[0].cpu().detach().numpy().transpose((1, 2, 0)), target_mask,smooth_mask)
+                
+
+                target_image = cv2.warpAffine(target_image_parsing, mat_rev, orisize)
+                # target_image_parsing = cv2.warpAffine(swaped_img, mat_rev, orisize)
+            else:
+                target_image = cv2.warpAffine(swaped_img, mat_rev, orisize)[..., ::-1]
+        else:
+            target_image = cv2.warpAffine(swaped_img, mat_rev, orisize)
+        # source_image   = cv2.warpAffine(source_img, mat_rev, orisize)
+
         img_white = cv2.warpAffine(img_white, mat_rev, orisize)
 
 
@@ -31,16 +130,39 @@ def reverse2wholeimage(swaped_imgs, mats, crop_size, oriimg, logoclass, save_pat
 
         img_mask = img_white
 
-        kernel = np.ones((10,10),np.uint8)
+        # if use_mask:
+        #     kernel = np.ones((40,40),np.uint8)
+        #     img_mask = cv2.erode(img_mask,kernel,iterations = 1)
+        # else:
+        kernel = np.ones((40,40),np.uint8)
         img_mask = cv2.erode(img_mask,kernel,iterations = 1)
+        kernel_size = (20, 20)
+        blur_size = tuple(2*i+1 for i in kernel_size)
+        img_mask = cv2.GaussianBlur(img_mask, blur_size, 0)
+
+        # kernel = np.ones((10,10),np.uint8)
+        # img_mask = cv2.erode(img_mask,kernel,iterations = 1)
+
+
 
         img_mask /= 255
 
         img_mask = np.reshape(img_mask, [img_mask.shape[0],img_mask.shape[1],1])
-        target_image = np.array(target_image, dtype=np.float)[..., ::-1] * 255
+
+        # pasing mask
+
+        # target_image_parsing = postprocess(target_image, source_image, tgt_mask)
+
+        if use_mask:
+            target_image = np.array(target_image, dtype=np.float) * 255
+        else:
+            target_image = np.array(target_image, dtype=np.float)[..., ::-1] * 255
+
 
         img_mask_list.append(img_mask)
         target_image_list.append(target_image)
+        
+
     # target_image /= 255
     # target_image = 0
     img = np.array(oriimg, dtype=np.float)
@@ -51,8 +173,3 @@ def reverse2wholeimage(swaped_imgs, mats, crop_size, oriimg, logoclass, save_pat
     if not no_simswaplogo:
         final_img = logoclass.apply_frames(final_img)
     cv2.imwrite(save_path, final_img)
-
-    # cv2.imwrite('E:\\lny\\SimSwap-main\\output\\img_div.jpg', img * 255)
-    # cv2.imwrite('E:\\lny\\SimSwap-main\\output\\ori_img.jpg', oriimg)
-    
-    

util/videoswap.py

@@ -11,14 +11,15 @@ from moviepy.editor import AudioFileClip, VideoFileClip
 from moviepy.video.io.ImageSequenceClip import ImageSequenceClip
 import  time
 from util.add_watermark import watermark_image
-
+from util.norm import SpecificNorm
+from parsing_model.model import BiSeNet
 
 def _totensor(array):
     tensor = torch.from_numpy(array)
     img = tensor.transpose(0, 1).transpose(0, 2).contiguous()
     return img.float().div(255)
 
-def video_swap(video_path, id_vetor, swap_model, detect_model, save_path, temp_results_dir='./temp_results', crop_size=224, no_simswaplogo = False):
+def video_swap(video_path, id_vetor, swap_model, detect_model, save_path, temp_results_dir='./temp_results', crop_size=224, no_simswaplogo = False,use_mask =False):
     video_forcheck = VideoFileClip(video_path)
     if video_forcheck.audio is None:
         no_audio = True
@@ -45,6 +46,17 @@ def video_swap(video_path, id_vetor, swap_model, detect_model, save_path, temp_r
     if  os.path.exists(temp_results_dir):
             shutil.rmtree(temp_results_dir)
 
+    spNorm =SpecificNorm()
+    if use_mask:
+        n_classes = 19
+        net = BiSeNet(n_classes=n_classes)
+        net.cuda()
+        save_pth = os.path.join('./parsing_model/checkpoint', '79999_iter.pth')
+        net.load_state_dict(torch.load(save_pth))
+        net.eval()
+    else:
+        net =None
+
     # while ret:
     for frame_index in tqdm(range(frame_count)): 
         ret, frame = video.read()
@@ -58,7 +70,7 @@ def video_swap(video_path, id_vetor, swap_model, detect_model, save_path, temp_r
                 frame_align_crop_list = detect_results[0]
                 frame_mat_list = detect_results[1]
                 swap_result_list = []
-
+                frame_align_crop_tenor_list = []
                 for frame_align_crop in frame_align_crop_list:
 
                     # BGR TO RGB
@@ -68,10 +80,12 @@ def video_swap(video_path, id_vetor, swap_model, detect_model, save_path, temp_r
 
                     swap_result = swap_model(None, frame_align_crop_tenor, id_vetor, None, True)[0]
                     swap_result_list.append(swap_result)
+                    frame_align_crop_tenor_list.append(frame_align_crop_tenor)
 
                     
 
-                reverse2wholeimage(swap_result_list, frame_mat_list, crop_size, frame, logoclass,os.path.join(temp_results_dir, 'frame_{:0>7d}.jpg'.format(frame_index)),no_simswaplogo)
+                reverse2wholeimage(frame_align_crop_tenor_list,swap_result_list, frame_mat_list, crop_size, frame, logoclass,\
+                    os.path.join(temp_results_dir, 'frame_{:0>7d}.jpg'.format(frame_index)),no_simswaplogo,pasring_model =net,use_mask=use_mask, norm = spNorm)
 
             else:
                 if not os.path.exists(temp_results_dir):
@@ -95,5 +109,5 @@ def video_swap(video_path, id_vetor, swap_model, detect_model, save_path, temp_r
         clips = clips.set_audio(video_audio_clip)
 
 
-    clips.write_videofile(save_path)
+    clips.write_videofile(save_path,audio_codec='aac')
 

util/videoswap_multispecific.py

@@ -13,13 +13,14 @@ import  time
 from util.add_watermark import watermark_image
 from util.norm import SpecificNorm
 import torch.nn.functional as F
+from parsing_model.model import BiSeNet
 
 def _totensor(array):
     tensor = torch.from_numpy(array)
     img = tensor.transpose(0, 1).transpose(0, 2).contiguous()
     return img.float().div(255)
 
-def video_swap(video_path, target_id_norm_list,source_specific_id_nonorm_list,id_thres, swap_model, detect_model, save_path, temp_results_dir='./temp_results', crop_size=224, no_simswaplogo = False):
+def video_swap(video_path, target_id_norm_list,source_specific_id_nonorm_list,id_thres, swap_model, detect_model, save_path, temp_results_dir='./temp_results', crop_size=224, no_simswaplogo = False,use_mask =False):
     video_forcheck = VideoFileClip(video_path)
     if video_forcheck.audio is None:
         no_audio = True
@@ -49,6 +50,16 @@ def video_swap(video_path, target_id_norm_list,source_specific_id_nonorm_list,id
     spNorm =SpecificNorm()
     mse = torch.nn.MSELoss().cuda()
 
+    if use_mask:
+        n_classes = 19
+        net = BiSeNet(n_classes=n_classes)
+        net.cuda()
+        save_pth = os.path.join('./parsing_model/checkpoint', '79999_iter.pth')
+        net.load_state_dict(torch.load(save_pth))
+        net.eval()
+    else:
+        net =None
+
     # while ret:
     for frame_index in tqdm(range(frame_count)): 
         ret, frame = video.read()
@@ -85,12 +96,13 @@ def video_swap(video_path, target_id_norm_list,source_specific_id_nonorm_list,id
 
                 swap_result_list = [] 
                 swap_result_matrix_list = []
-
+                swap_result_ori_pic_list = []
                 for tmp_index, min_index in enumerate(min_indexs):
                     if min_value[tmp_index] < id_thres:
                         swap_result = swap_model(None, frame_align_crop_tenor_list[tmp_index], target_id_norm_list[min_index], None, True)[0]
                         swap_result_list.append(swap_result)
                         swap_result_matrix_list.append(frame_mat_list[tmp_index])
+                        swap_result_ori_pic_list.append(frame_align_crop_tenor_list[tmp_index])
                     else:
                         pass
 
@@ -98,7 +110,8 @@ def video_swap(video_path, target_id_norm_list,source_specific_id_nonorm_list,id
 
                 if len(swap_result_list) !=0:
                     
-                    reverse2wholeimage(swap_result_list, swap_result_matrix_list, crop_size, frame, logoclass,os.path.join(temp_results_dir, 'frame_{:0>7d}.jpg'.format(frame_index)),no_simswaplogo)
+                    reverse2wholeimage(swap_result_ori_pic_list,swap_result_list, swap_result_matrix_list, crop_size, frame, logoclass,\
+                        os.path.join(temp_results_dir, 'frame_{:0>7d}.jpg'.format(frame_index)),no_simswaplogo,pasring_model =net,use_mask=use_mask, norm = spNorm)
                 else:
                     if not os.path.exists(temp_results_dir):
                         os.mkdir(temp_results_dir)
@@ -129,5 +142,5 @@ def video_swap(video_path, target_id_norm_list,source_specific_id_nonorm_list,id
         clips = clips.set_audio(video_audio_clip)
 
 
-    clips.write_videofile(save_path)
+    clips.write_videofile(save_path,audio_codec='aac')
 

util/videoswap_specific.py

@@ -13,13 +13,14 @@ import  time
 from util.add_watermark import watermark_image
 from util.norm import SpecificNorm
 import torch.nn.functional as F
+from parsing_model.model import BiSeNet
 
 def _totensor(array):
     tensor = torch.from_numpy(array)
     img = tensor.transpose(0, 1).transpose(0, 2).contiguous()
     return img.float().div(255)
 
-def video_swap(video_path, id_vetor,specific_person_id_nonorm,id_thres, swap_model, detect_model, save_path, temp_results_dir='./temp_results', crop_size=224, no_simswaplogo = False):
+def video_swap(video_path, id_vetor,specific_person_id_nonorm,id_thres, swap_model, detect_model, save_path, temp_results_dir='./temp_results', crop_size=224, no_simswaplogo = False,use_mask =False):
     video_forcheck = VideoFileClip(video_path)
     if video_forcheck.audio is None:
         no_audio = True
@@ -49,6 +50,16 @@ def video_swap(video_path, id_vetor,specific_person_id_nonorm,id_thres, swap_mod
     spNorm =SpecificNorm()
     mse = torch.nn.MSELoss().cuda()
 
+    if use_mask:
+        n_classes = 19
+        net = BiSeNet(n_classes=n_classes)
+        net.cuda()
+        save_pth = os.path.join('./parsing_model/checkpoint', '79999_iter.pth')
+        net.load_state_dict(torch.load(save_pth))
+        net.eval()
+    else:
+        net =None
+
     # while ret:
     for frame_index in tqdm(range(frame_count)): 
         ret, frame = video.read()
@@ -83,7 +94,8 @@ def video_swap(video_path, id_vetor,specific_person_id_nonorm,id_thres, swap_mod
                 if min_value < id_thres:
                     swap_result = swap_model(None, frame_align_crop_tenor_list[min_index], id_vetor, None, True)[0]
                 
-                    reverse2wholeimage([swap_result], [frame_mat_list[min_index]], crop_size, frame, logoclass,os.path.join(temp_results_dir, 'frame_{:0>7d}.jpg'.format(frame_index)),no_simswaplogo)
+                    reverse2wholeimage([frame_align_crop_tenor_list[min_index]], [swap_result], [frame_mat_list[min_index]], crop_size, frame, logoclass,\
+                        os.path.join(temp_results_dir, 'frame_{:0>7d}.jpg'.format(frame_index)),no_simswaplogo,pasring_model =net,use_mask= use_mask, norm = spNorm)
                 else:
                     if not os.path.exists(temp_results_dir):
                         os.mkdir(temp_results_dir)
@@ -114,5 +126,5 @@ def video_swap(video_path, id_vetor,specific_person_id_nonorm,id_thres, swap_mod
         clips = clips.set_audio(video_audio_clip)
 
 
-    clips.write_videofile(save_path)
+    clips.write_videofile(save_path,audio_codec='aac')