diff --git a/.gitignore b/.gitignore index fd548c0..5b6e4b9 100644 --- a/.gitignore +++ b/.gitignore @@ -135,4 +135,8 @@ checkpoints/ *.zip *.avi *.pdf -*.pptx \ No newline at end of file +*.pptx + +*.pth +*.onnx +wandb/ \ No newline at end of file diff --git a/MultiSpecific.ipynb b/MultiSpecific.ipynb index b7b0fd6..ef13d1c 100644 --- a/MultiSpecific.ipynb +++ b/MultiSpecific.ipynb @@ -1 +1 @@ -{"nbformat":4,"nbformat_minor":0,"metadata":{"colab":{"name":"MultiSpecific.ipynb","provenance":[],"collapsed_sections":[],"authorship_tag":"ABX9TyNw8SfPWhG77cf/e7YZd178"},"kernelspec":{"name":"python3","display_name":"Python 3"},"language_info":{"name":"python"},"accelerator":"GPU"},"cells":[{"cell_type":"markdown","metadata":{"id":"7_gtFoV8BuRx"},"source":["This is an example of SimSwap on processing video with multiple faces with designated sources.\n","\n","Code path: https://github.com/neuralchen/SimSwap\n","Paper path: https://arxiv.org/pdf/2106.06340v1.pdf."]},{"cell_type":"code","metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"0Y1RfpzsCAl9","executionInfo":{"status":"ok","timestamp":1625380781426,"user_tz":-480,"elapsed":586,"user":{"displayName":"José Lampreia","photoUrl":"","userId":"16015278604201270582"}},"outputId":"2a897b34-72f1-4515-ac6f-2f0e2d4ea4f7"},"source":["## make sure you are using a runtime with GPU\n","## you can check at Runtime/Change runtime type in the top bar.\n","!nvidia-smi"],"execution_count":1,"outputs":[{"output_type":"stream","text":["Sun Jul 4 06:39:39 2021 \n","+-----------------------------------------------------------------------------+\n","| NVIDIA-SMI 465.27 Driver Version: 460.32.03 CUDA Version: 11.2 |\n","|-------------------------------+----------------------+----------------------+\n","| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |\n","| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |\n","| | | MIG M. |\n","|===============================+======================+======================|\n","| 0 Tesla T4 Off | 00000000:00:04.0 Off | 0 |\n","| N/A 45C P8 9W / 70W | 0MiB / 15109MiB | 0% Default |\n","| | | N/A |\n","+-------------------------------+----------------------+----------------------+\n"," \n","+-----------------------------------------------------------------------------+\n","| Processes: |\n","| GPU GI CI PID Type Process name GPU Memory |\n","| ID ID Usage |\n","|=============================================================================|\n","| No running processes found |\n","+-----------------------------------------------------------------------------+\n"],"name":"stdout"}]},{"cell_type":"markdown","metadata":{"id":"0Qzzx2UpDkqw"},"source":["All file changes make by this notebook are temporary. \n","You can try to mount your own google drive to store files if you wang.\n"]},{"cell_type":"code","metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"VA_4CeWZCHLP","executionInfo":{"status":"ok","timestamp":1625380786661,"user_tz":-480,"elapsed":4693,"user":{"displayName":"José Lampreia","photoUrl":"","userId":"16015278604201270582"}},"outputId":"d0665552-be92-45b0-aab2-f84c619a51fb"},"source":["!git clone https://github.com/neuralchen/SimSwap\n","!cd SimSwap && git pull"],"execution_count":2,"outputs":[{"output_type":"stream","text":["Cloning into 'SimSwap'...\n","remote: Enumerating objects: 667, done.\u001b[K\n","remote: Counting objects: 100% (48/48), done.\u001b[K\n","remote: Compressing objects: 100% (35/35), done.\u001b[K\n","remote: Total 667 (delta 19), reused 28 (delta 13), pack-reused 619\u001b[K\n","Receiving objects: 100% (667/667), 132.14 MiB | 44.44 MiB/s, done.\n","Resolving deltas: 100% (292/292), done.\n","Already up to date.\n"],"name":"stdout"}]},{"cell_type":"code","metadata":{"id":"Y5K4au_UCkKn","colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"status":"ok","timestamp":1625380797906,"user_tz":-480,"elapsed":11253,"user":{"displayName":"José Lampreia","photoUrl":"","userId":"16015278604201270582"}},"outputId":"7429f153-bc6d-48c2-eb3c-21f1f02fede9"},"source":["!pip install insightface==0.2.1 onnxruntime moviepy\n","!pip install googledrivedownloader\n","!pip install imageio==2.4.1"],"execution_count":3,"outputs":[{"output_type":"stream","text":["Collecting insightface==0.2.1\n"," Downloading https://files.pythonhosted.org/packages/ee/1e/6395bbe0db665f187c8e49266cda54fcf661f182192370d409423e4943e4/insightface-0.2.1-py2.py3-none-any.whl\n","Collecting onnxruntime\n","\u001b[?25l Downloading https://files.pythonhosted.org/packages/f9/76/3d0f8bb2776961c7335693df06eccf8d099e48fa6fb552c7546867192603/onnxruntime-1.8.0-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (4.5MB)\n","\u001b[K |████████████████████████████████| 4.5MB 37.5MB/s \n","\u001b[?25hRequirement already satisfied: moviepy in /usr/local/lib/python3.7/dist-packages (0.2.3.5)\n","Requirement already satisfied: scikit-learn in /usr/local/lib/python3.7/dist-packages (from insightface==0.2.1) (0.22.2.post1)\n","Requirement already satisfied: requests in /usr/local/lib/python3.7/dist-packages (from insightface==0.2.1) (2.23.0)\n","Collecting onnx\n","\u001b[?25l Downloading https://files.pythonhosted.org/packages/3f/9b/54c950d3256e27f970a83cd0504efb183a24312702deed0179453316dbd0/onnx-1.9.0-cp37-cp37m-manylinux2010_x86_64.whl (12.2MB)\n","\u001b[K |████████████████████████████████| 12.2MB 32.2MB/s \n","\u001b[?25hRequirement already satisfied: matplotlib in /usr/local/lib/python3.7/dist-packages (from insightface==0.2.1) (3.2.2)\n","Requirement already satisfied: Pillow in /usr/local/lib/python3.7/dist-packages (from insightface==0.2.1) (7.1.2)\n","Requirement already satisfied: scikit-image in /usr/local/lib/python3.7/dist-packages (from insightface==0.2.1) (0.16.2)\n","Requirement already satisfied: opencv-python in /usr/local/lib/python3.7/dist-packages (from insightface==0.2.1) (4.1.2.30)\n","Requirement already satisfied: tqdm in /usr/local/lib/python3.7/dist-packages (from insightface==0.2.1) (4.41.1)\n","Requirement already satisfied: scipy in /usr/local/lib/python3.7/dist-packages (from insightface==0.2.1) (1.4.1)\n","Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from insightface==0.2.1) (1.19.5)\n","Requirement already satisfied: easydict in /usr/local/lib/python3.7/dist-packages (from insightface==0.2.1) (1.9)\n","Requirement already satisfied: flatbuffers in /usr/local/lib/python3.7/dist-packages (from onnxruntime) (1.12)\n","Requirement already satisfied: protobuf in /usr/local/lib/python3.7/dist-packages (from onnxruntime) (3.12.4)\n","Requirement already satisfied: imageio<3.0,>=2.1.2 in /usr/local/lib/python3.7/dist-packages (from moviepy) (2.4.1)\n","Requirement already satisfied: decorator<5.0,>=4.0.2 in /usr/local/lib/python3.7/dist-packages (from moviepy) (4.4.2)\n","Requirement already satisfied: joblib>=0.11 in /usr/local/lib/python3.7/dist-packages (from scikit-learn->insightface==0.2.1) (1.0.1)\n","Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests->insightface==0.2.1) (3.0.4)\n","Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests->insightface==0.2.1) (1.24.3)\n","Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests->insightface==0.2.1) (2021.5.30)\n","Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests->insightface==0.2.1) (2.10)\n","Requirement already satisfied: typing-extensions>=3.6.2.1 in /usr/local/lib/python3.7/dist-packages (from onnx->insightface==0.2.1) (3.7.4.3)\n","Requirement already satisfied: six in /usr/local/lib/python3.7/dist-packages (from onnx->insightface==0.2.1) (1.15.0)\n","Requirement already satisfied: python-dateutil>=2.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib->insightface==0.2.1) (2.8.1)\n","Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.7/dist-packages (from matplotlib->insightface==0.2.1) (0.10.0)\n","Requirement already satisfied: pyparsing!=2.0.4,!=2.1.2,!=2.1.6,>=2.0.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib->insightface==0.2.1) (2.4.7)\n","Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib->insightface==0.2.1) (1.3.1)\n","Requirement already satisfied: networkx>=2.0 in /usr/local/lib/python3.7/dist-packages (from scikit-image->insightface==0.2.1) (2.5.1)\n","Requirement already satisfied: PyWavelets>=0.4.0 in /usr/local/lib/python3.7/dist-packages (from scikit-image->insightface==0.2.1) (1.1.1)\n","Requirement already satisfied: setuptools in /usr/local/lib/python3.7/dist-packages (from protobuf->onnxruntime) (57.0.0)\n","Installing collected packages: onnx, insightface, onnxruntime\n","Successfully installed insightface-0.2.1 onnx-1.9.0 onnxruntime-1.8.0\n","Requirement already satisfied: googledrivedownloader in /usr/local/lib/python3.7/dist-packages (0.4)\n","Requirement already satisfied: imageio==2.4.1 in /usr/local/lib/python3.7/dist-packages (2.4.1)\n","Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from imageio==2.4.1) (1.19.5)\n","Requirement already satisfied: pillow in /usr/local/lib/python3.7/dist-packages (from imageio==2.4.1) (7.1.2)\n"],"name":"stdout"}]},{"cell_type":"code","metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"gQ7ZoIbLFCye","executionInfo":{"status":"ok","timestamp":1625380798405,"user_tz":-480,"elapsed":533,"user":{"displayName":"José Lampreia","photoUrl":"","userId":"16015278604201270582"}},"outputId":"8448a0a3-a19e-44c2-a044-f4d3f9152e91"},"source":["import os\n","os.chdir(\"SimSwap\")\n","!ls"],"execution_count":4,"outputs":[{"output_type":"stream","text":[" crop_224\t simswaplogo\n"," data\t\t test_one_image.py\n"," demo_file\t test_video_swapmulti.py\n"," docs\t\t test_video_swap_multispecific.py\n"," insightface_func test_video_swapsingle.py\n"," LICENSE\t test_video_swapspecific.py\n"," models\t\t test_wholeimage_swapmulti.py\n"," options\t test_wholeimage_swap_multispecific.py\n"," output\t\t test_wholeimage_swapsingle.py\n"," README.md\t test_wholeimage_swapspecific.py\n","'SimSwap colab.ipynb' util\n"],"name":"stdout"}]},{"cell_type":"code","metadata":{"id":"ZvGp-p0nOmKE","executionInfo":{"status":"ok","timestamp":1625380798407,"user_tz":-480,"elapsed":17,"user":{"displayName":"José Lampreia","photoUrl":"","userId":"16015278604201270582"}}},"source":["## You can upload filed manually\n","# from google.colab import drive\n","# drive.mount('/content/gdrive')"],"execution_count":5,"outputs":[]},{"cell_type":"code","metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"gLti1J0pEFjJ","executionInfo":{"status":"ok","timestamp":1625380813268,"user_tz":-480,"elapsed":14876,"user":{"displayName":"José Lampreia","photoUrl":"","userId":"16015278604201270582"}},"outputId":"99dc9306-9b9a-475d-cc7d-a3f423bd1e81"},"source":["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","\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"],"execution_count":6,"outputs":[{"output_type":"stream","text":["--2021-07-04 06:39:56-- https://github.com/neuralchen/SimSwap/releases/download/1.0/arcface_checkpoint.tar\n","Resolving github.com (github.com)... 140.82.114.3\n","Connecting to github.com (github.com)|140.82.114.3|:443... connected.\n","HTTP request sent, awaiting response... 302 Found\n","Location: https://github-releases.githubusercontent.com/374891081/e17b9d00-dcb8-11eb-8c4f-1412bcea78a6?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIWNJYAX4CSVEH53A%2F20210704%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20210704T063956Z&X-Amz-Expires=300&X-Amz-Signature=b6d431c65405e894ddc994061c5fe8fe87db4e71e702513aec01f398a1004825&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=374891081&response-content-disposition=attachment%3B%20filename%3Darcface_checkpoint.tar&response-content-type=application%2Foctet-stream [following]\n","--2021-07-04 06:39:56-- https://github-releases.githubusercontent.com/374891081/e17b9d00-dcb8-11eb-8c4f-1412bcea78a6?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIWNJYAX4CSVEH53A%2F20210704%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20210704T063956Z&X-Amz-Expires=300&X-Amz-Signature=b6d431c65405e894ddc994061c5fe8fe87db4e71e702513aec01f398a1004825&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=374891081&response-content-disposition=attachment%3B%20filename%3Darcface_checkpoint.tar&response-content-type=application%2Foctet-stream\n","Resolving github-releases.githubusercontent.com (github-releases.githubusercontent.com)... 185.199.108.154, 185.199.109.154, 185.199.110.154, ...\n","Connecting to github-releases.githubusercontent.com (github-releases.githubusercontent.com)|185.199.108.154|:443... connected.\n","HTTP request sent, awaiting response... 200 OK\n","Length: 766871429 (731M) [application/octet-stream]\n","Saving to: ‘./arcface_model/arcface_checkpoint.tar’\n","\n","arcface_checkpoint. 100%[===================>] 731.34M 64.4MB/s in 11s \n","\n","2021-07-04 06:40:07 (68.4 MB/s) - ‘./arcface_model/arcface_checkpoint.tar’ saved [766871429/766871429]\n","\n","--2021-07-04 06:40:07-- https://github.com/neuralchen/SimSwap/releases/download/1.0/checkpoints.zip\n","Resolving github.com (github.com)... 140.82.113.3\n","Connecting to github.com (github.com)|140.82.113.3|:443... connected.\n","HTTP request sent, awaiting response... 302 Found\n","Location: https://github-releases.githubusercontent.com/374891081/a8dac400-dcb6-11eb-933f-977cd7f5f554?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIWNJYAX4CSVEH53A%2F20210704%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20210704T063831Z&X-Amz-Expires=300&X-Amz-Signature=3fd2850d03abb9301bf5ba5969d82eb73cb0b940b85e45de2e1e34f1ba2eaf09&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=374891081&response-content-disposition=attachment%3B%20filename%3Dcheckpoints.zip&response-content-type=application%2Foctet-stream [following]\n","--2021-07-04 06:40:07-- https://github-releases.githubusercontent.com/374891081/a8dac400-dcb6-11eb-933f-977cd7f5f554?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIWNJYAX4CSVEH53A%2F20210704%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20210704T063831Z&X-Amz-Expires=300&X-Amz-Signature=3fd2850d03abb9301bf5ba5969d82eb73cb0b940b85e45de2e1e34f1ba2eaf09&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=374891081&response-content-disposition=attachment%3B%20filename%3Dcheckpoints.zip&response-content-type=application%2Foctet-stream\n","Resolving github-releases.githubusercontent.com (github-releases.githubusercontent.com)... 185.199.108.154, 185.199.109.154, 185.199.110.154, ...\n","Connecting to github-releases.githubusercontent.com (github-releases.githubusercontent.com)|185.199.108.154|:443... connected.\n","HTTP request sent, awaiting response... 200 OK\n","Length: 256461775 (245M) [application/octet-stream]\n","Saving to: ‘checkpoints.zip’\n","\n","checkpoints.zip 100%[===================>] 244.58M 219MB/s in 1.1s \n","\n","2021-07-04 06:40:08 (219 MB/s) - ‘checkpoints.zip’ saved [256461775/256461775]\n","\n","Archive: ./checkpoints.zip\n"," creating: ./checkpoints/people/\n"," inflating: ./checkpoints/people/iter.txt \n"," inflating: ./checkpoints/people/latest_net_D1.pth \n"," inflating: ./checkpoints/people/latest_net_D2.pth \n"," inflating: ./checkpoints/people/latest_net_G.pth \n"," inflating: ./checkpoints/people/loss_log.txt \n"," inflating: ./checkpoints/people/opt.txt \n"," creating: ./checkpoints/people/web/\n"," creating: ./checkpoints/people/web/images/\n"],"name":"stdout"}]},{"cell_type":"code","metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"fJ9DYRrCPIUL","executionInfo":{"status":"ok","timestamp":1625380821122,"user_tz":-480,"elapsed":7869,"user":{"displayName":"José Lampreia","photoUrl":"","userId":"16015278604201270582"}},"outputId":"a3d1d841-440c-4244-8045-cb0ce3cc81fd"},"source":["!wget --no-check-certificate \"https://sh23tw.dm.files.1drv.com/y4mmGiIkNVigkSwOKDcV3nwMJulRGhbtHdkheehR5TArc52UjudUYNXAEvKCii2O5LAmzGCGK6IfleocxuDeoKxDZkNzDRSt4ZUlEt8GlSOpCXAFEkBwaZimtWGDRbpIGpb_pz9Nq5jATBQpezBS6G_UtspWTkgrXHHxhviV2nWy8APPx134zOZrUIbkSF6xnsqzs3uZ_SEX_m9Rey0ykpx9w\" -O antelope.zip\n","!unzip ./antelope.zip -d ./insightface_func/models/"],"execution_count":7,"outputs":[{"output_type":"stream","text":["--2021-07-04 06:40:11-- https://sh23tw.dm.files.1drv.com/y4mmGiIkNVigkSwOKDcV3nwMJulRGhbtHdkheehR5TArc52UjudUYNXAEvKCii2O5LAmzGCGK6IfleocxuDeoKxDZkNzDRSt4ZUlEt8GlSOpCXAFEkBwaZimtWGDRbpIGpb_pz9Nq5jATBQpezBS6G_UtspWTkgrXHHxhviV2nWy8APPx134zOZrUIbkSF6xnsqzs3uZ_SEX_m9Rey0ykpx9w\n","Resolving sh23tw.dm.files.1drv.com (sh23tw.dm.files.1drv.com)... 13.107.42.12\n","Connecting to sh23tw.dm.files.1drv.com (sh23tw.dm.files.1drv.com)|13.107.42.12|:443... connected.\n","HTTP request sent, awaiting response... 200 OK\n","Length: 248024513 (237M) [application/zip]\n","Saving to: ‘antelope.zip’\n","\n","antelope.zip 100%[===================>] 236.53M 52.4MB/s in 4.7s \n","\n","2021-07-04 06:40:16 (49.9 MB/s) - ‘antelope.zip’ saved [248024513/248024513]\n","\n","Archive: ./antelope.zip\n"," creating: ./insightface_func/models/antelope/\n"," inflating: ./insightface_func/models/antelope/glintr100.onnx \n"," inflating: ./insightface_func/models/antelope/scrfd_10g_bnkps.onnx \n"],"name":"stdout"}]},{"cell_type":"code","metadata":{"id":"PfSsND36EMvn","colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"status":"ok","timestamp":1625380827902,"user_tz":-480,"elapsed":6811,"user":{"displayName":"José Lampreia","photoUrl":"","userId":"16015278604201270582"}},"outputId":"8130e97d-b4a0-4988-85fb-e2bc3e755259"},"source":["import cv2\n","import torch\n","import fractions\n","import numpy as np\n","from PIL import Image\n","import torch.nn.functional as F\n","from torchvision import transforms\n","from models.models import create_model\n","from options.test_options import TestOptions\n","from insightface_func.face_detect_crop_multi import Face_detect_crop\n","from util.videoswap_multispecific import video_swap\n","import os\n","import glob"],"execution_count":8,"outputs":[{"output_type":"stream","text":["Imageio: 'ffmpeg-linux64-v3.3.1' was not found on your computer; downloading it now.\n","Try 1. Download from https://github.com/imageio/imageio-binaries/raw/master/ffmpeg/ffmpeg-linux64-v3.3.1 (43.8 MB)\n","Downloading: 8192/45929032 bytes (0.0%)\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b3432448/45929032 bytes (7.5%)\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b7036928/45929032 bytes (15.3%)\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b10641408/45929032 bytes (23.2%)\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b14278656/45929032 bytes (31.1%)\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b18104320/45929032 bytes (39.4%)\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b21954560/45929032 bytes (47.8%)\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b25780224/45929032 bytes (56.1%)\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b29736960/45929032 bytes (64.7%)\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b33488896/45929032 bytes (72.9%)\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b37093376/45929032 bytes (80.8%)\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b40689664/45929032 bytes (88.6%)\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b44392448/45929032 bytes (96.7%)\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b45929032/45929032 bytes (100.0%)\n"," Done\n","File saved as /root/.imageio/ffmpeg/ffmpeg-linux64-v3.3.1.\n"],"name":"stdout"}]},{"cell_type":"code","metadata":{"id":"rxSbZ2EDNDlf","executionInfo":{"status":"ok","timestamp":1625380827903,"user_tz":-480,"elapsed":12,"user":{"displayName":"José Lampreia","photoUrl":"","userId":"16015278604201270582"}}},"source":["def lcm(a, b): return abs(a * b) / fractions.gcd(a, b) if a and b else 0\n","\n","transformer = transforms.Compose([\n"," transforms.ToTensor(),\n"," #transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n"," ])\n","\n","transformer_Arcface = transforms.Compose([\n"," transforms.ToTensor(),\n"," transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n"," ])\n"],"execution_count":9,"outputs":[]},{"cell_type":"code","metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"ye8iS0UVPMRg","executionInfo":{"status":"ok","timestamp":1625380828574,"user_tz":-480,"elapsed":680,"user":{"displayName":"José Lampreia","photoUrl":"","userId":"16015278604201270582"}},"outputId":"cb5a4b02-b1d0-4ff8-f542-5c1b3c9703d9"},"source":["!ls ./checkpoints"],"execution_count":10,"outputs":[{"output_type":"stream","text":["people\n"],"name":"stdout"}]},{"cell_type":"code","metadata":{"id":"wwJOwR9LNKRz","colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"status":"ok","timestamp":1625380828576,"user_tz":-480,"elapsed":13,"user":{"displayName":"José Lampreia","photoUrl":"","userId":"16015278604201270582"}},"outputId":"0f92f785-4d9c-4130-b24d-76871b2dafba"},"source":["opt = TestOptions()\n","opt.initialize()\n","opt.parser.add_argument('-f') ## dummy arg to avoid bug\n","opt = opt.parse()\n","opt.multisepcific_dir = './demo_file/multispecific' ## or replace it with folder from your own google drive\n"," ## and remember to follow the dir structure in usage.md\n","opt.video_path = './demo_file/multi_people_1080p.mp4' ## or replace it with video from your own google drive\n","opt.output_path = './output/multi_test_multispecific.mp4'\n","opt.temp_path = './tmp'\n","opt.Arc_path = './arcface_model/arcface_checkpoint.tar'\n","opt.name = 'people'\n","opt.isTrain = False\n","\n","crop_size = 224\n"],"execution_count":11,"outputs":[{"output_type":"stream","text":["------------ Options -------------\n","Arc_path: models/BEST_checkpoint.tar\n","aspect_ratio: 1.0\n","batchSize: 8\n","checkpoints_dir: ./checkpoints\n","cluster_path: features_clustered_010.npy\n","data_type: 32\n","dataroot: ./datasets/cityscapes/\n","display_winsize: 512\n","engine: None\n","export_onnx: None\n","f: /root/.local/share/jupyter/runtime/kernel-19937219-895d-4d02-9a72-5cfa0e889adf.json\n","feat_num: 3\n","fineSize: 512\n","fp16: False\n","gpu_ids: [0]\n","how_many: 50\n","id_thres: 0.03\n","image_size: 224\n","input_nc: 3\n","instance_feat: False\n","isTrain: False\n","label_feat: False\n","label_nc: 0\n","latent_size: 512\n","loadSize: 1024\n","load_features: False\n","local_rank: 0\n","max_dataset_size: inf\n","model: pix2pixHD\n","multisepcific_dir: ./demo_file/multispecific\n","nThreads: 2\n","n_blocks_global: 6\n","n_blocks_local: 3\n","n_clusters: 10\n","n_downsample_E: 4\n","n_downsample_global: 3\n","n_local_enhancers: 1\n","name: people\n","nef: 16\n","netG: global\n","ngf: 64\n","niter_fix_global: 0\n","no_flip: False\n","no_instance: False\n","no_simswaplogo: False\n","norm: batch\n","norm_G: spectralspadesyncbatch3x3\n","ntest: inf\n","onnx: None\n","output_nc: 3\n","output_path: ./output/\n","phase: test\n","pic_a_path: ./crop_224/gdg.jpg\n","pic_b_path: ./crop_224/zrf.jpg\n","pic_specific_path: ./crop_224/zrf.jpg\n","resize_or_crop: scale_width\n","results_dir: ./results/\n","semantic_nc: 3\n","serial_batches: False\n","temp_path: ./temp_results\n","tf_log: False\n","use_dropout: False\n","use_encoded_image: False\n","verbose: False\n","video_path: ./demo_file/multi_people_1080p.mp4\n","which_epoch: latest\n","-------------- End ----------------\n"],"name":"stdout"}]},{"cell_type":"code","metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"UFt8zQrAMq9F","executionInfo":{"status":"ok","timestamp":1625381428564,"user_tz":-480,"elapsed":599996,"user":{"displayName":"José Lampreia","photoUrl":"","userId":"16015278604201270582"}},"outputId":"46188013-3f6d-4174-9959-d1fd203dcc0d"},"source":["pic_specific = opt.pic_specific_path\n","crop_size = 224\n","multisepcific_dir = opt.multisepcific_dir\n","\n","torch.nn.Module.dump_patches = True\n","model = create_model(opt)\n","model.eval()\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","# The specific person to be swapped(source)\n","source_specific_id_nonorm_list = []\n","source_path = os.path.join(multisepcific_dir,'SRC_*')\n","source_specific_images_path = sorted(glob.glob(source_path))\n","\n","for source_specific_image_path in source_specific_images_path:\n"," specific_person_whole = cv2.imread(source_specific_image_path)\n"," specific_person_align_crop, _ = app.get(specific_person_whole,crop_size)\n"," specific_person_align_crop_pil = Image.fromarray(cv2.cvtColor(specific_person_align_crop[0],cv2.COLOR_BGR2RGB)) \n"," specific_person = transformer_Arcface(specific_person_align_crop_pil)\n"," specific_person = specific_person.view(-1, specific_person.shape[0], specific_person.shape[1], specific_person.shape[2])\n"," # convert numpy to tensor\n"," specific_person = specific_person.cuda()\n"," #create latent id\n"," specific_person_downsample = F.interpolate(specific_person, scale_factor=0.5)\n"," specific_person_id_nonorm = model.netArc(specific_person_downsample)\n"," source_specific_id_nonorm_list.append(specific_person_id_nonorm.clone())\n","\n","# The person who provides id information (list)\n","target_id_norm_list = []\n","target_path = os.path.join(multisepcific_dir,'DST_*')\n","target_images_path = sorted(glob.glob(target_path))\n","\n","for target_image_path in target_images_path:\n"," img_a_whole = cv2.imread(target_image_path)\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"," # convert numpy to tensor\n"," img_id = img_id.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 = F.normalize(latend_id, p=2, dim=1)\n"," target_id_norm_list.append(latend_id.clone())\n"," \n","assert len(target_id_norm_list) == len(source_specific_id_nonorm_list), \"The number of images in source and target directory must be same !!!\"\n","video_swap(opt.video_path, target_id_norm_list,source_specific_id_nonorm_list, opt.id_thres, \\\n"," model, app, opt.output_path,temp_results_dir=opt.temp_path,no_simswaplogo=opt.no_simswaplogo)"],"execution_count":12,"outputs":[{"output_type":"stream","text":["input mean and std: 127.5 127.5\n","find model: ./insightface_func/models/antelope/glintr100.onnx recognition\n","find model: ./insightface_func/models/antelope/scrfd_10g_bnkps.onnx detection\n","set det-size: (640, 640)\n"],"name":"stdout"},{"output_type":"stream","text":["\r 0%| | 0/594 [00:00>>> Building video ./output/multi_test_multispecific.mp4\n","[MoviePy] Writing audio in multi_test_multispecificTEMP_MPY_wvf_snd.mp3\n"],"name":"stdout"},{"output_type":"stream","text":["100%|██████████| 438/438 [00:00<00:00, 832.16it/s]"],"name":"stderr"},{"output_type":"stream","text":["[MoviePy] Done.\n","[MoviePy] Writing video ./output/multi_test_multispecific.mp4\n"],"name":"stdout"},{"output_type":"stream","text":["\n","100%|██████████| 595/595 [00:53<00:00, 11.11it/s]\n"],"name":"stderr"},{"output_type":"stream","text":["[MoviePy] Done.\n","[MoviePy] >>>> Video ready: ./output/multi_test_multispecific.mp4 \n","\n"],"name":"stdout"}]},{"cell_type":"code","metadata":{"id":"Rty2GsyZZrI6"},"source":[],"execution_count":null,"outputs":[]}]} \ No newline at end of file +{"nbformat":4,"nbformat_minor":0,"metadata":{"colab":{"name":"MultiSpecific.ipynb","provenance":[],"collapsed_sections":[],"authorship_tag":"ABX9TyNw8SfPWhG77cf/e7YZd178"},"kernelspec":{"name":"python3","display_name":"Python 3"},"language_info":{"name":"python"},"accelerator":"GPU"},"cells":[{"cell_type":"markdown","metadata":{"id":"7_gtFoV8BuRx"},"source":["This is an example of SimSwap on processing video with multiple faces with designated sources.\n","\n","Code path: https://github.com/neuralchen/SimSwap\n","Paper path: https://arxiv.org/pdf/2106.06340v1.pdf."]},{"cell_type":"code","metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"0Y1RfpzsCAl9","executionInfo":{"status":"ok","timestamp":1625380781426,"user_tz":-480,"elapsed":586,"user":{"displayName":"José Lampreia","photoUrl":"","userId":"16015278604201270582"}},"outputId":"2a897b34-72f1-4515-ac6f-2f0e2d4ea4f7"},"source":["## make sure you are using a runtime with GPU\n","## you can check at Runtime/Change runtime type in the top bar.\n","!nvidia-smi"],"execution_count":1,"outputs":[{"output_type":"stream","text":["Sun Jul 4 06:39:39 2021 \n","+-----------------------------------------------------------------------------+\n","| NVIDIA-SMI 465.27 Driver Version: 460.32.03 CUDA Version: 11.2 |\n","|-------------------------------+----------------------+----------------------+\n","| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |\n","| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |\n","| | | MIG M. |\n","|===============================+======================+======================|\n","| 0 Tesla T4 Off | 00000000:00:04.0 Off | 0 |\n","| N/A 45C P8 9W / 70W | 0MiB / 15109MiB | 0% Default |\n","| | | N/A |\n","+-------------------------------+----------------------+----------------------+\n"," \n","+-----------------------------------------------------------------------------+\n","| Processes: |\n","| GPU GI CI PID Type Process name GPU Memory |\n","| ID ID Usage |\n","|=============================================================================|\n","| No running processes found |\n","+-----------------------------------------------------------------------------+\n"],"name":"stdout"}]},{"cell_type":"markdown","metadata":{"id":"0Qzzx2UpDkqw"},"source":["All file changes make by this notebook are temporary. \n","You can try to mount your own google drive to store files if you wang.\n"]},{"cell_type":"code","metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"VA_4CeWZCHLP","executionInfo":{"status":"ok","timestamp":1625380786661,"user_tz":-480,"elapsed":4693,"user":{"displayName":"José Lampreia","photoUrl":"","userId":"16015278604201270582"}},"outputId":"d0665552-be92-45b0-aab2-f84c619a51fb"},"source":["!git clone https://github.com/neuralchen/SimSwap\n","!cd SimSwap && git pull"],"execution_count":2,"outputs":[{"output_type":"stream","text":["Cloning into 'SimSwap'...\n","remote: Enumerating objects: 667, done.\u001b[K\n","remote: Counting objects: 100% (48/48), done.\u001b[K\n","remote: Compressing objects: 100% (35/35), done.\u001b[K\n","remote: Total 667 (delta 19), reused 28 (delta 13), pack-reused 619\u001b[K\n","Receiving objects: 100% (667/667), 132.14 MiB | 44.44 MiB/s, done.\n","Resolving deltas: 100% (292/292), done.\n","Already up to date.\n"],"name":"stdout"}]},{"cell_type":"code","metadata":{"id":"Y5K4au_UCkKn","colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"status":"ok","timestamp":1625380797906,"user_tz":-480,"elapsed":11253,"user":{"displayName":"José Lampreia","photoUrl":"","userId":"16015278604201270582"}},"outputId":"7429f153-bc6d-48c2-eb3c-21f1f02fede9"},"source":["!pip install insightface==0.2.1 onnxruntime moviepy\n","!pip install googledrivedownloader\n","!pip install imageio==2.4.1"],"execution_count":3,"outputs":[{"output_type":"stream","text":["Collecting insightface==0.2.1\n"," Downloading https://files.pythonhosted.org/packages/ee/1e/6395bbe0db665f187c8e49266cda54fcf661f182192370d409423e4943e4/insightface-0.2.1-py2.py3-none-any.whl\n","Collecting onnxruntime\n","\u001b[?25l Downloading https://files.pythonhosted.org/packages/f9/76/3d0f8bb2776961c7335693df06eccf8d099e48fa6fb552c7546867192603/onnxruntime-1.8.0-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (4.5MB)\n","\u001b[K |████████████████████████████████| 4.5MB 37.5MB/s \n","\u001b[?25hRequirement already satisfied: moviepy in /usr/local/lib/python3.7/dist-packages (0.2.3.5)\n","Requirement already satisfied: scikit-learn in /usr/local/lib/python3.7/dist-packages (from insightface==0.2.1) (0.22.2.post1)\n","Requirement already satisfied: requests in /usr/local/lib/python3.7/dist-packages (from insightface==0.2.1) (2.23.0)\n","Collecting onnx\n","\u001b[?25l Downloading https://files.pythonhosted.org/packages/3f/9b/54c950d3256e27f970a83cd0504efb183a24312702deed0179453316dbd0/onnx-1.9.0-cp37-cp37m-manylinux2010_x86_64.whl (12.2MB)\n","\u001b[K |████████████████████████████████| 12.2MB 32.2MB/s \n","\u001b[?25hRequirement already satisfied: matplotlib in /usr/local/lib/python3.7/dist-packages (from insightface==0.2.1) (3.2.2)\n","Requirement already satisfied: Pillow in /usr/local/lib/python3.7/dist-packages (from insightface==0.2.1) (7.1.2)\n","Requirement already satisfied: scikit-image in /usr/local/lib/python3.7/dist-packages (from insightface==0.2.1) (0.16.2)\n","Requirement already satisfied: opencv-python in /usr/local/lib/python3.7/dist-packages (from insightface==0.2.1) (4.1.2.30)\n","Requirement already satisfied: tqdm in /usr/local/lib/python3.7/dist-packages (from insightface==0.2.1) (4.41.1)\n","Requirement already satisfied: scipy in /usr/local/lib/python3.7/dist-packages (from insightface==0.2.1) (1.4.1)\n","Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from insightface==0.2.1) (1.19.5)\n","Requirement already satisfied: easydict in /usr/local/lib/python3.7/dist-packages (from insightface==0.2.1) (1.9)\n","Requirement already satisfied: flatbuffers in /usr/local/lib/python3.7/dist-packages (from onnxruntime) (1.12)\n","Requirement already satisfied: protobuf in /usr/local/lib/python3.7/dist-packages (from onnxruntime) (3.12.4)\n","Requirement already satisfied: imageio<3.0,>=2.1.2 in /usr/local/lib/python3.7/dist-packages (from moviepy) (2.4.1)\n","Requirement already satisfied: decorator<5.0,>=4.0.2 in /usr/local/lib/python3.7/dist-packages (from moviepy) (4.4.2)\n","Requirement already satisfied: joblib>=0.11 in /usr/local/lib/python3.7/dist-packages (from scikit-learn->insightface==0.2.1) (1.0.1)\n","Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests->insightface==0.2.1) (3.0.4)\n","Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests->insightface==0.2.1) (1.24.3)\n","Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests->insightface==0.2.1) (2021.5.30)\n","Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests->insightface==0.2.1) (2.10)\n","Requirement already satisfied: typing-extensions>=3.6.2.1 in /usr/local/lib/python3.7/dist-packages (from onnx->insightface==0.2.1) (3.7.4.3)\n","Requirement already satisfied: six in /usr/local/lib/python3.7/dist-packages (from onnx->insightface==0.2.1) (1.15.0)\n","Requirement already satisfied: python-dateutil>=2.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib->insightface==0.2.1) (2.8.1)\n","Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.7/dist-packages (from matplotlib->insightface==0.2.1) (0.10.0)\n","Requirement already satisfied: pyparsing!=2.0.4,!=2.1.2,!=2.1.6,>=2.0.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib->insightface==0.2.1) (2.4.7)\n","Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib->insightface==0.2.1) (1.3.1)\n","Requirement already satisfied: networkx>=2.0 in /usr/local/lib/python3.7/dist-packages (from scikit-image->insightface==0.2.1) (2.5.1)\n","Requirement already satisfied: PyWavelets>=0.4.0 in /usr/local/lib/python3.7/dist-packages (from scikit-image->insightface==0.2.1) (1.1.1)\n","Requirement already satisfied: setuptools in /usr/local/lib/python3.7/dist-packages (from protobuf->onnxruntime) (57.0.0)\n","Installing collected packages: onnx, insightface, onnxruntime\n","Successfully installed insightface-0.2.1 onnx-1.9.0 onnxruntime-1.8.0\n","Requirement already satisfied: googledrivedownloader in /usr/local/lib/python3.7/dist-packages (0.4)\n","Requirement already satisfied: imageio==2.4.1 in /usr/local/lib/python3.7/dist-packages (2.4.1)\n","Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from imageio==2.4.1) (1.19.5)\n","Requirement already satisfied: pillow in /usr/local/lib/python3.7/dist-packages (from imageio==2.4.1) (7.1.2)\n"],"name":"stdout"}]},{"cell_type":"code","metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"gQ7ZoIbLFCye","executionInfo":{"status":"ok","timestamp":1625380798405,"user_tz":-480,"elapsed":533,"user":{"displayName":"José Lampreia","photoUrl":"","userId":"16015278604201270582"}},"outputId":"8448a0a3-a19e-44c2-a044-f4d3f9152e91"},"source":["import os\n","os.chdir(\"SimSwap\")\n","!ls"],"execution_count":4,"outputs":[{"output_type":"stream","text":[" crop_224\t simswaplogo\n"," data\t\t test_one_image.py\n"," demo_file\t test_video_swapmulti.py\n"," docs\t\t test_video_swap_multispecific.py\n"," insightface_func test_video_swapsingle.py\n"," LICENSE\t test_video_swapspecific.py\n"," models\t\t test_wholeimage_swapmulti.py\n"," options\t test_wholeimage_swap_multispecific.py\n"," output\t\t test_wholeimage_swapsingle.py\n"," README.md\t test_wholeimage_swapspecific.py\n","'SimSwap colab.ipynb' util\n"],"name":"stdout"}]},{"cell_type":"code","metadata":{"id":"ZvGp-p0nOmKE","executionInfo":{"status":"ok","timestamp":1625380798407,"user_tz":-480,"elapsed":17,"user":{"displayName":"José Lampreia","photoUrl":"","userId":"16015278604201270582"}}},"source":["## You can upload filed manually\n","# from google.colab import drive\n","# drive.mount('/content/gdrive')"],"execution_count":5,"outputs":[]},{"cell_type":"code","metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"gLti1J0pEFjJ","executionInfo":{"status":"ok","timestamp":1625380813268,"user_tz":-480,"elapsed":14876,"user":{"displayName":"José Lampreia","photoUrl":"","userId":"16015278604201270582"}},"outputId":"99dc9306-9b9a-475d-cc7d-a3f423bd1e81"},"source":["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","\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":6,"outputs":[{"output_type":"stream","text":["--2021-07-04 06:39:56-- https://github.com/neuralchen/SimSwap/releases/download/1.0/arcface_checkpoint.tar\n","Resolving github.com (github.com)... 140.82.114.3\n","Connecting to github.com (github.com)|140.82.114.3|:443... connected.\n","HTTP request sent, awaiting response... 302 Found\n","Location: https://github-releases.githubusercontent.com/374891081/e17b9d00-dcb8-11eb-8c4f-1412bcea78a6?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIWNJYAX4CSVEH53A%2F20210704%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20210704T063956Z&X-Amz-Expires=300&X-Amz-Signature=b6d431c65405e894ddc994061c5fe8fe87db4e71e702513aec01f398a1004825&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=374891081&response-content-disposition=attachment%3B%20filename%3Darcface_checkpoint.tar&response-content-type=application%2Foctet-stream [following]\n","--2021-07-04 06:39:56-- https://github-releases.githubusercontent.com/374891081/e17b9d00-dcb8-11eb-8c4f-1412bcea78a6?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIWNJYAX4CSVEH53A%2F20210704%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20210704T063956Z&X-Amz-Expires=300&X-Amz-Signature=b6d431c65405e894ddc994061c5fe8fe87db4e71e702513aec01f398a1004825&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=374891081&response-content-disposition=attachment%3B%20filename%3Darcface_checkpoint.tar&response-content-type=application%2Foctet-stream\n","Resolving github-releases.githubusercontent.com (github-releases.githubusercontent.com)... 185.199.108.154, 185.199.109.154, 185.199.110.154, ...\n","Connecting to github-releases.githubusercontent.com (github-releases.githubusercontent.com)|185.199.108.154|:443... connected.\n","HTTP request sent, awaiting response... 200 OK\n","Length: 766871429 (731M) [application/octet-stream]\n","Saving to: ‘./arcface_model/arcface_checkpoint.tar’\n","\n","arcface_checkpoint. 100%[===================>] 731.34M 64.4MB/s in 11s \n","\n","2021-07-04 06:40:07 (68.4 MB/s) - ‘./arcface_model/arcface_checkpoint.tar’ saved [766871429/766871429]\n","\n","--2021-07-04 06:40:07-- https://github.com/neuralchen/SimSwap/releases/download/1.0/checkpoints.zip\n","Resolving github.com (github.com)... 140.82.113.3\n","Connecting to github.com (github.com)|140.82.113.3|:443... connected.\n","HTTP request sent, awaiting response... 302 Found\n","Location: https://github-releases.githubusercontent.com/374891081/a8dac400-dcb6-11eb-933f-977cd7f5f554?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIWNJYAX4CSVEH53A%2F20210704%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20210704T063831Z&X-Amz-Expires=300&X-Amz-Signature=3fd2850d03abb9301bf5ba5969d82eb73cb0b940b85e45de2e1e34f1ba2eaf09&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=374891081&response-content-disposition=attachment%3B%20filename%3Dcheckpoints.zip&response-content-type=application%2Foctet-stream [following]\n","--2021-07-04 06:40:07-- https://github-releases.githubusercontent.com/374891081/a8dac400-dcb6-11eb-933f-977cd7f5f554?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIWNJYAX4CSVEH53A%2F20210704%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20210704T063831Z&X-Amz-Expires=300&X-Amz-Signature=3fd2850d03abb9301bf5ba5969d82eb73cb0b940b85e45de2e1e34f1ba2eaf09&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=374891081&response-content-disposition=attachment%3B%20filename%3Dcheckpoints.zip&response-content-type=application%2Foctet-stream\n","Resolving github-releases.githubusercontent.com (github-releases.githubusercontent.com)... 185.199.108.154, 185.199.109.154, 185.199.110.154, ...\n","Connecting to github-releases.githubusercontent.com (github-releases.githubusercontent.com)|185.199.108.154|:443... connected.\n","HTTP request sent, awaiting response... 200 OK\n","Length: 256461775 (245M) [application/octet-stream]\n","Saving to: ‘checkpoints.zip’\n","\n","checkpoints.zip 100%[===================>] 244.58M 219MB/s in 1.1s \n","\n","2021-07-04 06:40:08 (219 MB/s) - ‘checkpoints.zip’ saved [256461775/256461775]\n","\n","Archive: ./checkpoints.zip\n"," creating: ./checkpoints/people/\n"," inflating: ./checkpoints/people/iter.txt \n"," inflating: ./checkpoints/people/latest_net_D1.pth \n"," inflating: ./checkpoints/people/latest_net_D2.pth \n"," inflating: ./checkpoints/people/latest_net_G.pth \n"," inflating: ./checkpoints/people/loss_log.txt \n"," inflating: ./checkpoints/people/opt.txt \n"," creating: ./checkpoints/people/web/\n"," creating: ./checkpoints/people/web/images/\n"],"name":"stdout"}]},{"cell_type":"code","metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"fJ9DYRrCPIUL","executionInfo":{"status":"ok","timestamp":1625380821122,"user_tz":-480,"elapsed":7869,"user":{"displayName":"José Lampreia","photoUrl":"","userId":"16015278604201270582"}},"outputId":"a3d1d841-440c-4244-8045-cb0ce3cc81fd"},"source":["!wget --no-check-certificate \"https://sh23tw.dm.files.1drv.com/y4mmGiIkNVigkSwOKDcV3nwMJulRGhbtHdkheehR5TArc52UjudUYNXAEvKCii2O5LAmzGCGK6IfleocxuDeoKxDZkNzDRSt4ZUlEt8GlSOpCXAFEkBwaZimtWGDRbpIGpb_pz9Nq5jATBQpezBS6G_UtspWTkgrXHHxhviV2nWy8APPx134zOZrUIbkSF6xnsqzs3uZ_SEX_m9Rey0ykpx9w\" -O antelope.zip\n","!unzip ./antelope.zip -d ./insightface_func/models/"],"execution_count":7,"outputs":[{"output_type":"stream","text":["--2021-07-04 06:40:11-- https://sh23tw.dm.files.1drv.com/y4mmGiIkNVigkSwOKDcV3nwMJulRGhbtHdkheehR5TArc52UjudUYNXAEvKCii2O5LAmzGCGK6IfleocxuDeoKxDZkNzDRSt4ZUlEt8GlSOpCXAFEkBwaZimtWGDRbpIGpb_pz9Nq5jATBQpezBS6G_UtspWTkgrXHHxhviV2nWy8APPx134zOZrUIbkSF6xnsqzs3uZ_SEX_m9Rey0ykpx9w\n","Resolving sh23tw.dm.files.1drv.com (sh23tw.dm.files.1drv.com)... 13.107.42.12\n","Connecting to sh23tw.dm.files.1drv.com (sh23tw.dm.files.1drv.com)|13.107.42.12|:443... connected.\n","HTTP request sent, awaiting response... 200 OK\n","Length: 248024513 (237M) [application/zip]\n","Saving to: ‘antelope.zip’\n","\n","antelope.zip 100%[===================>] 236.53M 52.4MB/s in 4.7s \n","\n","2021-07-04 06:40:16 (49.9 MB/s) - ‘antelope.zip’ saved [248024513/248024513]\n","\n","Archive: ./antelope.zip\n"," creating: ./insightface_func/models/antelope/\n"," inflating: ./insightface_func/models/antelope/glintr100.onnx \n"," inflating: ./insightface_func/models/antelope/scrfd_10g_bnkps.onnx \n"],"name":"stdout"}]},{"cell_type":"code","metadata":{"id":"PfSsND36EMvn","colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"status":"ok","timestamp":1625380827902,"user_tz":-480,"elapsed":6811,"user":{"displayName":"José Lampreia","photoUrl":"","userId":"16015278604201270582"}},"outputId":"8130e97d-b4a0-4988-85fb-e2bc3e755259"},"source":["import cv2\n","import torch\n","import fractions\n","import numpy as np\n","from PIL import Image\n","import torch.nn.functional as F\n","from torchvision import transforms\n","from models.models import create_model\n","from options.test_options import TestOptions\n","from insightface_func.face_detect_crop_multi import Face_detect_crop\n","from util.videoswap_multispecific import video_swap\n","import os\n","import glob"],"execution_count":8,"outputs":[{"output_type":"stream","text":["Imageio: 'ffmpeg-linux64-v3.3.1' was not found on your computer; downloading it now.\n","Try 1. Download from https://github.com/imageio/imageio-binaries/raw/master/ffmpeg/ffmpeg-linux64-v3.3.1 (43.8 MB)\n","Downloading: 8192/45929032 bytes (0.0%)\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b3432448/45929032 bytes (7.5%)\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b7036928/45929032 bytes (15.3%)\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b10641408/45929032 bytes (23.2%)\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b14278656/45929032 bytes (31.1%)\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b18104320/45929032 bytes (39.4%)\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b21954560/45929032 bytes (47.8%)\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b25780224/45929032 bytes (56.1%)\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b29736960/45929032 bytes (64.7%)\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b33488896/45929032 bytes (72.9%)\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b37093376/45929032 bytes (80.8%)\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b40689664/45929032 bytes (88.6%)\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b44392448/45929032 bytes (96.7%)\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b45929032/45929032 bytes (100.0%)\n"," Done\n","File saved as /root/.imageio/ffmpeg/ffmpeg-linux64-v3.3.1.\n"],"name":"stdout"}]},{"cell_type":"code","metadata":{"id":"rxSbZ2EDNDlf","executionInfo":{"status":"ok","timestamp":1625380827903,"user_tz":-480,"elapsed":12,"user":{"displayName":"José Lampreia","photoUrl":"","userId":"16015278604201270582"}}},"source":["def lcm(a, b): return abs(a * b) / fractions.gcd(a, b) if a and b else 0\n","\n","transformer = transforms.Compose([\n"," transforms.ToTensor(),\n"," #transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n"," ])\n","\n","transformer_Arcface = transforms.Compose([\n"," transforms.ToTensor(),\n"," transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n"," ])\n"],"execution_count":9,"outputs":[]},{"cell_type":"code","metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"ye8iS0UVPMRg","executionInfo":{"status":"ok","timestamp":1625380828574,"user_tz":-480,"elapsed":680,"user":{"displayName":"José Lampreia","photoUrl":"","userId":"16015278604201270582"}},"outputId":"cb5a4b02-b1d0-4ff8-f542-5c1b3c9703d9"},"source":["!ls ./checkpoints"],"execution_count":10,"outputs":[{"output_type":"stream","text":["people\n"],"name":"stdout"}]},{"cell_type":"code","metadata":{"id":"wwJOwR9LNKRz","colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"status":"ok","timestamp":1625380828576,"user_tz":-480,"elapsed":13,"user":{"displayName":"José Lampreia","photoUrl":"","userId":"16015278604201270582"}},"outputId":"0f92f785-4d9c-4130-b24d-76871b2dafba"},"source":["opt = TestOptions()\n","opt.initialize()\n","opt.parser.add_argument('-f') ## dummy arg to avoid bug\n","opt = opt.parse()\n","opt.multisepcific_dir = './demo_file/multispecific' ## or replace it with folder from your own google drive\n"," ## and remember to follow the dir structure in usage.md\n","opt.video_path = './demo_file/multi_people_1080p.mp4' ## or replace it with video from your own google drive\n","opt.output_path = './output/multi_test_multispecific.mp4'\n","opt.temp_path = './tmp'\n","opt.Arc_path = './arcface_model/arcface_checkpoint.tar'\n","opt.name = 'people'\n","opt.isTrain = False\n","opt.use_mask = True ## new feature up-to-date\n","\n","crop_size = opt.crop_size\n"],"execution_count":11,"outputs":[{"output_type":"stream","text":["------------ Options -------------\n","Arc_path: models/BEST_checkpoint.tar\n","aspect_ratio: 1.0\n","batchSize: 8\n","checkpoints_dir: ./checkpoints\n","cluster_path: features_clustered_010.npy\n","data_type: 32\n","dataroot: ./datasets/cityscapes/\n","display_winsize: 512\n","engine: None\n","export_onnx: None\n","f: /root/.local/share/jupyter/runtime/kernel-19937219-895d-4d02-9a72-5cfa0e889adf.json\n","feat_num: 3\n","fineSize: 512\n","fp16: False\n","gpu_ids: [0]\n","how_many: 50\n","id_thres: 0.03\n","image_size: 224\n","input_nc: 3\n","instance_feat: False\n","isTrain: False\n","label_feat: False\n","label_nc: 0\n","latent_size: 512\n","loadSize: 1024\n","load_features: False\n","local_rank: 0\n","max_dataset_size: inf\n","model: pix2pixHD\n","multisepcific_dir: ./demo_file/multispecific\n","nThreads: 2\n","n_blocks_global: 6\n","n_blocks_local: 3\n","n_clusters: 10\n","n_downsample_E: 4\n","n_downsample_global: 3\n","n_local_enhancers: 1\n","name: people\n","nef: 16\n","netG: global\n","ngf: 64\n","niter_fix_global: 0\n","no_flip: False\n","no_instance: False\n","no_simswaplogo: False\n","norm: batch\n","norm_G: spectralspadesyncbatch3x3\n","ntest: inf\n","onnx: None\n","output_nc: 3\n","output_path: ./output/\n","phase: test\n","pic_a_path: ./crop_224/gdg.jpg\n","pic_b_path: ./crop_224/zrf.jpg\n","pic_specific_path: ./crop_224/zrf.jpg\n","resize_or_crop: scale_width\n","results_dir: ./results/\n","semantic_nc: 3\n","serial_batches: False\n","temp_path: ./temp_results\n","tf_log: False\n","use_dropout: False\n","use_encoded_image: False\n","verbose: False\n","video_path: ./demo_file/multi_people_1080p.mp4\n","which_epoch: latest\n","-------------- End ----------------\n"],"name":"stdout"}]},{"cell_type":"code","metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"UFt8zQrAMq9F","executionInfo":{"status":"ok","timestamp":1625381428564,"user_tz":-480,"elapsed":599996,"user":{"displayName":"José Lampreia","photoUrl":"","userId":"16015278604201270582"}},"outputId":"46188013-3f6d-4174-9959-d1fd203dcc0d"},"source":["pic_specific = opt.pic_specific_path\n","crop_size = opt.crop_size\n","multisepcific_dir = opt.multisepcific_dir\n","\n","torch.nn.Module.dump_patches = True\n","model = create_model(opt)\n","model.eval()\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","# The specific person to be swapped(source)\n","source_specific_id_nonorm_list = []\n","source_path = os.path.join(multisepcific_dir,'SRC_*')\n","source_specific_images_path = sorted(glob.glob(source_path))\n","\n","with torch.no_grad():\n"," for source_specific_image_path in source_specific_images_path:\n"," specific_person_whole = cv2.imread(source_specific_image_path)\n"," specific_person_align_crop, _ = app.get(specific_person_whole,crop_size)\n"," specific_person_align_crop_pil = Image.fromarray(cv2.cvtColor(specific_person_align_crop[0],cv2.COLOR_BGR2RGB)) \n"," specific_person = transformer_Arcface(specific_person_align_crop_pil)\n"," specific_person = specific_person.view(-1, specific_person.shape[0], specific_person.shape[1], specific_person.shape [2])\n"," # convert numpy to tensor\n"," specific_person = specific_person.cuda()\n"," #create latent id\n"," specific_person_downsample = F.interpolate(specific_person, size=(112,112))\n"," specific_person_id_nonorm = model.netArc(specific_person_downsample)\n"," source_specific_id_nonorm_list.append(specific_person_id_nonorm.clone())\n","\n"," # The person who provides id information (list)\n"," target_id_norm_list = []\n"," target_path = os.path.join(multisepcific_dir,'DST_*')\n"," target_images_path = sorted(glob.glob(target_path))\n","\n"," for target_image_path in target_images_path:\n"," img_a_whole = cv2.imread(target_image_path)\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"," # convert numpy to tensor\n"," img_id = img_id.cuda()\n"," #create latent id\n"," img_id_downsample = F.interpolate(img_id, size=(112,112))\n"," latend_id = model.netArc(img_id_downsample)\n"," latend_id = F.normalize(latend_id, p=2, dim=1)\n"," target_id_norm_list.append(latend_id.clone())\n"," \n"," assert len(target_id_norm_list) == len(source_specific_id_nonorm_list), \"The number of images in source and target directory must be same !!!\"\n"," video_swap(opt.video_path, target_id_norm_list,source_specific_id_nonorm_list, opt.id_thres, \\\n"," model, app, opt.output_path,temp_results_dir=opt.temp_path,no_simswaplogo=opt.no_simswaplogo,use_mask=opt.use_mask)"],"execution_count":12,"outputs":[{"output_type":"stream","text":["input mean and std: 127.5 127.5\n","find model: ./insightface_func/models/antelope/glintr100.onnx recognition\n","find model: ./insightface_func/models/antelope/scrfd_10g_bnkps.onnx detection\n","set det-size: (640, 640)\n"],"name":"stdout"},{"output_type":"stream","text":["\r 0%| | 0/594 [00:00>>> Building video ./output/multi_test_multispecific.mp4\n","[MoviePy] Writing audio in multi_test_multispecificTEMP_MPY_wvf_snd.mp3\n"],"name":"stdout"},{"output_type":"stream","text":["100%|██████████| 438/438 [00:00<00:00, 832.16it/s]"],"name":"stderr"},{"output_type":"stream","text":["[MoviePy] Done.\n","[MoviePy] Writing video ./output/multi_test_multispecific.mp4\n"],"name":"stdout"},{"output_type":"stream","text":["\n","100%|██████████| 595/595 [00:53<00:00, 11.11it/s]\n"],"name":"stderr"},{"output_type":"stream","text":["[MoviePy] Done.\n","[MoviePy] >>>> Video ready: ./output/multi_test_multispecific.mp4 \n","\n"],"name":"stdout"}]},{"cell_type":"code","metadata":{"id":"Rty2GsyZZrI6"},"source":[],"execution_count":null,"outputs":[]}]} \ No newline at end of file diff --git a/README.md b/README.md index f01c43c..d7842c5 100644 --- a/README.md +++ b/README.md @@ -2,9 +2,14 @@ ## Proceedings of the 28th ACM International Conference on Multimedia **The official repository with Pytorch** -*Our method can realize **arbitrary face swapping** on images and videos with **one single trained model**.* +**Our method can realize **arbitrary face swapping** on images and videos with **one single trained model**.** -Currently, only the test code is available. Training scripts are coming soon +Training and test code are now available! +[ google colab logo](https://colab.research.google.com/github/neuralchen/SimSwap/blob/main/train.ipynb) + +We are working with our incoming paper SimSwap++, keeping expecting! + +The high resolution version of ***SimSwap-HQ*** is supported! [![simswaplogo](/docs/img/logo1.png)](https://github.com/neuralchen/SimSwap) @@ -16,16 +21,34 @@ 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.*** +***In the event of violation of the legal and ethical requirements of the user's country or region, this code repository is exempt from liability*** + +***Please do not ignore the content at the end of this README!*** + +If you find this project useful, please star it. It is the greatest appreciation of our work. + ## Top News -**`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. +**`2022-04-21`**: For resource limited users, we provide the cropped VGGFace2-224 dataset [[Google Driver] VGGFace2-224 (10.8G)](https://drive.google.com/file/d/19pWvdEHS-CEG6tW3PdxdtZ5QEymVjImc/view?usp=sharing) [[Baidu Driver] ](https://pan.baidu.com/s/1OiwLJHVBSYB4AY2vEcfN0A) [Password: lrod]. -**`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. +**`2022-04-20`**: Training scripts are now available. We highly recommend that you guys train the simswap model with our released high quality dataset [VGGFace2-HQ](https://github.com/NNNNAI/VGGFace2-HQ). -**`2021-07-02`**: We have added a hyper parameter to allow users to choose whether to add the simswap logo as a watermark, please go to the section "About watermark of simswap logo" of [Inference for image or video face swapping](./docs/guidance/usage.md) for details. +**`2021-11-24`**: We have trained a beta version of ***SimSwap-HQ*** on [VGGFace2-HQ](https://github.com/NNNNAI/VGGFace2-HQ) and open sourced the checkpoint of this model (if you think the Simswap 512 is cool, please star our [VGGFace2-HQ](https://github.com/NNNNAI/VGGFace2-HQ) repo). Please don’t forget to go to [Preparation](./docs/guidance/preparation.md) and [Inference for image or video face swapping](./docs/guidance/usage.md) to check the latest set up. -**`2021-06-20`**: We release the scripts for arbitrary video and image processing, and a colab demo. +**`2021-11-23`**: The google drive link of [VGGFace2-HQ](https://github.com/NNNNAI/VGGFace2-HQ) is released. + +**`2021-11-17`**: We released a high resolution face dataset [VGGFace2-HQ](https://github.com/NNNNAI/VGGFace2-HQ) and the method to generate this dataset. This dataset is for research purpose. + +**`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)) + +## The first open source high resolution dataset for face swapping!!! +## High Resolution Dataset [VGGFace2-HQ](https://github.com/NNNNAI/VGGFace2-HQ) + +[![logo](./docs/img/vggface2_hq_compare.png)](https://github.com/NNNNAI/VGGFace2-HQ) ## Dependencies - python3.6+ @@ -38,14 +61,51 @@ Our paper can be downloaded from [[Arxiv]](https://arxiv.org/pdf/2106.06340v1.pd - moviepy - insightface -## Usage +## Training + +[Preparation](./docs/guidance/preparation.md) + +The training script is slightly different from the original version, e.g., we replace the patch discriminator with the projected discriminator, which saves a lot of hardware overhead and achieves slightly better results. + +In order to ensure the normal training, the batch size must be greater than 1. + +Friendly reminder, due to the difference in training settings, the user-trained model will have subtle differences in visual effects from the pre-trained model we provide. + +- Train 224 models with VGGFace2 224*224 [[Google Driver] VGGFace2-224 (10.8G)](https://drive.google.com/file/d/19pWvdEHS-CEG6tW3PdxdtZ5QEymVjImc/view?usp=sharing) [[Baidu Driver] ](https://pan.baidu.com/s/1OiwLJHVBSYB4AY2vEcfN0A) [Password: lrod] + +For faster convergence and better results, a large batch size (more than 16) is recommended! + +***We recommend training more than 400K iterations (batch size is 16), 600K~800K will be better, more iterations will not be recommended.*** + + +``` +python train.py --name simswap224_test --batchSize 8 --gpu_ids 0 --dataset /path/to/VGGFace2HQ --Gdeep False +``` + +[Colab demo for training 224 model](https://colab.research.google.com/github/neuralchen/SimSwap/blob/main/train.ipynb) + +For faster convergence and better results, a large batch size (more than 16) is recommended! + +- Train 512 models with VGGFace2-HQ 512*512 [VGGFace2-HQ](https://github.com/NNNNAI/VGGFace2-HQ). +``` +python train.py --name simswap512_test --batchSize 16 --gpu_ids 0 --dataset /path/to/VGGFace2HQ --Gdeep True +``` + + + +## Inference with a pretrained SimSwap model [Preparation](./docs/guidance/preparation.md) [Inference for image or video face swapping](./docs/guidance/usage.md) [Colab demo](https://colab.research.google.com/github/neuralchen/SimSwap/blob/main/SimSwap%20colab.ipynb) -Training: **coming soon** +
Stronger feature
+ +[Colab for 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) + ## Video @@ -56,6 +116,10 @@ Training: **coming soon** +
+ + +
## Results ![Results1](/docs/img/results1.PNG) @@ -71,7 +135,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) @@ -101,14 +165,7 @@ For academic and non-commercial use only.The whole project is under the CC-BY-NC Yanhao Ge}, title = {SimSwap: An Efficient Framework For High Fidelity Face Swapping}, booktitle = {{MM} '20: The 28th {ACM} International Conference on Multimedia}, - pages = {2003--2011}, - publisher = {{ACM}}, - year = {2020}, - url = {https://doi.org/10.1145/3394171.3413630}, - doi = {10.1145/3394171.3413630}, - timestamp = {Thu, 15 Oct 2020 16:32:08 +0200}, - biburl = {https://dblp.org/rec/conf/mm/ChenCNG20.bib}, - bibsource = {dblp computer science bibliography, https://dblp.org} + year = {2020} } ``` @@ -120,7 +177,19 @@ For academic and non-commercial use only.The whole project is under the CC-BY-NC [![title](/docs/img/title.png)](https://github.com/neuralchen/ASMAGAN) +**Please visit our AAAI2021 sketch based rendering project** + +[![logo](./docs/img/girl2.gif)](https://github.com/TZYSJTU/Sketch-Generation-with-Drawing-Process-Guided-by-Vector-Flow-and-Grayscale) +[![title](/docs/img/girl2-RGB.png)](https://github.com/TZYSJTU/Sketch-Generation-with-Drawing-Process-Guided-by-Vector-Flow-and-Grayscale) + +**Please visit our high resolution face dataset VGGFace2-HQ** + +[![logo](./docs/img/vggface2_hq_compare.png)](https://github.com/NNNNAI/VGGFace2-HQ) + Learn about our other projects + +[[VGGFace2-HQ]](https://github.com/NNNNAI/VGGFace2-HQ); + [[RainNet]](https://neuralchen.github.io/RainNet); [[Sketch Generation]](https://github.com/TZYSJTU/Sketch-Generation-with-Drawing-Process-Guided-by-Vector-Flow-and-Grayscale); @@ -140,5 +209,8 @@ Learn about our other projects ## Acknowledgements +* [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) diff --git a/SimSwap colab.ipynb b/SimSwap colab.ipynb index 04934ac..6df70c1 100644 --- a/SimSwap colab.ipynb +++ b/SimSwap colab.ipynb @@ -228,7 +228,8 @@ "\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" + "!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": [ @@ -407,36 +408,37 @@ "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", + "crop_size = opt.crop_size\n", "\n", "torch.nn.Module.dump_patches = True\n", "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, size=(112,112))\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": [ diff --git a/cog.yaml b/cog.yaml new file mode 100644 index 0000000..280c369 --- /dev/null +++ b/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" diff --git a/data/CelebA_class.py b/data/CelebA_class.py deleted file mode 100644 index bb546a2..0000000 --- a/data/CelebA_class.py +++ /dev/null @@ -1,94 +0,0 @@ -import torch -from torch.utils.data import Dataset -import os -import numpy as np -import random -from torchvision import transforms -from PIL import Image -import cv2 - -class FaceDataSet(Dataset): - def __init__(self, dataset_path, batch_size): - super(FaceDataSet, self).__init__() - - - - '''picture_dir_list = [] - for i in range(self.people_num): - picture_dir_list.append('/data/home/renwangchen/vgg_align_224/'+self.people_list[i]) - - self.people_pic_list = [] - for i in range(self.people_num): - pic_list = os.listdir(picture_dir_list[i]) - person_pic_list = [] - for j in range(len(pic_list)): - pic_dir = os.path.join(picture_dir_list[i], pic_list[j]) - person_pic_list.append(pic_dir) - self.people_pic_list.append(person_pic_list)''' - - pic_dir = '/data/home/renwangchen/CelebA_224/' - latent_dir = '/data/home/renwangchen/CelebA_latent/' - - tmp_list = os.listdir(pic_dir) - self.pic_list = [] - self.latent_list = [] - for i in range(len(tmp_list)): - self.pic_list.append(pic_dir + tmp_list[i]) - self.latent_list.append(latent_dir + tmp_list[i][:-3] + 'npy') - - self.pic_list = self.pic_list[:29984] - '''for i in range(29984): - print(self.pic_list[i])''' - self.latent_list = self.latent_list[:29984] - - self.people_num = len(self.pic_list) - - self.type = 1 - self.bs = batch_size - self.count = 0 - - self.transformer = transforms.Compose([ - transforms.ToTensor(), - transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) - ]) - - def __getitem__(self, index): - p1 = random.randint(0, self.people_num - 1) - p2 = p1 - - if self.type == 0: - # load pictures from the same folder - pass - else: - # load pictures from different folders - p2 = p1 - while p2 == p1: - p2 = random.randint(0, self.people_num - 1) - - pic_id_dir = self.pic_list[p1] - pic_att_dir = self.pic_list[p2] - latent_id_dir = self.latent_list[p1] - latent_att_dir = self.latent_list[p2] - - img_id = Image.open(pic_id_dir).convert('RGB') - img_id = self.transformer(img_id) - latent_id = np.load(latent_id_dir) - latent_id = latent_id / np.linalg.norm(latent_id) - latent_id = torch.from_numpy(latent_id) - - img_att = Image.open(pic_att_dir).convert('RGB') - img_att = self.transformer(img_att) - latent_att = np.load(latent_att_dir) - latent_att = latent_att / np.linalg.norm(latent_att) - latent_att = torch.from_numpy(latent_att) - - self.count += 1 - data_type = self.type - if self.count == self.bs: - self.type = 1 - self.type - self.count = 0 - - return img_id, img_att, latent_id, latent_att, data_type - - def __len__(self): - return len(self.pic_list) \ No newline at end of file diff --git a/data/aligned_dataset.py b/data/aligned_dataset.py deleted file mode 100644 index 29785c1..0000000 --- a/data/aligned_dataset.py +++ /dev/null @@ -1,76 +0,0 @@ -import os.path -from data.base_dataset import BaseDataset, get_params, get_transform, normalize -from data.image_folder import make_dataset -from PIL import Image - -class AlignedDataset(BaseDataset): - def initialize(self, opt): - self.opt = opt - self.root = opt.dataroot - - ### input A (label maps) - dir_A = '_A' if self.opt.label_nc == 0 else '_label' - self.dir_A = os.path.join(opt.dataroot, opt.phase + dir_A) - self.A_paths = sorted(make_dataset(self.dir_A)) - - ### input B (real images) - if opt.isTrain or opt.use_encoded_image: - dir_B = '_B' if self.opt.label_nc == 0 else '_img' - self.dir_B = os.path.join(opt.dataroot, opt.phase + dir_B) - self.B_paths = sorted(make_dataset(self.dir_B)) - - ### instance maps - if not opt.no_instance: - self.dir_inst = os.path.join(opt.dataroot, opt.phase + '_inst') - self.inst_paths = sorted(make_dataset(self.dir_inst)) - - ### load precomputed instance-wise encoded features - if opt.load_features: - self.dir_feat = os.path.join(opt.dataroot, opt.phase + '_feat') - print('----------- loading features from %s ----------' % self.dir_feat) - self.feat_paths = sorted(make_dataset(self.dir_feat)) - - self.dataset_size = len(self.A_paths) - - def __getitem__(self, index): - ### input A (label maps) - A_path = self.A_paths[index] - A = Image.open(A_path) - params = get_params(self.opt, A.size) - if self.opt.label_nc == 0: - transform_A = get_transform(self.opt, params) - A_tensor = transform_A(A.convert('RGB')) - else: - transform_A = get_transform(self.opt, params, method=Image.NEAREST, normalize=False) - A_tensor = transform_A(A) * 255.0 - - B_tensor = inst_tensor = feat_tensor = 0 - ### input B (real images) - if self.opt.isTrain or self.opt.use_encoded_image: - B_path = self.B_paths[index] - B = Image.open(B_path).convert('RGB') - transform_B = get_transform(self.opt, params) - B_tensor = transform_B(B) - - ### if using instance maps - if not self.opt.no_instance: - inst_path = self.inst_paths[index] - inst = Image.open(inst_path) - inst_tensor = transform_A(inst) - - if self.opt.load_features: - feat_path = self.feat_paths[index] - feat = Image.open(feat_path).convert('RGB') - norm = normalize() - feat_tensor = norm(transform_A(feat)) - - input_dict = {'label': A_tensor, 'inst': inst_tensor, 'image': B_tensor, - 'feat': feat_tensor, 'path': A_path} - - return input_dict - - def __len__(self): - return len(self.A_paths) // self.opt.batchSize * self.opt.batchSize - - def name(self): - return 'AlignedDataset' \ No newline at end of file diff --git a/data/base_dataset.py b/data/base_dataset.py deleted file mode 100644 index ece8813..0000000 --- a/data/base_dataset.py +++ /dev/null @@ -1,90 +0,0 @@ -import torch.utils.data as data -from PIL import Image -import torchvision.transforms as transforms -import numpy as np -import random - -class BaseDataset(data.Dataset): - def __init__(self): - super(BaseDataset, self).__init__() - - def name(self): - return 'BaseDataset' - - def initialize(self, opt): - pass - -def get_params(opt, size): - w, h = size - new_h = h - new_w = w - if opt.resize_or_crop == 'resize_and_crop': - new_h = new_w = opt.loadSize - elif opt.resize_or_crop == 'scale_width_and_crop': - new_w = opt.loadSize - new_h = opt.loadSize * h // w - - x = random.randint(0, np.maximum(0, new_w - opt.fineSize)) - y = random.randint(0, np.maximum(0, new_h - opt.fineSize)) - - flip = random.random() > 0.5 - return {'crop_pos': (x, y), 'flip': flip} - -def get_transform(opt, params, method=Image.BICUBIC, normalize=True): - transform_list = [] - if 'resize' in opt.resize_or_crop: - osize = [opt.loadSize, opt.loadSize] - transform_list.append(transforms.Scale(osize, method)) - elif 'scale_width' in opt.resize_or_crop: - transform_list.append(transforms.Lambda(lambda img: __scale_width(img, opt.loadSize, method))) - - if 'crop' in opt.resize_or_crop: - transform_list.append(transforms.Lambda(lambda img: __crop(img, params['crop_pos'], opt.fineSize))) - - if opt.resize_or_crop == 'none': - base = float(2 ** opt.n_downsample_global) - if opt.netG == 'local': - base *= (2 ** opt.n_local_enhancers) - transform_list.append(transforms.Lambda(lambda img: __make_power_2(img, base, method))) - - if opt.isTrain and not opt.no_flip: - transform_list.append(transforms.Lambda(lambda img: __flip(img, params['flip']))) - - transform_list += [transforms.ToTensor()] - - if normalize: - transform_list += [transforms.Normalize((0.5, 0.5, 0.5), - (0.5, 0.5, 0.5))] - return transforms.Compose(transform_list) - -def normalize(): - return transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) - -def __make_power_2(img, base, method=Image.BICUBIC): - ow, oh = img.size - h = int(round(oh / base) * base) - w = int(round(ow / base) * base) - if (h == oh) and (w == ow): - return img - return img.resize((w, h), method) - -def __scale_width(img, target_width, method=Image.BICUBIC): - ow, oh = img.size - if (ow == target_width): - return img - w = target_width - h = int(target_width * oh / ow) - return img.resize((w, h), method) - -def __crop(img, pos, size): - ow, oh = img.size - x1, y1 = pos - tw = th = size - if (ow > tw or oh > th): - return img.crop((x1, y1, x1 + tw, y1 + th)) - return img - -def __flip(img, flip): - if flip: - return img.transpose(Image.FLIP_LEFT_RIGHT) - return img diff --git a/data/data_loader.py b/data/data_loader.py deleted file mode 100644 index 2a4433a..0000000 --- a/data/data_loader.py +++ /dev/null @@ -1,7 +0,0 @@ - -def CreateDataLoader(opt): - from data.custom_dataset_data_loader import CustomDatasetDataLoader - data_loader = CustomDatasetDataLoader() - print(data_loader.name()) - data_loader.initialize(opt) - return data_loader diff --git a/data/data_loader_Swapping.py b/data/data_loader_Swapping.py new file mode 100644 index 0000000..8788b95 --- /dev/null +++ b/data/data_loader_Swapping.py @@ -0,0 +1,125 @@ +import os +import glob +import torch +import random +from PIL import Image +from torch.utils import data +from torchvision import transforms as T + +class data_prefetcher(): + def __init__(self, loader): + self.loader = loader + self.dataiter = iter(loader) + self.stream = torch.cuda.Stream() + self.mean = torch.tensor([0.485, 0.456, 0.406]).cuda().view(1,3,1,1) + self.std = torch.tensor([0.229, 0.224, 0.225]).cuda().view(1,3,1,1) + # With Amp, it isn't necessary to manually convert data to half. + # if args.fp16: + # self.mean = self.mean.half() + # self.std = self.std.half() + self.num_images = len(loader) + self.preload() + + def preload(self): + try: + self.src_image1, self.src_image2 = next(self.dataiter) + except StopIteration: + self.dataiter = iter(self.loader) + self.src_image1, self.src_image2 = next(self.dataiter) + + with torch.cuda.stream(self.stream): + self.src_image1 = self.src_image1.cuda(non_blocking=True) + self.src_image1 = self.src_image1.sub_(self.mean).div_(self.std) + self.src_image2 = self.src_image2.cuda(non_blocking=True) + self.src_image2 = self.src_image2.sub_(self.mean).div_(self.std) + + def next(self): + torch.cuda.current_stream().wait_stream(self.stream) + src_image1 = self.src_image1 + src_image2 = self.src_image2 + self.preload() + return src_image1, src_image2 + + def __len__(self): + """Return the number of images.""" + return self.num_images + +class SwappingDataset(data.Dataset): + """Dataset class for the Artworks dataset and content dataset.""" + + def __init__(self, + image_dir, + img_transform, + subffix='jpg', + random_seed=1234): + """Initialize and preprocess the Swapping dataset.""" + self.image_dir = image_dir + self.img_transform = img_transform + self.subffix = subffix + self.dataset = [] + self.random_seed = random_seed + self.preprocess() + self.num_images = len(self.dataset) + + def preprocess(self): + """Preprocess the Swapping dataset.""" + print("processing Swapping dataset images...") + + temp_path = os.path.join(self.image_dir,'*/') + pathes = glob.glob(temp_path) + self.dataset = [] + for dir_item in pathes: + join_path = glob.glob(os.path.join(dir_item,'*.jpg')) + print("processing %s"%dir_item,end='\r') + temp_list = [] + for item in join_path: + temp_list.append(item) + self.dataset.append(temp_list) + random.seed(self.random_seed) + random.shuffle(self.dataset) + print('Finished preprocessing the Swapping dataset, total dirs number: %d...'%len(self.dataset)) + + def __getitem__(self, index): + """Return two src domain images and two dst domain images.""" + dir_tmp1 = self.dataset[index] + dir_tmp1_len = len(dir_tmp1) + + filename1 = dir_tmp1[random.randint(0,dir_tmp1_len-1)] + filename2 = dir_tmp1[random.randint(0,dir_tmp1_len-1)] + image1 = self.img_transform(Image.open(filename1)) + image2 = self.img_transform(Image.open(filename2)) + return image1, image2 + + def __len__(self): + """Return the number of images.""" + return self.num_images + +def GetLoader( dataset_roots, + batch_size=16, + dataloader_workers=8, + random_seed = 1234 + ): + """Build and return a data loader.""" + + num_workers = dataloader_workers + data_root = dataset_roots + random_seed = random_seed + + c_transforms = [] + + c_transforms.append(T.ToTensor()) + c_transforms = T.Compose(c_transforms) + + content_dataset = SwappingDataset( + data_root, + c_transforms, + "jpg", + random_seed) + content_data_loader = data.DataLoader(dataset=content_dataset,batch_size=batch_size, + drop_last=True,shuffle=True,num_workers=num_workers,pin_memory=True) + prefetcher = data_prefetcher(content_data_loader) + return prefetcher + +def denorm(x): + out = (x + 1) / 2 + return out.clamp_(0, 1) \ No newline at end of file diff --git a/docs/css/page.css b/docs/css/page.css index 254eb91..614330a 100644 --- a/docs/css/page.css +++ b/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; */ } \ No newline at end of file diff --git a/docs/guidance/preparation.md b/docs/guidance/preparation.md index 63fd8a6..56dca98 100644 --- a/docs/guidance/preparation.md +++ b/docs/guidance/preparation.md @@ -9,11 +9,14 @@ conda activate simswap conda install pytorch==1.8.0 torchvision==0.9.0 torchaudio==0.8.0 cudatoolkit=10.2 -c pytorch (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. +- These settings are tested valid on both Windows and Ubuntu. ### Pretrained model There are two archive files in the drive: **checkpoints.zip** and **arcface_checkpoint.tar** @@ -24,5 +27,11 @@ There are two archive files in the drive: **checkpoints.zip** and **arcface_chec [[Google Drive]](https://drive.google.com/drive/folders/1jV6_0FIMPC53FZ2HzZNJZGMe55bbu17R?usp=sharing) [[Baidu Drive]](https://pan.baidu.com/s/1wFV11RVZMHqd-ky4YpLdcA) Password: ```jd2v``` +**Simswap 512 (optional)** + +The checkpoint of **Simswap 512 beta version** has been uploaded in [Github release](https://github.com/neuralchen/SimSwap/releases/download/512_beta/512.zip).If you want to experience Simswap 512, feel free to try. +- **Unzip 512.zip, place it in the root dir ./checkpoints**. + + ### 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). diff --git a/docs/guidance/usage.md b/docs/guidance/usage.md index 558ce46..fd60869 100644 --- a/docs/guidance/usage.md +++ b/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 --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 --crop_size 224 --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 --crop_size 224 --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 --crop_size 224 --use_mask --pic_specific_path ./demo_file/specific1.png --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 --crop_size 224 --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,47 @@ 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 --crop_size 224 --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 --crop_size 224 --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 --crop_size 224 --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 --crop_size 224 --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 using Simswap 512 (beta version) +We trained a beta version of Simswap 512 on [VGGFace2-HQ](https://github.com/NNNNAI/VGGFace2-HQ) and open sourced the model (if you think the Simswap 512 is cool, please star our [VGGFace2-HQ](https://github.com/NNNNAI/VGGFace2-HQ) repo). + +The usage of applying Simswap 512 is to modify the value of the argument: "***--crop_size***" to 512 , take the command line of "Swap **multi specific** face with **multi specific id** within one image." as an example, the following command line can get the result without watermark: +``` +python test_wholeimage_swap_multispecific.py --crop_size 512 --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 +``` +The effect of Simswap 512 is shown below. + + ### 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 --crop_size 224 --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. + + + ### Parameters | Parameters | Function | | :---- | :---- | @@ -95,6 +109,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. diff --git a/docs/img/1.png b/docs/img/1.png deleted file mode 100644 index fbb9a0e..0000000 Binary files a/docs/img/1.png and /dev/null differ diff --git a/docs/img/2.png b/docs/img/2.png deleted file mode 100644 index 94e6ff7..0000000 Binary files a/docs/img/2.png and /dev/null differ diff --git a/docs/img/HRGT_201009539_201009571.webp b/docs/img/HRGT_201009539_201009571.webp deleted file mode 100644 index ebaf772..0000000 Binary files a/docs/img/HRGT_201009539_201009571.webp and /dev/null differ diff --git a/docs/img/HRGT_201010607_201010639.webp b/docs/img/HRGT_201010607_201010639.webp deleted file mode 100644 index 8effc98..0000000 Binary files a/docs/img/HRGT_201010607_201010639.webp and /dev/null differ diff --git a/docs/img/HRGT_201011587_201011619.webp b/docs/img/HRGT_201011587_201011619.webp deleted file mode 100644 index 5922a29..0000000 Binary files a/docs/img/HRGT_201011587_201011619.webp and /dev/null differ diff --git a/docs/img/HRGT_201108607_201108655.webp b/docs/img/HRGT_201108607_201108655.webp deleted file mode 100644 index 4ad3cda..0000000 Binary files a/docs/img/HRGT_201108607_201108655.webp and /dev/null differ diff --git a/docs/img/HRGT_201109091_201109123.webp b/docs/img/HRGT_201109091_201109123.webp deleted file mode 100644 index 40d59f2..0000000 Binary files a/docs/img/HRGT_201109091_201109123.webp and /dev/null differ diff --git a/docs/img/HRGT_201110059_201110091.webp b/docs/img/HRGT_201110059_201110091.webp deleted file mode 100644 index 049d592..0000000 Binary files a/docs/img/HRGT_201110059_201110091.webp and /dev/null differ diff --git a/docs/img/LRGT_201009539_201009571.webp b/docs/img/LRGT_201009539_201009571.webp deleted file mode 100644 index 66a9ff7..0000000 Binary files a/docs/img/LRGT_201009539_201009571.webp and /dev/null differ diff --git a/docs/img/LRGT_201010607_201010639.webp b/docs/img/LRGT_201010607_201010639.webp deleted file mode 100644 index 7848f33..0000000 Binary files a/docs/img/LRGT_201010607_201010639.webp and /dev/null differ diff --git a/docs/img/LRGT_201011587_201011619.webp b/docs/img/LRGT_201011587_201011619.webp deleted file mode 100644 index 6dcc9fe..0000000 Binary files a/docs/img/LRGT_201011587_201011619.webp and /dev/null differ diff --git a/docs/img/LRGT_201108607_201108655.webp b/docs/img/LRGT_201108607_201108655.webp deleted file mode 100644 index 73de61e..0000000 Binary files a/docs/img/LRGT_201108607_201108655.webp and /dev/null differ diff --git a/docs/img/LRGT_201109091_201109123.webp b/docs/img/LRGT_201109091_201109123.webp deleted file mode 100644 index 33b83c3..0000000 Binary files a/docs/img/LRGT_201109091_201109123.webp and /dev/null differ diff --git a/docs/img/girl2-RGB.png b/docs/img/girl2-RGB.png new file mode 100644 index 0000000..7dfee55 Binary files /dev/null and b/docs/img/girl2-RGB.png differ diff --git a/docs/img/girl2.gif b/docs/img/girl2.gif new file mode 100644 index 0000000..fc5c5e0 Binary files /dev/null and b/docs/img/girl2.gif differ diff --git a/docs/img/id/Iron_man.jpg b/docs/img/id/Iron_man.jpg new file mode 100644 index 0000000..77a0f20 Binary files /dev/null and b/docs/img/id/Iron_man.jpg differ diff --git a/docs/img/id/wuyifan.png b/docs/img/id/wuyifan.png new file mode 100644 index 0000000..f2c2e8a Binary files /dev/null and b/docs/img/id/wuyifan.png differ diff --git a/docs/img/id/zhoujielun.jpg b/docs/img/id/zhoujielun.jpg index 6e5da6c..795d14b 100644 Binary files a/docs/img/id/zhoujielun.jpg and b/docs/img/id/zhoujielun.jpg differ diff --git a/docs/img/mama_mask_short.webp b/docs/img/mama_mask_short.webp new file mode 100644 index 0000000..cc7c210 Binary files /dev/null and b/docs/img/mama_mask_short.webp differ diff --git a/docs/img/mama_mask_wuyifan_short.webp b/docs/img/mama_mask_wuyifan_short.webp new file mode 100644 index 0000000..c423572 Binary files /dev/null and b/docs/img/mama_mask_wuyifan_short.webp differ diff --git a/docs/img/result_whole_swap_multispecific_512.jpg b/docs/img/result_whole_swap_multispecific_512.jpg new file mode 100644 index 0000000..37560d3 Binary files /dev/null and b/docs/img/result_whole_swap_multispecific_512.jpg differ diff --git a/docs/img/vggface2_hq_compare.png b/docs/img/vggface2_hq_compare.png new file mode 100644 index 0000000..4f0b427 Binary files /dev/null and b/docs/img/vggface2_hq_compare.png differ diff --git a/docs/index.html b/docs/index.html index c5387cf..2b48b1a 100644 --- a/docs/index.html +++ b/docs/index.html @@ -150,19 +150,19 @@

Single Face Video Swap

-
+
- +
- +
@@ -172,11 +172,23 @@
- - - +
+
+
+
+
+ +
+
+
+
+ +
+
+
+ +
+
@@ -319,6 +331,15 @@ + \ No newline at end of file diff --git a/docs/js/which-image.js b/docs/js/which-image.js index 698f7d1..867bf20 100644 --- a/docs/js/which-image.js +++ b/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'; + + } + } \ No newline at end of file diff --git a/download-weights.sh b/download-weights.sh new file mode 100755 index 0000000..54e6ea5 --- /dev/null +++ b/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 diff --git a/insightface_func/face_detect_crop_multi.py b/insightface_func/face_detect_crop_multi.py index acbf615..d75c36b 100644 --- a/insightface_func/face_detect_crop_multi.py +++ b/insightface_func/face_detect_crop_multi.py @@ -1,3 +1,11 @@ +''' +Author: Naiyuan liu +Github: https://github.com/NNNNAI +Date: 2021-11-23 17:03:58 +LastEditors: Naiyuan liu +LastEditTime: 2021-11-24 16:45:41 +Description: +''' from __future__ import division import collections import numpy as np @@ -6,7 +14,7 @@ import os import os.path as osp import cv2 from insightface.model_zoo import model_zoo -from insightface.utils import face_align +from insightface_func.utils import face_align_ffhqandnewarc as face_align __all__ = ['Face_detect_crop', 'Face'] @@ -40,8 +48,9 @@ class Face_detect_crop: self.det_model = self.models['detection'] - def prepare(self, ctx_id, det_thresh=0.5, det_size=(640, 640)): + def prepare(self, ctx_id, det_thresh=0.5, det_size=(640, 640), mode ='None'): self.det_thresh = det_thresh + self.mode = mode assert det_size is not None print('set det-size:', det_size) self.det_size = det_size @@ -73,7 +82,7 @@ class Face_detect_crop: kps = None if kpss is not None: kps = kpss[i] - M, _ = face_align.estimate_norm(kps, crop_size, mode ='None') + M, _ = face_align.estimate_norm(kps, crop_size, mode = self.mode) align_img = cv2.warpAffine(img, M, (crop_size, crop_size), borderValue=0.0) align_img_list.append(align_img) M_list.append(M) diff --git a/insightface_func/face_detect_crop_single.py b/insightface_func/face_detect_crop_single.py index 1fdd5d9..f3f648f 100644 --- a/insightface_func/face_detect_crop_single.py +++ b/insightface_func/face_detect_crop_single.py @@ -1,3 +1,11 @@ +''' +Author: Naiyuan liu +Github: https://github.com/NNNNAI +Date: 2021-11-23 17:03:58 +LastEditors: Naiyuan liu +LastEditTime: 2021-11-24 16:46:04 +Description: +''' from __future__ import division import collections import numpy as np @@ -6,7 +14,7 @@ import os import os.path as osp import cv2 from insightface.model_zoo import model_zoo -from insightface.utils import face_align +from insightface_func.utils import face_align_ffhqandnewarc as face_align __all__ = ['Face_detect_crop', 'Face'] @@ -40,8 +48,9 @@ class Face_detect_crop: self.det_model = self.models['detection'] - def prepare(self, ctx_id, det_thresh=0.5, det_size=(640, 640)): + def prepare(self, ctx_id, det_thresh=0.5, det_size=(640, 640), mode ='None'): self.det_thresh = det_thresh + self.mode = mode assert det_size is not None print('set det-size:', det_size) self.det_size = det_size @@ -82,7 +91,7 @@ class Face_detect_crop: kps = None if kpss is not None: kps = kpss[best_index] - M, _ = face_align.estimate_norm(kps, crop_size, mode ='None') + M, _ = face_align.estimate_norm(kps, crop_size, mode = self.mode) align_img = cv2.warpAffine(img, M, (crop_size, crop_size), borderValue=0.0) return [align_img], [M] diff --git a/insightface_func/utils/face_align_ffhqandnewarc.py b/insightface_func/utils/face_align_ffhqandnewarc.py new file mode 100644 index 0000000..2f5aeb7 --- /dev/null +++ b/insightface_func/utils/face_align_ffhqandnewarc.py @@ -0,0 +1,159 @@ +''' +Author: Naiyuan liu +Github: https://github.com/NNNNAI +Date: 2021-11-15 19:42:42 +LastEditors: Naiyuan liu +LastEditTime: 2021-11-15 20:01:47 +Description: +''' + +import cv2 +import numpy as np +from skimage import transform as trans + +src1 = np.array([[51.642, 50.115], [57.617, 49.990], [35.740, 69.007], + [51.157, 89.050], [57.025, 89.702]], + dtype=np.float32) +#<--left +src2 = np.array([[45.031, 50.118], [65.568, 50.872], [39.677, 68.111], + [45.177, 86.190], [64.246, 86.758]], + dtype=np.float32) + +#---frontal +src3 = np.array([[39.730, 51.138], [72.270, 51.138], [56.000, 68.493], + [42.463, 87.010], [69.537, 87.010]], + dtype=np.float32) + +#-->right +src4 = np.array([[46.845, 50.872], [67.382, 50.118], [72.737, 68.111], + [48.167, 86.758], [67.236, 86.190]], + dtype=np.float32) + +#-->right profile +src5 = np.array([[54.796, 49.990], [60.771, 50.115], [76.673, 69.007], + [55.388, 89.702], [61.257, 89.050]], + dtype=np.float32) + +src = np.array([src1, src2, src3, src4, src5]) +src_map = src + +ffhq_src = np.array([[192.98138, 239.94708], [318.90277, 240.1936], [256.63416, 314.01935], + [201.26117, 371.41043], [313.08905, 371.15118]]) +ffhq_src = np.expand_dims(ffhq_src, axis=0) + +# arcface_src = np.array( +# [[38.2946, 51.6963], [73.5318, 51.5014], [56.0252, 71.7366], +# [41.5493, 92.3655], [70.7299, 92.2041]], +# dtype=np.float32) + +# arcface_src = np.expand_dims(arcface_src, axis=0) + +# In[66]: + + +# lmk is prediction; src is template +def estimate_norm(lmk, image_size=112, mode='ffhq'): + assert lmk.shape == (5, 2) + tform = trans.SimilarityTransform() + lmk_tran = np.insert(lmk, 2, values=np.ones(5), axis=1) + min_M = [] + min_index = [] + min_error = float('inf') + if mode == 'ffhq': + # assert image_size == 112 + src = ffhq_src * image_size / 512 + else: + src = src_map * image_size / 112 + for i in np.arange(src.shape[0]): + tform.estimate(lmk, src[i]) + M = tform.params[0:2, :] + results = np.dot(M, lmk_tran.T) + results = results.T + error = np.sum(np.sqrt(np.sum((results - src[i])**2, axis=1))) + # print(error) + if error < min_error: + min_error = error + min_M = M + min_index = i + return min_M, min_index + + +def norm_crop(img, landmark, image_size=112, mode='ffhq'): + if mode == 'Both': + M_None, _ = estimate_norm(landmark, image_size, mode = 'newarc') + M_ffhq, _ = estimate_norm(landmark, image_size, mode='ffhq') + warped_None = cv2.warpAffine(img, M_None, (image_size, image_size), borderValue=0.0) + warped_ffhq = cv2.warpAffine(img, M_ffhq, (image_size, image_size), borderValue=0.0) + return warped_ffhq, warped_None + else: + M, pose_index = estimate_norm(landmark, image_size, mode) + warped = cv2.warpAffine(img, M, (image_size, image_size), borderValue=0.0) + return warped + +def square_crop(im, S): + if im.shape[0] > im.shape[1]: + height = S + width = int(float(im.shape[1]) / im.shape[0] * S) + scale = float(S) / im.shape[0] + else: + width = S + height = int(float(im.shape[0]) / im.shape[1] * S) + scale = float(S) / im.shape[1] + resized_im = cv2.resize(im, (width, height)) + det_im = np.zeros((S, S, 3), dtype=np.uint8) + det_im[:resized_im.shape[0], :resized_im.shape[1], :] = resized_im + return det_im, scale + + +def transform(data, center, output_size, scale, rotation): + scale_ratio = scale + rot = float(rotation) * np.pi / 180.0 + #translation = (output_size/2-center[0]*scale_ratio, output_size/2-center[1]*scale_ratio) + t1 = trans.SimilarityTransform(scale=scale_ratio) + cx = center[0] * scale_ratio + cy = center[1] * scale_ratio + t2 = trans.SimilarityTransform(translation=(-1 * cx, -1 * cy)) + t3 = trans.SimilarityTransform(rotation=rot) + t4 = trans.SimilarityTransform(translation=(output_size / 2, + output_size / 2)) + t = t1 + t2 + t3 + t4 + M = t.params[0:2] + cropped = cv2.warpAffine(data, + M, (output_size, output_size), + borderValue=0.0) + return cropped, M + + +def trans_points2d(pts, M): + new_pts = np.zeros(shape=pts.shape, dtype=np.float32) + for i in range(pts.shape[0]): + pt = pts[i] + new_pt = np.array([pt[0], pt[1], 1.], dtype=np.float32) + new_pt = np.dot(M, new_pt) + #print('new_pt', new_pt.shape, new_pt) + new_pts[i] = new_pt[0:2] + + return new_pts + + +def trans_points3d(pts, M): + scale = np.sqrt(M[0][0] * M[0][0] + M[0][1] * M[0][1]) + #print(scale) + new_pts = np.zeros(shape=pts.shape, dtype=np.float32) + for i in range(pts.shape[0]): + pt = pts[i] + new_pt = np.array([pt[0], pt[1], 1.], dtype=np.float32) + new_pt = np.dot(M, new_pt) + #print('new_pt', new_pt.shape, new_pt) + new_pts[i][0:2] = new_pt[0:2] + new_pts[i][2] = pts[i][2] * scale + + return new_pts + + +def trans_points(pts, M): + if pts.shape[1] == 2: + return trans_points2d(pts, M) + else: + return trans_points3d(pts, M) + diff --git a/models/base_model.py b/models/base_model.py index f3f6b53..0a6474a 100644 --- a/models/base_model.py +++ b/models/base_model.py @@ -37,15 +37,21 @@ class BaseModel(torch.nn.Module): def save(self, label): pass - + # helper saving function that can be used by subclasses - def save_network(self, network, network_label, epoch_label, gpu_ids): - save_filename = '%s_net_%s.pth' % (epoch_label, network_label) + def save_network(self, network, network_label, epoch_label, gpu_ids=None): + save_filename = '{}_net_{}.pth'.format(epoch_label, network_label) save_path = os.path.join(self.save_dir, save_filename) torch.save(network.cpu().state_dict(), save_path) - if len(gpu_ids) and torch.cuda.is_available(): + if torch.cuda.is_available(): network.cuda() + def save_optim(self, network, network_label, epoch_label, gpu_ids=None): + save_filename = '{}_optim_{}.pth'.format(epoch_label, network_label) + save_path = os.path.join(self.save_dir, save_filename) + torch.save(network.state_dict(), save_path) + + # helper loading function that can be used by subclasses def load_network(self, network, network_label, epoch_label, save_dir=''): save_filename = '%s_net_%s.pth' % (epoch_label, network_label) @@ -63,6 +69,47 @@ class BaseModel(torch.nn.Module): except: pretrained_dict = torch.load(save_path) model_dict = network.state_dict() + try: + pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict} + network.load_state_dict(pretrained_dict) + if self.opt.verbose: + print('Pretrained network %s has excessive layers; Only loading layers that are used' % network_label) + except: + print('Pretrained network %s has fewer layers; The following are not initialized:' % network_label) + for k, v in pretrained_dict.items(): + if v.size() == model_dict[k].size(): + model_dict[k] = v + + if sys.version_info >= (3,0): + not_initialized = set() + else: + from sets import Set + not_initialized = Set() + + for k, v in model_dict.items(): + if k not in pretrained_dict or v.size() != pretrained_dict[k].size(): + not_initialized.add(k.split('.')[0]) + + print(sorted(not_initialized)) + network.load_state_dict(model_dict) + + # helper loading function that can be used by subclasses + def load_optim(self, network, network_label, epoch_label, save_dir=''): + save_filename = '%s_optim_%s.pth' % (epoch_label, network_label) + if not save_dir: + save_dir = self.save_dir + save_path = os.path.join(save_dir, save_filename) + if not os.path.isfile(save_path): + print('%s not exists yet!' % save_path) + if network_label == 'G': + raise('Generator must exist!') + else: + #network.load_state_dict(torch.load(save_path)) + try: + network.load_state_dict(torch.load(save_path, map_location=torch.device("cpu"))) + except: + pretrained_dict = torch.load(save_path, map_location=torch.device("cpu")) + model_dict = network.state_dict() try: pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict} network.load_state_dict(pretrained_dict) diff --git a/models/fs_model.py b/models/fs_model.py index 49cb066..548dd92 100644 --- a/models/fs_model.py +++ b/models/fs_model.py @@ -4,10 +4,8 @@ import torch.nn as nn import torch.nn.functional as F import os from torch.autograd import Variable -from util.image_pool import ImagePool from .base_model import BaseModel from . import networks -from .fs_networks import Generator_Adain_Upsample, Discriminator class SpecificNorm(nn.Module): def __init__(self, epsilon=1e-8): @@ -52,6 +50,11 @@ class fsModel(BaseModel): device = torch.device("cuda:0") + if opt.crop_size == 224: + from .fs_networks import Generator_Adain_Upsample, Discriminator + elif opt.crop_size == 512: + from .fs_networks_512 import Generator_Adain_Upsample, Discriminator + # Generator network self.netG = Generator_Adain_Upsample(input_nc=3, output_nc=3, latent_size=512, n_blocks=9, deep=False) self.netG.to(device) @@ -197,7 +200,7 @@ class fsModel(BaseModel): #G_ID - img_fake_down = F.interpolate(img_fake, scale_factor=0.5) + img_fake_down = F.interpolate(img_fake, size=(112,112)) img_fake_down = self.spNorm(img_fake_down) latent_fake = self.netArc(img_fake_down) loss_G_ID = (1 - self.cosin_metric(latent_fake, latent_id)) diff --git a/models/fs_networks_512.py b/models/fs_networks_512.py new file mode 100644 index 0000000..1b9f9c8 --- /dev/null +++ b/models/fs_networks_512.py @@ -0,0 +1,232 @@ +''' +Author: Naiyuan liu +Github: https://github.com/NNNNAI +Date: 2021-11-23 16:55:48 +LastEditors: Naiyuan liu +LastEditTime: 2021-11-24 16:58:06 +Description: +''' +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" + +import torch +import torch.nn as nn + + +class InstanceNorm(nn.Module): + def __init__(self, epsilon=1e-8): + """ + @notice: avoid in-place ops. + https://discuss.pytorch.org/t/encounter-the-runtimeerror-one-of-the-variables-needed-for-gradient-computation-has-been-modified-by-an-inplace-operation/836/3 + """ + super(InstanceNorm, self).__init__() + self.epsilon = epsilon + + def forward(self, x): + x = x - torch.mean(x, (2, 3), True) + tmp = torch.mul(x, x) # or x ** 2 + tmp = torch.rsqrt(torch.mean(tmp, (2, 3), True) + self.epsilon) + return x * tmp + +class ApplyStyle(nn.Module): + """ + @ref: https://github.com/lernapparat/lernapparat/blob/master/style_gan/pytorch_style_gan.ipynb + """ + def __init__(self, latent_size, channels): + super(ApplyStyle, self).__init__() + self.linear = nn.Linear(latent_size, channels * 2) + + def forward(self, x, latent): + style = self.linear(latent) # style => [batch_size, n_channels*2] + shape = [-1, 2, x.size(1), 1, 1] + style = style.view(shape) # [batch_size, 2, n_channels, ...] + #x = x * (style[:, 0] + 1.) + style[:, 1] + x = x * (style[:, 0] * 1 + 1.) + style[:, 1] * 1 + return x + +class ResnetBlock_Adain(nn.Module): + def __init__(self, dim, latent_size, padding_type, activation=nn.ReLU(True)): + super(ResnetBlock_Adain, self).__init__() + + p = 0 + conv1 = [] + if padding_type == 'reflect': + conv1 += [nn.ReflectionPad2d(1)] + elif padding_type == 'replicate': + conv1 += [nn.ReplicationPad2d(1)] + elif padding_type == 'zero': + p = 1 + else: + raise NotImplementedError('padding [%s] is not implemented' % padding_type) + conv1 += [nn.Conv2d(dim, dim, kernel_size=3, padding = p), InstanceNorm()] + self.conv1 = nn.Sequential(*conv1) + self.style1 = ApplyStyle(latent_size, dim) + self.act1 = activation + + p = 0 + conv2 = [] + if padding_type == 'reflect': + conv2 += [nn.ReflectionPad2d(1)] + elif padding_type == 'replicate': + conv2 += [nn.ReplicationPad2d(1)] + elif padding_type == 'zero': + p = 1 + else: + raise NotImplementedError('padding [%s] is not implemented' % padding_type) + conv2 += [nn.Conv2d(dim, dim, kernel_size=3, padding=p), InstanceNorm()] + self.conv2 = nn.Sequential(*conv2) + self.style2 = ApplyStyle(latent_size, dim) + + + def forward(self, x, dlatents_in_slice): + y = self.conv1(x) + y = self.style1(y, dlatents_in_slice) + y = self.act1(y) + y = self.conv2(y) + y = self.style2(y, dlatents_in_slice) + out = x + y + return out + + + +class Generator_Adain_Upsample(nn.Module): + def __init__(self, input_nc, output_nc, latent_size, n_blocks=6, deep=False, + norm_layer=nn.BatchNorm2d, + padding_type='reflect'): + assert (n_blocks >= 0) + super(Generator_Adain_Upsample, self).__init__() + activation = nn.ReLU(True) + self.deep = deep + + self.first_layer = nn.Sequential(nn.ReflectionPad2d(3), nn.Conv2d(input_nc, 32, kernel_size=7, padding=0), + norm_layer(32), activation) + ### downsample + self.down0 = nn.Sequential(nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1), + norm_layer(64), activation) + self.down1 = nn.Sequential(nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1), + norm_layer(128), activation) + self.down2 = nn.Sequential(nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1), + norm_layer(256), activation) + self.down3 = nn.Sequential(nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1), + norm_layer(512), activation) + if self.deep: + self.down4 = nn.Sequential(nn.Conv2d(512, 512, kernel_size=3, stride=2, padding=1), + norm_layer(512), activation) + + ### resnet blocks + BN = [] + for i in range(n_blocks): + BN += [ + ResnetBlock_Adain(512, latent_size=latent_size, padding_type=padding_type, activation=activation)] + self.BottleNeck = nn.Sequential(*BN) + + if self.deep: + self.up4 = nn.Sequential( + nn.Upsample(scale_factor=2, mode='bilinear'), + nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1), + nn.BatchNorm2d(512), activation + ) + self.up3 = nn.Sequential( + nn.Upsample(scale_factor=2, mode='bilinear'), + nn.Conv2d(512, 256, kernel_size=3, stride=1, padding=1), + nn.BatchNorm2d(256), activation + ) + self.up2 = nn.Sequential( + nn.Upsample(scale_factor=2, mode='bilinear'), + nn.Conv2d(256, 128, kernel_size=3, stride=1, padding=1), + nn.BatchNorm2d(128), activation + ) + self.up1 = nn.Sequential( + nn.Upsample(scale_factor=2, mode='bilinear'), + nn.Conv2d(128, 64, kernel_size=3, stride=1, padding=1), + nn.BatchNorm2d(64), activation + ) + self.up0 = nn.Sequential( + nn.Upsample(scale_factor=2, mode='bilinear'), + nn.Conv2d(64, 32, kernel_size=3, stride=1, padding=1), + nn.BatchNorm2d(32), activation + ) + self.last_layer = nn.Sequential(nn.ReflectionPad2d(3), nn.Conv2d(32, output_nc, kernel_size=7, padding=0), + nn.Tanh()) + + def forward(self, input, dlatents): + x = input # 3*224*224 + + skip0 = self.first_layer(x) + skip1 = self.down0(skip0) + skip2 = self.down1(skip1) + skip3 = self.down2(skip2) + if self.deep: + skip4 = self.down3(skip3) + x = self.down4(skip4) + else: + x = self.down3(skip3) + + for i in range(len(self.BottleNeck)): + x = self.BottleNeck[i](x, dlatents) + + if self.deep: + x = self.up4(x) + x = self.up3(x) + x = self.up2(x) + x = self.up1(x) + x = self.up0(x) + x = self.last_layer(x) + x = (x + 1) / 2 + + return x + +class Discriminator(nn.Module): + def __init__(self, input_nc, norm_layer=nn.BatchNorm2d, use_sigmoid=False): + super(Discriminator, self).__init__() + + kw = 4 + padw = 1 + self.down1 = nn.Sequential( + nn.Conv2d(input_nc, 64, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True) + ) + self.down2 = nn.Sequential( + nn.Conv2d(64, 128, kernel_size=kw, stride=2, padding=padw), + norm_layer(128), nn.LeakyReLU(0.2, True) + ) + self.down3 = nn.Sequential( + nn.Conv2d(128, 256, kernel_size=kw, stride=2, padding=padw), + norm_layer(256), nn.LeakyReLU(0.2, True) + ) + self.down4 = nn.Sequential( + nn.Conv2d(256, 512, kernel_size=kw, stride=2, padding=padw), + norm_layer(512), nn.LeakyReLU(0.2, True) + ) + self.conv1 = nn.Sequential( + nn.Conv2d(512, 512, kernel_size=kw, stride=1, padding=padw), + norm_layer(512), + nn.LeakyReLU(0.2, True) + ) + + if use_sigmoid: + self.conv2 = nn.Sequential( + nn.Conv2d(512, 1, kernel_size=kw, stride=1, padding=padw), nn.Sigmoid() + ) + else: + self.conv2 = nn.Sequential( + nn.Conv2d(512, 1, kernel_size=kw, stride=1, padding=padw) + ) + + def forward(self, input): + out = [] + x = self.down1(input) + out.append(x) + x = self.down2(x) + out.append(x) + x = self.down3(x) + out.append(x) + x = self.down4(x) + out.append(x) + x = self.conv1(x) + out.append(x) + x = self.conv2(x) + out.append(x) + + return out diff --git a/models/fs_networks_fix.py b/models/fs_networks_fix.py new file mode 100644 index 0000000..c7b0525 --- /dev/null +++ b/models/fs_networks_fix.py @@ -0,0 +1,172 @@ +""" +Copyright (C) 2019 NVIDIA Corporation. All rights reserved. +Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +""" + +import torch +import torch.nn as nn + + +class InstanceNorm(nn.Module): + def __init__(self, epsilon=1e-8): + """ + @notice: avoid in-place ops. + https://discuss.pytorch.org/t/encounter-the-runtimeerror-one-of-the-variables-needed-for-gradient-computation-has-been-modified-by-an-inplace-operation/836/3 + """ + super(InstanceNorm, self).__init__() + self.epsilon = epsilon + + def forward(self, x): + x = x - torch.mean(x, (2, 3), True) + tmp = torch.mul(x, x) # or x ** 2 + tmp = torch.rsqrt(torch.mean(tmp, (2, 3), True) + self.epsilon) + return x * tmp + +class ApplyStyle(nn.Module): + """ + @ref: https://github.com/lernapparat/lernapparat/blob/master/style_gan/pytorch_style_gan.ipynb + """ + def __init__(self, latent_size, channels): + super(ApplyStyle, self).__init__() + self.linear = nn.Linear(latent_size, channels * 2) + + def forward(self, x, latent): + style = self.linear(latent) # style => [batch_size, n_channels*2] + shape = [-1, 2, x.size(1), 1, 1] + style = style.view(shape) # [batch_size, 2, n_channels, ...] + #x = x * (style[:, 0] + 1.) + style[:, 1] + x = x * (style[:, 0] * 1 + 1.) + style[:, 1] * 1 + return x + +class ResnetBlock_Adain(nn.Module): + def __init__(self, dim, latent_size, padding_type, activation=nn.ReLU(True)): + super(ResnetBlock_Adain, self).__init__() + + p = 0 + conv1 = [] + if padding_type == 'reflect': + conv1 += [nn.ReflectionPad2d(1)] + elif padding_type == 'replicate': + conv1 += [nn.ReplicationPad2d(1)] + elif padding_type == 'zero': + p = 1 + else: + raise NotImplementedError('padding [%s] is not implemented' % padding_type) + conv1 += [nn.Conv2d(dim, dim, kernel_size=3, padding = p), InstanceNorm()] + self.conv1 = nn.Sequential(*conv1) + self.style1 = ApplyStyle(latent_size, dim) + self.act1 = activation + + p = 0 + conv2 = [] + if padding_type == 'reflect': + conv2 += [nn.ReflectionPad2d(1)] + elif padding_type == 'replicate': + conv2 += [nn.ReplicationPad2d(1)] + elif padding_type == 'zero': + p = 1 + else: + raise NotImplementedError('padding [%s] is not implemented' % padding_type) + conv2 += [nn.Conv2d(dim, dim, kernel_size=3, padding=p), InstanceNorm()] + self.conv2 = nn.Sequential(*conv2) + self.style2 = ApplyStyle(latent_size, dim) + + + def forward(self, x, dlatents_in_slice): + y = self.conv1(x) + y = self.style1(y, dlatents_in_slice) + y = self.act1(y) + y = self.conv2(y) + y = self.style2(y, dlatents_in_slice) + out = x + y + return out + + + +class Generator_Adain_Upsample(nn.Module): + def __init__(self, input_nc, output_nc, latent_size, n_blocks=6, deep=False, + norm_layer=nn.BatchNorm2d, + padding_type='reflect'): + assert (n_blocks >= 0) + super(Generator_Adain_Upsample, self).__init__() + + activation = nn.ReLU(True) + + self.deep = deep + + self.first_layer = nn.Sequential(nn.ReflectionPad2d(3), nn.Conv2d(input_nc, 64, kernel_size=7, padding=0), + norm_layer(64), activation) + ### downsample + self.down1 = nn.Sequential(nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1), + norm_layer(128), activation) + self.down2 = nn.Sequential(nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1), + norm_layer(256), activation) + self.down3 = nn.Sequential(nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1), + norm_layer(512), activation) + + if self.deep: + self.down4 = nn.Sequential(nn.Conv2d(512, 512, kernel_size=3, stride=2, padding=1), + norm_layer(512), activation) + + ### resnet blocks + BN = [] + for i in range(n_blocks): + BN += [ + ResnetBlock_Adain(512, latent_size=latent_size, padding_type=padding_type, activation=activation)] + self.BottleNeck = nn.Sequential(*BN) + + if self.deep: + self.up4 = nn.Sequential( + nn.Upsample(scale_factor=2, mode='bilinear',align_corners=False), + nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1), + nn.BatchNorm2d(512), activation + ) + self.up3 = nn.Sequential( + nn.Upsample(scale_factor=2, mode='bilinear',align_corners=False), + nn.Conv2d(512, 256, kernel_size=3, stride=1, padding=1), + nn.BatchNorm2d(256), activation + ) + self.up2 = nn.Sequential( + nn.Upsample(scale_factor=2, mode='bilinear',align_corners=False), + nn.Conv2d(256, 128, kernel_size=3, stride=1, padding=1), + nn.BatchNorm2d(128), activation + ) + self.up1 = nn.Sequential( + nn.Upsample(scale_factor=2, mode='bilinear',align_corners=False), + nn.Conv2d(128, 64, kernel_size=3, stride=1, padding=1), + nn.BatchNorm2d(64), activation + ) + self.last_layer = nn.Sequential(nn.ReflectionPad2d(3), nn.Conv2d(64, output_nc, kernel_size=7, padding=0)) + + def forward(self, input, dlatents): + x = input # 3*224*224 + + skip1 = self.first_layer(x) + skip2 = self.down1(skip1) + skip3 = self.down2(skip2) + if self.deep: + skip4 = self.down3(skip3) + x = self.down4(skip4) + else: + x = self.down3(skip3) + bot = [] + bot.append(x) + features = [] + for i in range(len(self.BottleNeck)): + x = self.BottleNeck[i](x, dlatents) + bot.append(x) + + if self.deep: + x = self.up4(x) + features.append(x) + x = self.up3(x) + features.append(x) + x = self.up2(x) + features.append(x) + x = self.up1(x) + features.append(x) + x = self.last_layer(x) + # x = (x + 1) / 2 + + # return x, bot, features, dlatents + return x \ No newline at end of file diff --git a/models/projected_model.py b/models/projected_model.py new file mode 100644 index 0000000..477f63e --- /dev/null +++ b/models/projected_model.py @@ -0,0 +1,122 @@ +#!/usr/bin/env python3 +# -*- coding:utf-8 -*- +############################################################# +# File: fs_model_fix_idnorm_donggp_saveoptim copy.py +# Created Date: Wednesday January 12th 2022 +# Author: Chen Xuanhong +# Email: chenxuanhongzju@outlook.com +# Last Modified: Thursday, 21st April 2022 8:13:37 pm +# Modified By: Chen Xuanhong +# Copyright (c) 2022 Shanghai Jiao Tong University +############################################################# + + +import torch +import torch.nn as nn + +from .base_model import BaseModel +from .fs_networks_fix import Generator_Adain_Upsample + +from pg_modules.projected_discriminator import ProjectedDiscriminator + +def compute_grad2(d_out, x_in): + batch_size = x_in.size(0) + grad_dout = torch.autograd.grad( + outputs=d_out.sum(), inputs=x_in, + create_graph=True, retain_graph=True, only_inputs=True + )[0] + grad_dout2 = grad_dout.pow(2) + assert(grad_dout2.size() == x_in.size()) + reg = grad_dout2.view(batch_size, -1).sum(1) + return reg + +class fsModel(BaseModel): + def name(self): + return 'fsModel' + + def initialize(self, opt): + BaseModel.initialize(self, opt) + # if opt.resize_or_crop != 'none' or not opt.isTrain: # when training at full res this causes OOM + self.isTrain = opt.isTrain + + # Generator network + self.netG = Generator_Adain_Upsample(input_nc=3, output_nc=3, latent_size=512, n_blocks=9, deep=opt.Gdeep) + self.netG.cuda() + + # Id network + netArc_checkpoint = opt.Arc_path + netArc_checkpoint = torch.load(netArc_checkpoint, map_location=torch.device("cpu")) + self.netArc = netArc_checkpoint['model'].module + self.netArc = self.netArc.cuda() + self.netArc.eval() + self.netArc.requires_grad_(False) + if not self.isTrain: + pretrained_path = opt.checkpoints_dir + self.load_network(self.netG, 'G', opt.which_epoch, pretrained_path) + return + self.netD = ProjectedDiscriminator(diffaug=False, interp224=False, **{}) + # self.netD.feature_network.requires_grad_(False) + self.netD.cuda() + + + if self.isTrain: + # define loss functions + self.criterionFeat = nn.L1Loss() + self.criterionRec = nn.L1Loss() + + + # initialize optimizers + + # optimizer G + params = list(self.netG.parameters()) + self.optimizer_G = torch.optim.Adam(params, lr=opt.lr, betas=(opt.beta1, 0.99),eps=1e-8) + + # optimizer D + params = list(self.netD.parameters()) + self.optimizer_D = torch.optim.Adam(params, lr=opt.lr, betas=(opt.beta1, 0.99),eps=1e-8) + + # load networks + if opt.continue_train: + pretrained_path = '' if not self.isTrain else opt.load_pretrain + # print (pretrained_path) + self.load_network(self.netG, 'G', opt.which_epoch, pretrained_path) + self.load_network(self.netD, 'D', opt.which_epoch, pretrained_path) + self.load_optim(self.optimizer_G, 'G', opt.which_epoch, pretrained_path) + self.load_optim(self.optimizer_D, 'D', opt.which_epoch, pretrained_path) + torch.cuda.empty_cache() + + def cosin_metric(self, x1, x2): + #return np.dot(x1, x2) / (np.linalg.norm(x1) * np.linalg.norm(x2)) + return torch.sum(x1 * x2, dim=1) / (torch.norm(x1, dim=1) * torch.norm(x2, dim=1)) + + + + def save(self, which_epoch): + self.save_network(self.netG, 'G', which_epoch) + self.save_network(self.netD, 'D', which_epoch) + self.save_optim(self.optimizer_G, 'G', which_epoch) + self.save_optim(self.optimizer_D, 'D', which_epoch) + '''if self.gen_features: + self.save_network(self.netE, 'E', which_epoch, self.gpu_ids)''' + + def update_fixed_params(self): + # after fixing the global generator for a number of iterations, also start finetuning it + params = list(self.netG.parameters()) + if self.gen_features: + params += list(self.netE.parameters()) + self.optimizer_G = torch.optim.Adam(params, lr=self.opt.lr, betas=(self.opt.beta1, 0.999)) + if self.opt.verbose: + print('------------ Now also finetuning global generator -----------') + + def update_learning_rate(self): + lrd = self.opt.lr / self.opt.niter_decay + lr = self.old_lr - lrd + for param_group in self.optimizer_D.param_groups: + param_group['lr'] = lr + for param_group in self.optimizer_G.param_groups: + param_group['lr'] = lr + if self.opt.verbose: + print('update learning rate: %f -> %f' % (self.old_lr, lr)) + self.old_lr = lr + + diff --git a/models/projectionhead.py b/models/projectionhead.py new file mode 100644 index 0000000..a0d18f1 --- /dev/null +++ b/models/projectionhead.py @@ -0,0 +1,14 @@ +import torch.nn as nn + +class ProjectionHead(nn.Module): + def __init__(self, proj_dim=256): + super(ProjectionHead, self).__init__() + + self.proj = nn.Sequential( + nn.Linear(proj_dim, proj_dim), + nn.ReLU(), + nn.Linear(proj_dim, proj_dim), + ) + + def forward(self, x): + return self.proj(x) \ No newline at end of file diff --git a/options/test_options.py b/options/test_options.py index 140e1c4..5b43c9a 100644 --- a/options/test_options.py +++ b/options/test_options.py @@ -1,3 +1,11 @@ +''' +Author: Naiyuan liu +Github: https://github.com/NNNNAI +Date: 2021-11-23 17:03:58 +LastEditors: Naiyuan liu +LastEditTime: 2021-11-23 17:08:08 +Description: +''' from .base_options import BaseOptions class TestOptions(BaseOptions): @@ -24,6 +32,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.parser.add_argument('--crop_size', type=int, default=224, help='Crop of size of input image') - - self.isTrain = False + self.isTrain = False \ No newline at end of file diff --git a/parsing_model/model.py b/parsing_model/model.py new file mode 100644 index 0000000..d434836 --- /dev/null +++ b/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() diff --git a/parsing_model/resnet.py b/parsing_model/resnet.py new file mode 100644 index 0000000..aa2bf95 --- /dev/null +++ b/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() diff --git a/pg_modules/blocks.py b/pg_modules/blocks.py new file mode 100644 index 0000000..78bd113 --- /dev/null +++ b/pg_modules/blocks.py @@ -0,0 +1,325 @@ +import functools +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.nn.utils import spectral_norm + + +### single layers + + +def conv2d(*args, **kwargs): + return spectral_norm(nn.Conv2d(*args, **kwargs)) + + +def convTranspose2d(*args, **kwargs): + return spectral_norm(nn.ConvTranspose2d(*args, **kwargs)) + + +def embedding(*args, **kwargs): + return spectral_norm(nn.Embedding(*args, **kwargs)) + + +def linear(*args, **kwargs): + return spectral_norm(nn.Linear(*args, **kwargs)) + + +def NormLayer(c, mode='batch'): + if mode == 'group': + return nn.GroupNorm(c//2, c) + elif mode == 'batch': + return nn.BatchNorm2d(c) + + +### Activations + + +class GLU(nn.Module): + def forward(self, x): + nc = x.size(1) + assert nc % 2 == 0, 'channels dont divide 2!' + nc = int(nc/2) + return x[:, :nc] * torch.sigmoid(x[:, nc:]) + + +class Swish(nn.Module): + def forward(self, feat): + return feat * torch.sigmoid(feat) + + +### Upblocks + + +class InitLayer(nn.Module): + def __init__(self, nz, channel, sz=4): + super().__init__() + + self.init = nn.Sequential( + convTranspose2d(nz, channel*2, sz, 1, 0, bias=False), + NormLayer(channel*2), + GLU(), + ) + + def forward(self, noise): + noise = noise.view(noise.shape[0], -1, 1, 1) + return self.init(noise) + + +def UpBlockSmall(in_planes, out_planes): + block = nn.Sequential( + nn.Upsample(scale_factor=2, mode='nearest'), + conv2d(in_planes, out_planes*2, 3, 1, 1, bias=False), + NormLayer(out_planes*2), GLU()) + return block + + +class UpBlockSmallCond(nn.Module): + def __init__(self, in_planes, out_planes, z_dim): + super().__init__() + self.in_planes = in_planes + self.out_planes = out_planes + self.up = nn.Upsample(scale_factor=2, mode='nearest') + self.conv = conv2d(in_planes, out_planes*2, 3, 1, 1, bias=False) + + which_bn = functools.partial(CCBN, which_linear=linear, input_size=z_dim) + self.bn = which_bn(2*out_planes) + self.act = GLU() + + def forward(self, x, c): + x = self.up(x) + x = self.conv(x) + x = self.bn(x, c) + x = self.act(x) + return x + + +def UpBlockBig(in_planes, out_planes): + block = nn.Sequential( + nn.Upsample(scale_factor=2, mode='nearest'), + conv2d(in_planes, out_planes*2, 3, 1, 1, bias=False), + NoiseInjection(), + NormLayer(out_planes*2), GLU(), + conv2d(out_planes, out_planes*2, 3, 1, 1, bias=False), + NoiseInjection(), + NormLayer(out_planes*2), GLU() + ) + return block + + +class UpBlockBigCond(nn.Module): + def __init__(self, in_planes, out_planes, z_dim): + super().__init__() + self.in_planes = in_planes + self.out_planes = out_planes + self.up = nn.Upsample(scale_factor=2, mode='nearest') + self.conv1 = conv2d(in_planes, out_planes*2, 3, 1, 1, bias=False) + self.conv2 = conv2d(out_planes, out_planes*2, 3, 1, 1, bias=False) + + which_bn = functools.partial(CCBN, which_linear=linear, input_size=z_dim) + self.bn1 = which_bn(2*out_planes) + self.bn2 = which_bn(2*out_planes) + self.act = GLU() + self.noise = NoiseInjection() + + def forward(self, x, c): + # block 1 + x = self.up(x) + x = self.conv1(x) + x = self.noise(x) + x = self.bn1(x, c) + x = self.act(x) + + # block 2 + x = self.conv2(x) + x = self.noise(x) + x = self.bn2(x, c) + x = self.act(x) + + return x + + +class SEBlock(nn.Module): + def __init__(self, ch_in, ch_out): + super().__init__() + self.main = nn.Sequential( + nn.AdaptiveAvgPool2d(4), + conv2d(ch_in, ch_out, 4, 1, 0, bias=False), + Swish(), + conv2d(ch_out, ch_out, 1, 1, 0, bias=False), + nn.Sigmoid(), + ) + + def forward(self, feat_small, feat_big): + return feat_big * self.main(feat_small) + + +### Downblocks + + +class SeparableConv2d(nn.Module): + def __init__(self, in_channels, out_channels, kernel_size, bias=False): + super(SeparableConv2d, self).__init__() + self.depthwise = conv2d(in_channels, in_channels, kernel_size=kernel_size, + groups=in_channels, bias=bias, padding=1) + self.pointwise = conv2d(in_channels, out_channels, + kernel_size=1, bias=bias) + + def forward(self, x): + out = self.depthwise(x) + out = self.pointwise(out) + return out + + +class DownBlock(nn.Module): + def __init__(self, in_planes, out_planes, separable=False): + super().__init__() + if not separable: + self.main = nn.Sequential( + conv2d(in_planes, out_planes, 4, 2, 1), + NormLayer(out_planes), + nn.LeakyReLU(0.2, inplace=True), + ) + else: + self.main = nn.Sequential( + SeparableConv2d(in_planes, out_planes, 3), + NormLayer(out_planes), + nn.LeakyReLU(0.2, inplace=True), + nn.AvgPool2d(2, 2), + ) + + def forward(self, feat): + return self.main(feat) + + +class DownBlockPatch(nn.Module): + def __init__(self, in_planes, out_planes, separable=False): + super().__init__() + self.main = nn.Sequential( + DownBlock(in_planes, out_planes, separable), + conv2d(out_planes, out_planes, 1, 1, 0, bias=False), + NormLayer(out_planes), + nn.LeakyReLU(0.2, inplace=True), + ) + + def forward(self, feat): + return self.main(feat) + + +### CSM + + +class ResidualConvUnit(nn.Module): + def __init__(self, cin, activation, bn): + super().__init__() + self.conv = nn.Conv2d(cin, cin, kernel_size=3, stride=1, padding=1, bias=True) + self.skip_add = nn.quantized.FloatFunctional() + + def forward(self, x): + return self.skip_add.add(self.conv(x), x) + + +class FeatureFusionBlock(nn.Module): + def __init__(self, features, activation, deconv=False, bn=False, expand=False, align_corners=True, lowest=False): + super().__init__() + + self.deconv = deconv + self.align_corners = align_corners + + self.expand = expand + out_features = features + if self.expand==True: + out_features = features//2 + + self.out_conv = nn.Conv2d(features, out_features, kernel_size=1, stride=1, padding=0, bias=True, groups=1) + self.skip_add = nn.quantized.FloatFunctional() + + def forward(self, *xs): + output = xs[0] + + if len(xs) == 2: + output = self.skip_add.add(output, xs[1]) + + output = nn.functional.interpolate( + output, scale_factor=2, mode="bilinear", align_corners=self.align_corners + ) + + output = self.out_conv(output) + + return output + + +### Misc + + +class NoiseInjection(nn.Module): + def __init__(self): + super().__init__() + self.weight = nn.Parameter(torch.zeros(1), requires_grad=True) + + def forward(self, feat, noise=None): + if noise is None: + batch, _, height, width = feat.shape + noise = torch.randn(batch, 1, height, width).to(feat.device) + + return feat + self.weight * noise + + +class CCBN(nn.Module): + ''' conditional batchnorm ''' + def __init__(self, output_size, input_size, which_linear, eps=1e-5, momentum=0.1): + super().__init__() + self.output_size, self.input_size = output_size, input_size + + # Prepare gain and bias layers + self.gain = which_linear(input_size, output_size) + self.bias = which_linear(input_size, output_size) + + # epsilon to avoid dividing by 0 + self.eps = eps + # Momentum + self.momentum = momentum + + self.register_buffer('stored_mean', torch.zeros(output_size)) + self.register_buffer('stored_var', torch.ones(output_size)) + + def forward(self, x, y): + # Calculate class-conditional gains and biases + gain = (1 + self.gain(y)).view(y.size(0), -1, 1, 1) + bias = self.bias(y).view(y.size(0), -1, 1, 1) + out = F.batch_norm(x, self.stored_mean, self.stored_var, None, None, + self.training, 0.1, self.eps) + return out * gain + bias + + +class Interpolate(nn.Module): + """Interpolation module.""" + + def __init__(self, size, mode='bilinear', align_corners=False): + """Init. + Args: + scale_factor (float): scaling + mode (str): interpolation mode + """ + super(Interpolate, self).__init__() + + self.interp = nn.functional.interpolate + self.size = size + self.mode = mode + self.align_corners = align_corners + + def forward(self, x): + """Forward pass. + Args: + x (tensor): input + Returns: + tensor: interpolated data + """ + + x = self.interp( + x, + size=self.size, + mode=self.mode, + align_corners=self.align_corners, + ) + + return x diff --git a/pg_modules/diffaug.py b/pg_modules/diffaug.py new file mode 100644 index 0000000..54020be --- /dev/null +++ b/pg_modules/diffaug.py @@ -0,0 +1,76 @@ +# Differentiable Augmentation for Data-Efficient GAN Training +# Shengyu Zhao, Zhijian Liu, Ji Lin, Jun-Yan Zhu, and Song Han +# https://arxiv.org/pdf/2006.10738 + +import torch +import torch.nn.functional as F + + +def DiffAugment(x, policy='', channels_first=True): + if policy: + if not channels_first: + x = x.permute(0, 3, 1, 2) + for p in policy.split(','): + for f in AUGMENT_FNS[p]: + x = f(x) + if not channels_first: + x = x.permute(0, 2, 3, 1) + x = x.contiguous() + return x + + +def rand_brightness(x): + x = x + (torch.rand(x.size(0), 1, 1, 1, dtype=x.dtype, device=x.device) - 0.5) + return x + + +def rand_saturation(x): + x_mean = x.mean(dim=1, keepdim=True) + x = (x - x_mean) * (torch.rand(x.size(0), 1, 1, 1, dtype=x.dtype, device=x.device) * 2) + x_mean + return x + + +def rand_contrast(x): + x_mean = x.mean(dim=[1, 2, 3], keepdim=True) + x = (x - x_mean) * (torch.rand(x.size(0), 1, 1, 1, dtype=x.dtype, device=x.device) + 0.5) + x_mean + return x + + +def rand_translation(x, ratio=0.125): + shift_x, shift_y = int(x.size(2) * ratio + 0.5), int(x.size(3) * ratio + 0.5) + translation_x = torch.randint(-shift_x, shift_x + 1, size=[x.size(0), 1, 1], device=x.device) + translation_y = torch.randint(-shift_y, shift_y + 1, size=[x.size(0), 1, 1], device=x.device) + grid_batch, grid_x, grid_y = torch.meshgrid( + torch.arange(x.size(0), dtype=torch.long, device=x.device), + torch.arange(x.size(2), dtype=torch.long, device=x.device), + torch.arange(x.size(3), dtype=torch.long, device=x.device), + ) + grid_x = torch.clamp(grid_x + translation_x + 1, 0, x.size(2) + 1) + grid_y = torch.clamp(grid_y + translation_y + 1, 0, x.size(3) + 1) + x_pad = F.pad(x, [1, 1, 1, 1, 0, 0, 0, 0]) + x = x_pad.permute(0, 2, 3, 1).contiguous()[grid_batch, grid_x, grid_y].permute(0, 3, 1, 2) + return x + + +def rand_cutout(x, ratio=0.2): + cutout_size = int(x.size(2) * ratio + 0.5), int(x.size(3) * ratio + 0.5) + offset_x = torch.randint(0, x.size(2) + (1 - cutout_size[0] % 2), size=[x.size(0), 1, 1], device=x.device) + offset_y = torch.randint(0, x.size(3) + (1 - cutout_size[1] % 2), size=[x.size(0), 1, 1], device=x.device) + grid_batch, grid_x, grid_y = torch.meshgrid( + torch.arange(x.size(0), dtype=torch.long, device=x.device), + torch.arange(cutout_size[0], dtype=torch.long, device=x.device), + torch.arange(cutout_size[1], dtype=torch.long, device=x.device), + ) + grid_x = torch.clamp(grid_x + offset_x - cutout_size[0] // 2, min=0, max=x.size(2) - 1) + grid_y = torch.clamp(grid_y + offset_y - cutout_size[1] // 2, min=0, max=x.size(3) - 1) + mask = torch.ones(x.size(0), x.size(2), x.size(3), dtype=x.dtype, device=x.device) + mask[grid_batch, grid_x, grid_y] = 0 + x = x * mask.unsqueeze(1) + return x + + +AUGMENT_FNS = { + 'color': [rand_brightness, rand_saturation, rand_contrast], + 'translation': [rand_translation], + 'cutout': [rand_cutout], +} diff --git a/pg_modules/projected_discriminator.py b/pg_modules/projected_discriminator.py new file mode 100644 index 0000000..d0c879f --- /dev/null +++ b/pg_modules/projected_discriminator.py @@ -0,0 +1,191 @@ +from functools import partial +import numpy as np +import torch +import torch.nn as nn + +from pg_modules.blocks import DownBlock, DownBlockPatch, conv2d +from pg_modules.projector import F_RandomProj +from pg_modules.diffaug import DiffAugment + + +class SingleDisc(nn.Module): + def __init__(self, nc=None, ndf=None, start_sz=256, end_sz=8, head=None, separable=False, patch=False): + super().__init__() + channel_dict = {4: 512, 8: 512, 16: 256, 32: 128, 64: 64, 128: 64, + 256: 32, 512: 16, 1024: 8} + + # interpolate for start sz that are not powers of two + if start_sz not in channel_dict.keys(): + sizes = np.array(list(channel_dict.keys())) + start_sz = sizes[np.argmin(abs(sizes - start_sz))] + self.start_sz = start_sz + + # if given ndf, allocate all layers with the same ndf + if ndf is None: + nfc = channel_dict + else: + nfc = {k: ndf for k, v in channel_dict.items()} + + # for feature map discriminators with nfc not in channel_dict + # this is the case for the pretrained backbone (midas.pretrained) + if nc is not None and head is None: + nfc[start_sz] = nc + + layers = [] + + # Head if the initial input is the full modality + if head: + layers += [conv2d(nc, nfc[256], 3, 1, 1, bias=False), + nn.LeakyReLU(0.2, inplace=True)] + + # Down Blocks + DB = partial(DownBlockPatch, separable=separable) if patch else partial(DownBlock, separable=separable) + while start_sz > end_sz: + layers.append(DB(nfc[start_sz], nfc[start_sz//2])) + start_sz = start_sz // 2 + + layers.append(conv2d(nfc[end_sz], 1, 4, 1, 0, bias=False)) + self.main = nn.Sequential(*layers) + + def forward(self, x, c): + return self.main(x) + + +class SingleDiscCond(nn.Module): + def __init__(self, nc=None, ndf=None, start_sz=256, end_sz=8, head=None, separable=False, patch=False, c_dim=1000, cmap_dim=64, embedding_dim=128): + super().__init__() + self.cmap_dim = cmap_dim + + # midas channels + channel_dict = {4: 512, 8: 512, 16: 256, 32: 128, 64: 64, 128: 64, + 256: 32, 512: 16, 1024: 8} + + # interpolate for start sz that are not powers of two + if start_sz not in channel_dict.keys(): + sizes = np.array(list(channel_dict.keys())) + start_sz = sizes[np.argmin(abs(sizes - start_sz))] + self.start_sz = start_sz + + # if given ndf, allocate all layers with the same ndf + if ndf is None: + nfc = channel_dict + else: + nfc = {k: ndf for k, v in channel_dict.items()} + + # for feature map discriminators with nfc not in channel_dict + # this is the case for the pretrained backbone (midas.pretrained) + if nc is not None and head is None: + nfc[start_sz] = nc + + layers = [] + + # Head if the initial input is the full modality + if head: + layers += [conv2d(nc, nfc[256], 3, 1, 1, bias=False), + nn.LeakyReLU(0.2, inplace=True)] + + # Down Blocks + DB = partial(DownBlockPatch, separable=separable) if patch else partial(DownBlock, separable=separable) + while start_sz > end_sz: + layers.append(DB(nfc[start_sz], nfc[start_sz//2])) + start_sz = start_sz // 2 + self.main = nn.Sequential(*layers) + + # additions for conditioning on class information + self.cls = conv2d(nfc[end_sz], self.cmap_dim, 4, 1, 0, bias=False) + self.embed = nn.Embedding(num_embeddings=c_dim, embedding_dim=embedding_dim) + self.embed_proj = nn.Sequential( + nn.Linear(self.embed.embedding_dim, self.cmap_dim), + nn.LeakyReLU(0.2, inplace=True), + ) + + def forward(self, x, c): + h = self.main(x) + out = self.cls(h) + + # conditioning via projection + cmap = self.embed_proj(self.embed(c.argmax(1))).unsqueeze(-1).unsqueeze(-1) + out = (out * cmap).sum(dim=1, keepdim=True) * (1 / np.sqrt(self.cmap_dim)) + + return out + + +class MultiScaleD(nn.Module): + def __init__( + self, + channels, + resolutions, + num_discs=4, + proj_type=2, # 0 = no projection, 1 = cross channel mixing, 2 = cross scale mixing + cond=0, + separable=False, + patch=False, + **kwargs, + ): + super().__init__() + + assert num_discs in [1, 2, 3, 4] + + # the first disc is on the lowest level of the backbone + self.disc_in_channels = channels[:num_discs] + self.disc_in_res = resolutions[:num_discs] + Disc = SingleDiscCond if cond else SingleDisc + + mini_discs = [] + for i, (cin, res) in enumerate(zip(self.disc_in_channels, self.disc_in_res)): + start_sz = res if not patch else 16 + mini_discs += [str(i), Disc(nc=cin, start_sz=start_sz, end_sz=8, separable=separable, patch=patch)], + self.mini_discs = nn.ModuleDict(mini_discs) + + def forward(self, features, c): + all_logits = [] + for k, disc in self.mini_discs.items(): + res = disc(features[k], c).view(features[k].size(0), -1) + all_logits.append(res) + + all_logits = torch.cat(all_logits, dim=1) + return all_logits + + +class ProjectedDiscriminator(torch.nn.Module): + def __init__( + self, + diffaug=True, + interp224=True, + backbone_kwargs={}, + **kwargs + ): + super().__init__() + self.diffaug = diffaug + self.interp224 = interp224 + self.feature_network = F_RandomProj(**backbone_kwargs) + self.discriminator = MultiScaleD( + channels=self.feature_network.CHANNELS, + resolutions=self.feature_network.RESOLUTIONS, + **backbone_kwargs, + ) + + def train(self, mode=True): + self.feature_network = self.feature_network.train(False) + self.discriminator = self.discriminator.train(mode) + return self + + def eval(self): + return self.train(False) + + def get_feature(self, x): + features = self.feature_network(x, get_features=True) + return features + + def forward(self, x, c): + # if self.diffaug: + # x = DiffAugment(x, policy='color,translation,cutout') + + # if self.interp224: + # x = F.interpolate(x, 224, mode='bilinear', align_corners=False) + + features,backbone_features = self.feature_network(x) + logits = self.discriminator(features, c) + + return logits,backbone_features + diff --git a/pg_modules/projector.py b/pg_modules/projector.py new file mode 100644 index 0000000..610a482 --- /dev/null +++ b/pg_modules/projector.py @@ -0,0 +1,158 @@ +import torch +import torch.nn as nn +import timm +from pg_modules.blocks import FeatureFusionBlock + + +def _make_scratch_ccm(scratch, in_channels, cout, expand=False): + # shapes + out_channels = [cout, cout*2, cout*4, cout*8] if expand else [cout]*4 + + scratch.layer0_ccm = nn.Conv2d(in_channels[0], out_channels[0], kernel_size=1, stride=1, padding=0, bias=True) + scratch.layer1_ccm = nn.Conv2d(in_channels[1], out_channels[1], kernel_size=1, stride=1, padding=0, bias=True) + scratch.layer2_ccm = nn.Conv2d(in_channels[2], out_channels[2], kernel_size=1, stride=1, padding=0, bias=True) + scratch.layer3_ccm = nn.Conv2d(in_channels[3], out_channels[3], kernel_size=1, stride=1, padding=0, bias=True) + + scratch.CHANNELS = out_channels + + return scratch + + +def _make_scratch_csm(scratch, in_channels, cout, expand): + scratch.layer3_csm = FeatureFusionBlock(in_channels[3], nn.ReLU(False), expand=expand, lowest=True) + scratch.layer2_csm = FeatureFusionBlock(in_channels[2], nn.ReLU(False), expand=expand) + scratch.layer1_csm = FeatureFusionBlock(in_channels[1], nn.ReLU(False), expand=expand) + scratch.layer0_csm = FeatureFusionBlock(in_channels[0], nn.ReLU(False)) + + # last refinenet does not expand to save channels in higher dimensions + scratch.CHANNELS = [cout, cout, cout*2, cout*4] if expand else [cout]*4 + + return scratch + + +def _make_efficientnet(model): + pretrained = nn.Module() + pretrained.layer0 = nn.Sequential(model.conv_stem, model.bn1, model.act1, *model.blocks[0:2]) + pretrained.layer1 = nn.Sequential(*model.blocks[2:3]) + pretrained.layer2 = nn.Sequential(*model.blocks[3:5]) + pretrained.layer3 = nn.Sequential(*model.blocks[5:9]) + return pretrained + + +def calc_channels(pretrained, inp_res=224): + channels = [] + tmp = torch.zeros(1, 3, inp_res, inp_res) + + # forward pass + tmp = pretrained.layer0(tmp) + channels.append(tmp.shape[1]) + tmp = pretrained.layer1(tmp) + channels.append(tmp.shape[1]) + tmp = pretrained.layer2(tmp) + channels.append(tmp.shape[1]) + tmp = pretrained.layer3(tmp) + channels.append(tmp.shape[1]) + + return channels + + +def _make_projector(im_res, cout, proj_type, expand=False): + assert proj_type in [0, 1, 2], "Invalid projection type" + + ### Build pretrained feature network + model = timm.create_model('tf_efficientnet_lite0', pretrained=True) + pretrained = _make_efficientnet(model) + + # determine resolution of feature maps, this is later used to calculate the number + # of down blocks in the discriminators. Interestingly, the best results are achieved + # by fixing this to 256, ie., we use the same number of down blocks per discriminator + # independent of the dataset resolution + im_res = 256 + pretrained.RESOLUTIONS = [im_res//4, im_res//8, im_res//16, im_res//32] + pretrained.CHANNELS = calc_channels(pretrained) + + if proj_type == 0: return pretrained, None + + ### Build CCM + scratch = nn.Module() + scratch = _make_scratch_ccm(scratch, in_channels=pretrained.CHANNELS, cout=cout, expand=expand) + pretrained.CHANNELS = scratch.CHANNELS + + if proj_type == 1: return pretrained, scratch + + ### build CSM + scratch = _make_scratch_csm(scratch, in_channels=scratch.CHANNELS, cout=cout, expand=expand) + + # CSM upsamples x2 so the feature map resolution doubles + pretrained.RESOLUTIONS = [res*2 for res in pretrained.RESOLUTIONS] + pretrained.CHANNELS = scratch.CHANNELS + + return pretrained, scratch + + +class F_RandomProj(nn.Module): + def __init__( + self, + im_res=256, + cout=64, + expand=True, + proj_type=2, # 0 = no projection, 1 = cross channel mixing, 2 = cross scale mixing + **kwargs, + ): + super().__init__() + self.proj_type = proj_type + self.cout = cout + self.expand = expand + + # build pretrained feature network and random decoder (scratch) + self.pretrained, self.scratch = _make_projector(im_res=im_res, cout=self.cout, proj_type=self.proj_type, expand=self.expand) + self.CHANNELS = self.pretrained.CHANNELS + self.RESOLUTIONS = self.pretrained.RESOLUTIONS + + def forward(self, x, get_features=False): + # predict feature maps + out0 = self.pretrained.layer0(x) + out1 = self.pretrained.layer1(out0) + out2 = self.pretrained.layer2(out1) + out3 = self.pretrained.layer3(out2) + + # start enumerating at the lowest layer (this is where we put the first discriminator) + backbone_features = { + '0': out0, + '1': out1, + '2': out2, + '3': out3, + } + if get_features: + return backbone_features + + if self.proj_type == 0: return backbone_features + + out0_channel_mixed = self.scratch.layer0_ccm(backbone_features['0']) + out1_channel_mixed = self.scratch.layer1_ccm(backbone_features['1']) + out2_channel_mixed = self.scratch.layer2_ccm(backbone_features['2']) + out3_channel_mixed = self.scratch.layer3_ccm(backbone_features['3']) + + out = { + '0': out0_channel_mixed, + '1': out1_channel_mixed, + '2': out2_channel_mixed, + '3': out3_channel_mixed, + } + + if self.proj_type == 1: return out + + # from bottom to top + out3_scale_mixed = self.scratch.layer3_csm(out3_channel_mixed) + out2_scale_mixed = self.scratch.layer2_csm(out3_scale_mixed, out2_channel_mixed) + out1_scale_mixed = self.scratch.layer1_csm(out2_scale_mixed, out1_channel_mixed) + out0_scale_mixed = self.scratch.layer0_csm(out1_scale_mixed, out0_channel_mixed) + + out = { + '0': out0_scale_mixed, + '1': out1_scale_mixed, + '2': out2_scale_mixed, + '3': out3_scale_mixed, + } + + return out, backbone_features diff --git a/predict.py b/predict.py new file mode 100644 index 0000000..3fe60e3 --- /dev/null +++ b/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 = opt.crop_size + 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, size=(112,112)) + 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 diff --git a/test_one_image.py b/test_one_image.py index 0915550..4eabd8e 100644 --- a/test_one_image.py +++ b/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, size=(112,112)) + 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) \ No newline at end of file + cv2.imwrite(opt.output_path + 'result.jpg',output) \ No newline at end of file diff --git a/test_video_swap_multispecific.py b/test_video_swap_multispecific.py index 5426ca8..fc9fd71 100644 --- a/test_video_swap_multispecific.py +++ b/test_video_swap_multispecific.py @@ -2,7 +2,6 @@ import cv2 import torch import fractions -import numpy as np from PIL import Image import torch.nn.functional as F from torchvision import transforms @@ -35,60 +34,66 @@ if __name__ == '__main__': opt = TestOptions().parse() pic_specific = opt.pic_specific_path start_epoch, epoch_iter = 1, 0 - crop_size = 224 + crop_size = opt.crop_size multisepcific_dir = opt.multisepcific_dir torch.nn.Module.dump_patches = True + if crop_size == 512: + opt.which_epoch = 550000 + opt.name = '512' + mode = 'ffhq' + else: + mode = 'None' model = create_model(opt) model.eval() app = Face_detect_crop(name='antelope', root='./insightface_func/models') - app.prepare(ctx_id= 0, det_thresh=0.6, det_size=(640,640)) + app.prepare(ctx_id= 0, det_thresh=0.6, det_size=(640,640),mode=mode) # 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)) - - 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, size=(112,112)) + 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, size=(112,112)) + 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,crop_size=crop_size) diff --git a/test_video_swapmulti.py b/test_video_swapmulti.py index 14a98ee..37d7751 100644 --- a/test_video_swapmulti.py +++ b/test_video_swapmulti.py @@ -1,3 +1,11 @@ +''' +Author: Naiyuan liu +Github: https://github.com/NNNNAI +Date: 2021-11-23 17:03:58 +LastEditors: Naiyuan liu +LastEditTime: 2021-11-24 19:00:34 +Description: +''' import cv2 import torch @@ -34,39 +42,47 @@ if __name__ == '__main__': opt = TestOptions().parse() start_epoch, epoch_iter = 1, 0 - crop_size = 224 + crop_size = opt.crop_size torch.nn.Module.dump_patches = True + + if crop_size == 512: + opt.which_epoch = 550000 + opt.name = '512' + mode = 'ffhq' + else: + mode = 'None' model = create_model(opt) model.eval() - app = Face_detect_crop(name='antelope', root='./insightface_func/models') - app.prepare(ctx_id= 0, det_thresh=0.6, det_size=(640,640)) + app.prepare(ctx_id= 0, det_thresh=0.6, det_size=(640,640),mode = mode) - 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, size=(112,112)) + 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,crop_size=crop_size) diff --git a/test_video_swapsingle.py b/test_video_swapsingle.py index b9cc525..221d0a3 100644 --- a/test_video_swapsingle.py +++ b/test_video_swapsingle.py @@ -1,3 +1,11 @@ +''' +Author: Naiyuan liu +Github: https://github.com/NNNNAI +Date: 2021-11-23 17:03:58 +LastEditors: Naiyuan liu +LastEditTime: 2021-11-24 19:00:38 +Description: +''' import cv2 import torch @@ -34,39 +42,46 @@ if __name__ == '__main__': opt = TestOptions().parse() start_epoch, epoch_iter = 1, 0 - crop_size = 224 + crop_size = opt.crop_size torch.nn.Module.dump_patches = True + if crop_size == 512: + opt.which_epoch = 550000 + opt.name = '512' + mode = 'ffhq' + else: + mode = 'None' model = create_model(opt) model.eval() app = Face_detect_crop(name='antelope', root='./insightface_func/models') - app.prepare(ctx_id= 0, det_thresh=0.6, det_size=(640,640)) + app.prepare(ctx_id= 0, det_thresh=0.6, det_size=(640,640),mode=mode) + 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, size=(112,112)) + 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,crop_size=crop_size) diff --git a/test_video_swapspecific.py b/test_video_swapspecific.py index dbcc4e0..3af643d 100644 --- a/test_video_swapspecific.py +++ b/test_video_swapspecific.py @@ -1,3 +1,11 @@ +''' +Author: Naiyuan liu +Github: https://github.com/NNNNAI +Date: 2021-11-23 17:03:58 +LastEditors: Naiyuan liu +LastEditTime: 2021-11-24 19:00:42 +Description: +''' import cv2 import torch @@ -34,51 +42,57 @@ if __name__ == '__main__': opt = TestOptions().parse() pic_specific = opt.pic_specific_path start_epoch, epoch_iter = 1, 0 - crop_size = 224 + crop_size = opt.crop_size torch.nn.Module.dump_patches = True + if crop_size == 512: + opt.which_epoch = 550000 + opt.name = '512' + mode = 'ffhq' + else: + mode = 'None' model = create_model(opt) model.eval() app = Face_detect_crop(name='antelope', root='./insightface_func/models') - app.prepare(ctx_id= 0, det_thresh=0.6, det_size=(640,640)) + app.prepare(ctx_id= 0, det_thresh=0.6, det_size=(640,640),mode=mode) + 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, size=(112,112)) + 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, size=(112,112)) + 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,crop_size=crop_size) diff --git a/test_wholeimage_swap_multispecific.py b/test_wholeimage_swap_multispecific.py index a5c6622..dca31af 100644 --- a/test_wholeimage_swap_multispecific.py +++ b/test_wholeimage_swap_multispecific.py @@ -1,3 +1,11 @@ +''' +Author: Naiyuan liu +Github: https://github.com/NNNNAI +Date: 2021-11-23 17:03:58 +LastEditors: Naiyuan liu +LastEditTime: 2021-11-24 19:19:22 +Description: +''' import cv2 import torch @@ -15,6 +23,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 @@ -37,11 +46,19 @@ if __name__ == '__main__': opt = TestOptions().parse() start_epoch, epoch_iter = 1, 0 - crop_size = 224 + crop_size = opt.crop_size multisepcific_dir = opt.multisepcific_dir torch.nn.Module.dump_patches = True + + if crop_size == 512: + opt.which_epoch = 550000 + opt.name = '512' + mode = 'ffhq' + else: + mode = 'None' + logoclass = watermark_image('./simswaplogo/simswaplogo.png') model = create_model(opt) model.eval() @@ -51,95 +68,109 @@ 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)) + app.prepare(ctx_id= 0, det_thresh=0.6, det_size=(640,640),mode = mode) - # 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, size=(112,112)) + 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, size=(112,112)) + 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, size=(112,112)) + 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)) diff --git a/test_wholeimage_swapmulti.py b/test_wholeimage_swapmulti.py index 6cd2d53..1938057 100644 --- a/test_wholeimage_swapmulti.py +++ b/test_wholeimage_swapmulti.py @@ -1,3 +1,11 @@ +''' +Author: Naiyuan liu +Github: https://github.com/NNNNAI +Date: 2021-11-23 17:03:58 +LastEditors: Naiyuan liu +LastEditTime: 2021-11-24 19:19:26 +Description: +''' import cv2 import torch @@ -12,6 +20,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 @@ -29,52 +39,72 @@ if __name__ == '__main__': opt = TestOptions().parse() start_epoch, epoch_iter = 1, 0 - crop_size = 224 + crop_size = opt.crop_size torch.nn.Module.dump_patches = True + if crop_size == 512: + opt.which_epoch = 550000 + opt.name = '512' + mode = 'ffhq' + else: + mode = 'None' 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)) + app.prepare(ctx_id= 0, det_thresh=0.6, det_size=(640,640),mode=mode) - 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, size=(112,112)) + 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 ! ************') diff --git a/test_wholeimage_swapsingle.py b/test_wholeimage_swapsingle.py index d5c28e3..e35d843 100644 --- a/test_wholeimage_swapsingle.py +++ b/test_wholeimage_swapsingle.py @@ -1,3 +1,11 @@ +''' +Author: Naiyuan liu +Github: https://github.com/NNNNAI +Date: 2021-11-23 17:03:58 +LastEditors: Naiyuan liu +LastEditTime: 2021-11-24 19:19:43 +Description: +''' import cv2 import torch @@ -12,6 +20,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 @@ -28,52 +38,73 @@ if __name__ == '__main__': opt = TestOptions().parse() start_epoch, epoch_iter = 1, 0 - crop_size = 224 + crop_size = opt.crop_size torch.nn.Module.dump_patches = True + if crop_size == 512: + opt.which_epoch = 550000 + opt.name = '512' + mode = 'ffhq' + else: + mode = 'None' 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)) + app.prepare(ctx_id= 0, det_thresh=0.6, det_size=(640,640),mode=mode) - 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, size=(112,112)) + 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 ! ************') diff --git a/test_wholeimage_swapspecific.py b/test_wholeimage_swapspecific.py index 0c3f1b4..b9167f8 100644 --- a/test_wholeimage_swapspecific.py +++ b/test_wholeimage_swapspecific.py @@ -1,3 +1,11 @@ +''' +Author: Naiyuan liu +Github: https://github.com/NNNNAI +Date: 2021-11-23 17:03:58 +LastEditors: Naiyuan liu +LastEditTime: 2021-11-24 19:19:47 +Description: +''' import cv2 import torch @@ -14,6 +22,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 @@ -36,9 +45,15 @@ if __name__ == '__main__': opt = TestOptions().parse() start_epoch, epoch_iter = 1, 0 - crop_size = 224 + crop_size = opt.crop_size torch.nn.Module.dump_patches = True + if crop_size == 512: + opt.which_epoch = 550000 + opt.name = '512' + mode = 'ffhq' + else: + mode = 'None' logoclass = watermark_image('./simswaplogo/simswaplogo.png') model = create_model(opt) model.eval() @@ -48,7 +63,7 @@ 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)) + app.prepare(ctx_id= 0, det_thresh=0.6, det_size=(640,640),mode=mode) pic_a = opt.pic_a_path pic_specific = opt.pic_specific_path @@ -64,7 +79,7 @@ if __name__ == '__main__': img_id = img_id.cuda() #create latent id - img_id_downsample = F.interpolate(img_id, scale_factor=0.5) + img_id_downsample = F.interpolate(img_id, size=(112,112)) latend_id = model.netArc(img_id_downsample) latend_id = F.normalize(latend_id, p=2, dim=1) @@ -80,7 +95,7 @@ if __name__ == '__main__': specific_person = specific_person.cuda() #create latent id - specific_person_downsample = F.interpolate(specific_person, scale_factor=0.5) + specific_person_downsample = F.interpolate(specific_person, size=(112,112)) specific_person_id_nonorm = model.netArc(specific_person_downsample) # specific_person_id_norm = F.normalize(specific_person_id_nonorm, p=2, dim=1) @@ -100,7 +115,7 @@ if __name__ == '__main__': 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_tenor_arcnorm_downsample = F.interpolate(b_align_crop_tenor_arcnorm, size=(112,112)) b_align_crop_id_nonorm = model.netArc(b_align_crop_tenor_arcnorm_downsample) id_compare_values.append(mse(b_align_crop_id_nonorm,specific_person_id_nonorm).detach().cpu().numpy()) @@ -110,11 +125,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(' ') diff --git a/train.ipynb b/train.ipynb new file mode 100644 index 0000000..01f8d10 --- /dev/null +++ b/train.ipynb @@ -0,0 +1,191 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "id": "fC7QoKePuJWu" + }, + "source": [ + "#Training Demo\n", + "This is a simple example for training the SimSwap 224*224 with VGGFace2-224.\n", + "\n", + "Code path: https://github.com/neuralchen/SimSwap\n", + "If you like the SimSwap project, please star it!\n", + "Paper path: https://arxiv.org/pdf/2106.06340v1.pdf or https://dl.acm.org/doi/10.1145/3394171.3413630" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": { + "id": "J8WrNaQHuUGC" + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Fri Apr 22 12:19:42 2022 \n", + "+-----------------------------------------------------------------------------+\n", + "| NVIDIA-SMI 456.71 Driver Version: 456.71 CUDA Version: 11.1 |\n", + "|-------------------------------+----------------------+----------------------+\n", + "| GPU Name TCC/WDDM | Bus-Id Disp.A | Volatile Uncorr. ECC |\n", + "| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |\n", + "|===============================+======================+======================|\n", + "| 0 TITAN Xp WDDM | 00000000:01:00.0 On | N/A |\n", + "| 23% 36C P8 15W / 250W | 1135MiB / 12288MiB | 4% Default |\n", + "+-------------------------------+----------------------+----------------------+\n", + " \n", + "+-----------------------------------------------------------------------------+\n", + "| Processes: |\n", + "| GPU GI CI PID Type Process name GPU Memory |\n", + "| ID ID Usage |\n", + "|=============================================================================|\n", + "| 0 N/A N/A 1232 C+G Insufficient Permissions N/A |\n", + "| 0 N/A N/A 1240 C+G Insufficient Permissions N/A |\n", + "| 0 N/A N/A 1528 C+G ...y\\ShellExperienceHost.exe N/A |\n", + "| 0 N/A N/A 7296 C+G Insufficient Permissions N/A |\n", + "| 0 N/A N/A 8280 C+G C:\\Windows\\explorer.exe N/A |\n", + "| 0 N/A N/A 9532 C+G ...artMenuExperienceHost.exe N/A |\n", + "| 0 N/A N/A 9896 C+G ...5n1h2txyewy\\SearchApp.exe N/A |\n", + "| 0 N/A N/A 11040 C+G ...2txyewy\\TextInputHost.exe N/A |\n", + "| 0 N/A N/A 11424 C+G Insufficient Permissions N/A |\n", + "| 0 N/A N/A 13112 C+G ...icrosoft VS Code\\Code.exe N/A |\n", + "| 0 N/A N/A 18720 C+G ...-2.9.15\\GitHubDesktop.exe N/A |\n", + "| 0 N/A N/A 22996 C+G ...bbwe\\Microsoft.Photos.exe N/A |\n", + "| 0 N/A N/A 23512 C+G ...me\\Application\\chrome.exe N/A |\n", + "| 0 N/A N/A 25892 C+G Insufficient Permissions N/A |\n", + "+-----------------------------------------------------------------------------+\n" + ] + } + ], + "source": [ + "!nvidia-smi" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "Z6BtQIgWuoqt" + }, + "source": [ + "Installation\n", + "All file changes made by this notebook are temporary. You can try to mount your own google drive to store files if you want." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "wdQJ9d8N8Tnf" + }, + "source": [ + "#Get Scripts" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "9jZWwt97uvIe" + }, + "outputs": [], + "source": [ + "!git clone https://github.com/neuralchen/SimSwap\n", + "!cd SimSwap && git pull" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "ATLrrbso8Y-Y" + }, + "source": [ + "# Install Blocks" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "rwvbPhtOvZAL" + }, + "outputs": [], + "source": [ + "!pip install googledrivedownloader\n", + "!pip install timm\n", + "!wget -P SimSwap/arcface_model https://github.com/neuralchen/SimSwap/releases/download/1.0/arcface_checkpoint.tar" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "hleVtHIJ_QUK" + }, + "source": [ + "#Download the Training Dataset\n", + "We employ the cropped VGGFace2-224 dataset for this toy training demo.\n", + "\n", + "You can download the dataset from our google driver https://drive.google.com/file/d/19pWvdEHS-CEG6tW3PdxdtZ5QEymVjImc/view?usp=sharing\n", + "\n", + "***Please check the dataset in dir /content/TrainingData***\n", + "\n", + "***If dataset already exists in /content/TrainingData, please do not run blow scripts!***\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "h2tyjBl0Llxp" + }, + "outputs": [], + "source": [ + "!wget --load-cookies /tmp/cookies.txt \"https://docs.google.com/uc?export=download&confirm=$(wget --quiet --save-cookies /tmp/cookies.txt --keep-session-cookies --no-check-certificate 'https://docs.google.com/uc?export=download&id=19pWvdEHS-CEG6tW3PdxdtZ5QEymVjImc' -O- | sed -rn 's/.*confirm=([0-9A-Za-z_]+).*/\\1\\n/p')&id=19pWvdEHS-CEG6tW3PdxdtZ5QEymVjImc\" -O /content/TrainingData/vggface2_crop_arcfacealign_224.tar && rm -rf /tmp/cookies.txt\n", + "%cd /content/\n", + "!tar -xzvf /content/TrainingData/vggface2_crop_arcfacealign_224.tar\n", + "!rm /content/TrainingData/vggface2_crop_arcfacealign_224.tar" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "o5SNDWzA8LjJ" + }, + "source": [ + "#Trainig\n", + "Batch size must larger than 1!" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "XCxHa4oW507s" + }, + "outputs": [], + "source": [ + "%cd /content/SimSwap\n", + "!ls\n", + "!python train.py --name simswap224_test --gpu_ids 0 --dataset /content/TrainingData/vggface2_crop_arcfacealign_224 --Gdeep False" + ] + } + ], + "metadata": { + "accelerator": "GPU", + "colab": { + "collapsed_sections": [], + "name": "train.ipynb", + "provenance": [] + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.8.5" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} diff --git a/train.py b/train.py new file mode 100644 index 0000000..ebe5fde --- /dev/null +++ b/train.py @@ -0,0 +1,290 @@ +#!/usr/bin/env python3 +# -*- coding:utf-8 -*- +############################################################# +# File: train.py +# Created Date: Monday December 27th 2021 +# Author: Chen Xuanhong +# Email: chenxuanhongzju@outlook.com +# Last Modified: Friday, 22nd April 2022 10:49:26 am +# Modified By: Chen Xuanhong +# Copyright (c) 2021 Shanghai Jiao Tong University +############################################################# + +import os +import time +import random +import argparse +import numpy as np + +import torch +import torch.nn.functional as F +from torch.backends import cudnn +import torch.utils.tensorboard as tensorboard + +from util import util +from util.plot import plot_batch + +from models.projected_model import fsModel +from data.data_loader_Swapping import GetLoader + +def str2bool(v): + return v.lower() in ('true') + +class TrainOptions: + def __init__(self): + self.parser = argparse.ArgumentParser() + self.initialized = False + + def initialize(self): + self.parser.add_argument('--name', type=str, default='simswap', help='name of the experiment. It decides where to store samples and models') + self.parser.add_argument('--gpu_ids', default='0') + self.parser.add_argument('--checkpoints_dir', type=str, default='./checkpoints', help='models are saved here') + self.parser.add_argument('--isTrain', type=str2bool, default='True') + + # input/output sizes + self.parser.add_argument('--batchSize', type=int, default=4, help='input batch size') + + # for displays + self.parser.add_argument('--use_tensorboard', type=str2bool, default='False') + + # for training + self.parser.add_argument('--dataset', type=str, default="/path/to/VGGFace2", help='path to the face swapping dataset') + self.parser.add_argument('--continue_train', type=str2bool, default='False', help='continue training: load the latest model') + self.parser.add_argument('--load_pretrain', type=str, default='./checkpoints/simswap224_test', help='load the pretrained model from the specified location') + self.parser.add_argument('--which_epoch', type=str, default='10000', help='which epoch to load? set to latest to use latest cached model') + self.parser.add_argument('--phase', type=str, default='train', help='train, val, test, etc') + self.parser.add_argument('--niter', type=int, default=10000, help='# of iter at starting learning rate') + self.parser.add_argument('--niter_decay', type=int, default=10000, help='# of iter to linearly decay learning rate to zero') + self.parser.add_argument('--beta1', type=float, default=0.0, help='momentum term of adam') + self.parser.add_argument('--lr', type=float, default=0.0004, help='initial learning rate for adam') + self.parser.add_argument('--Gdeep', type=str2bool, default='False') + + # for discriminators + self.parser.add_argument('--lambda_feat', type=float, default=10.0, help='weight for feature matching loss') + self.parser.add_argument('--lambda_id', type=float, default=30.0, help='weight for id loss') + self.parser.add_argument('--lambda_rec', type=float, default=10.0, help='weight for reconstruction loss') + + self.parser.add_argument("--Arc_path", type=str, default='arcface_model/arcface_checkpoint.tar', help="run ONNX model via TRT") + self.parser.add_argument("--total_step", type=int, default=1000000, help='total training step') + self.parser.add_argument("--log_frep", type=int, default=200, help='frequence for printing log information') + self.parser.add_argument("--sample_freq", type=int, default=1000, help='frequence for sampling') + self.parser.add_argument("--model_freq", type=int, default=10000, help='frequence for saving the model') + + + + + self.isTrain = True + + def parse(self, save=True): + if not self.initialized: + self.initialize() + self.opt = self.parser.parse_args() + self.opt.isTrain = self.isTrain # train or test + + args = vars(self.opt) + + print('------------ Options -------------') + for k, v in sorted(args.items()): + print('%s: %s' % (str(k), str(v))) + print('-------------- End ----------------') + + # save to the disk + if self.opt.isTrain: + expr_dir = os.path.join(self.opt.checkpoints_dir, self.opt.name) + util.mkdirs(expr_dir) + if save and not self.opt.continue_train: + file_name = os.path.join(expr_dir, 'opt.txt') + with open(file_name, 'wt') as opt_file: + opt_file.write('------------ Options -------------\n') + for k, v in sorted(args.items()): + opt_file.write('%s: %s\n' % (str(k), str(v))) + opt_file.write('-------------- End ----------------\n') + return self.opt + + +if __name__ == '__main__': + + opt = TrainOptions().parse() + iter_path = os.path.join(opt.checkpoints_dir, opt.name, 'iter.txt') + + sample_path = os.path.join(opt.checkpoints_dir, opt.name, 'samples') + + if not os.path.exists(sample_path): + os.makedirs(sample_path) + + log_path = os.path.join(opt.checkpoints_dir, opt.name, 'summary') + + if not os.path.exists(log_path): + os.makedirs(log_path) + + if opt.continue_train: + try: + start_epoch, epoch_iter = np.loadtxt(iter_path , delimiter=',', dtype=int) + except: + start_epoch, epoch_iter = 1, 0 + print('Resuming from epoch %d at iteration %d' % (start_epoch, epoch_iter)) + else: + start_epoch, epoch_iter = 1, 0 + + os.environ['CUDA_VISIBLE_DEVICES'] = str(opt.gpu_ids) + print("GPU used : ", str(opt.gpu_ids)) + + + cudnn.benchmark = True + + + + model = fsModel() + + model.initialize(opt) + + ##################################################### + if opt.use_tensorboard: + tensorboard_writer = tensorboard.SummaryWriter(log_path) + logger = tensorboard_writer + + log_name = os.path.join(opt.checkpoints_dir, opt.name, 'loss_log.txt') + + with open(log_name, "a") as log_file: + now = time.strftime("%c") + log_file.write('================ Training Loss (%s) ================\n' % now) + + optimizer_G, optimizer_D = model.optimizer_G, model.optimizer_D + + loss_avg = 0 + refresh_count = 0 + imagenet_std = torch.Tensor([0.229, 0.224, 0.225]).view(3,1,1) + imagenet_mean = torch.Tensor([0.485, 0.456, 0.406]).view(3,1,1) + + train_loader = GetLoader(opt.dataset,opt.batchSize,8,1234) + + randindex = [i for i in range(opt.batchSize)] + random.shuffle(randindex) + + if not opt.continue_train: + start = 0 + else: + start = int(opt.which_epoch) + total_step = opt.total_step + import datetime + print("Start to train at %s"%(datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S'))) + + from util.logo_class import logo_class + logo_class.print_start_training() + model.netD.feature_network.requires_grad_(False) + + # Training Cycle + for step in range(start, total_step): + model.netG.train() + for interval in range(2): + random.shuffle(randindex) + src_image1, src_image2 = train_loader.next() + + if step%2 == 0: + img_id = src_image2 + else: + img_id = src_image2[randindex] + + img_id_112 = F.interpolate(img_id,size=(112,112), mode='bicubic') + latent_id = model.netArc(img_id_112) + latent_id = F.normalize(latent_id, p=2, dim=1) + if interval: + + img_fake = model.netG(src_image1, latent_id) + gen_logits,_ = model.netD(img_fake.detach(), None) + loss_Dgen = (F.relu(torch.ones_like(gen_logits) + gen_logits)).mean() + + real_logits,_ = model.netD(src_image2,None) + loss_Dreal = (F.relu(torch.ones_like(real_logits) - real_logits)).mean() + + loss_D = loss_Dgen + loss_Dreal + optimizer_D.zero_grad() + loss_D.backward() + optimizer_D.step() + else: + + # model.netD.requires_grad_(True) + img_fake = model.netG(src_image1, latent_id) + # G loss + gen_logits,feat = model.netD(img_fake, None) + + loss_Gmain = (-gen_logits).mean() + img_fake_down = F.interpolate(img_fake, size=(112,112), mode='bicubic') + latent_fake = model.netArc(img_fake_down) + latent_fake = F.normalize(latent_fake, p=2, dim=1) + loss_G_ID = (1 - model.cosin_metric(latent_fake, latent_id)).mean() + real_feat = model.netD.get_feature(src_image1) + feat_match_loss = model.criterionFeat(feat["3"],real_feat["3"]) + loss_G = loss_Gmain + loss_G_ID * opt.lambda_id + feat_match_loss * opt.lambda_feat + + + if step%2 == 0: + #G_Rec + loss_G_Rec = model.criterionRec(img_fake, src_image1) * opt.lambda_rec + loss_G += loss_G_Rec + + optimizer_G.zero_grad() + loss_G.backward() + optimizer_G.step() + + + ############## Display results and errors ########## + ### print out errors + # Print out log info + if (step + 1) % opt.log_frep == 0: + # errors = {k: v.data.item() if not isinstance(v, int) else v for k, v in loss_dict.items()} + errors = { + "G_Loss":loss_Gmain.item(), + "G_ID":loss_G_ID.item(), + "G_Rec":loss_G_Rec.item(), + "G_feat_match":feat_match_loss.item(), + "D_fake":loss_Dgen.item(), + "D_real":loss_Dreal.item(), + "D_loss":loss_D.item() + } + if opt.use_tensorboard: + for tag, value in errors.items(): + logger.add_scalar(tag, value, step) + message = '( step: %d, ) ' % (step) + for k, v in errors.items(): + message += '%s: %.3f ' % (k, v) + + print(message) + with open(log_name, "a") as log_file: + log_file.write('%s\n' % message) + + ### display output images + if (step + 1) % opt.sample_freq == 0: + model.netG.eval() + with torch.no_grad(): + imgs = list() + zero_img = (torch.zeros_like(src_image1[0,...])) + imgs.append(zero_img.cpu().numpy()) + save_img = ((src_image1.cpu())* imagenet_std + imagenet_mean).numpy() + for r in range(opt.batchSize): + imgs.append(save_img[r,...]) + arcface_112 = F.interpolate(src_image2,size=(112,112), mode='bicubic') + id_vector_src1 = model.netArc(arcface_112) + id_vector_src1 = F.normalize(id_vector_src1, p=2, dim=1) + + for i in range(opt.batchSize): + + imgs.append(save_img[i,...]) + image_infer = src_image1[i, ...].repeat(opt.batchSize, 1, 1, 1) + img_fake = model.netG(image_infer, id_vector_src1).cpu() + + img_fake = img_fake * imagenet_std + img_fake = img_fake + imagenet_mean + img_fake = img_fake.numpy() + for j in range(opt.batchSize): + imgs.append(img_fake[j,...]) + print("Save test data") + imgs = np.stack(imgs, axis = 0).transpose(0,2,3,1) + plot_batch(imgs, os.path.join(sample_path, 'step_'+str(step+1)+'.jpg')) + + ### save latest model + if (step+1) % opt.model_freq==0: + print('saving the latest model (steps %d)' % (step+1)) + model.save(step+1) + np.savetxt(iter_path, (step+1, total_step), delimiter=',', fmt='%d') + wandb.finish() \ No newline at end of file diff --git a/util/json_config.py b/util/json_config.py new file mode 100644 index 0000000..c68fbff --- /dev/null +++ b/util/json_config.py @@ -0,0 +1,15 @@ +import json + + +def readConfig(path): + with open(path,'r') as cf: + nodelocaltionstr = cf.read() + nodelocaltioninf = json.loads(nodelocaltionstr) + if isinstance(nodelocaltioninf,str): + nodelocaltioninf = json.loads(nodelocaltioninf) + return nodelocaltioninf + +def writeConfig(path, info): + with open(path, 'w') as cf: + configjson = json.dumps(info, indent=4) + cf.writelines(configjson) \ No newline at end of file diff --git a/util/logo_class.py b/util/logo_class.py new file mode 100644 index 0000000..044dce3 --- /dev/null +++ b/util/logo_class.py @@ -0,0 +1,44 @@ +#!/usr/bin/env python3 +# -*- coding:utf-8 -*- +############################################################# +# File: logo_class.py +# Created Date: Tuesday June 29th 2021 +# Author: Chen Xuanhong +# Email: chenxuanhongzju@outlook.com +# Last Modified: Monday, 11th October 2021 12:39:55 am +# Modified By: Chen Xuanhong +# Copyright (c) 2021 Shanghai Jiao Tong University +############################################################# + +class logo_class: + + @staticmethod + def print_group_logo(): + logo_str = """ + +███╗ ██╗██████╗ ███████╗██╗ ██████╗ ███████╗ ██╗████████╗██╗ ██╗ +████╗ ██║██╔══██╗██╔════╝██║██╔════╝ ██╔════╝ ██║╚══██╔══╝██║ ██║ +██╔██╗ ██║██████╔╝███████╗██║██║ ███╗ ███████╗ ██║ ██║ ██║ ██║ +██║╚██╗██║██╔══██╗╚════██║██║██║ ██║ ╚════██║██ ██║ ██║ ██║ ██║ +██║ ╚████║██║ ██║███████║██║╚██████╔╝ ███████║╚█████╔╝ ██║ ╚██████╔╝ +╚═╝ ╚═══╝╚═╝ ╚═╝╚══════╝╚═╝ ╚═════╝ ╚══════╝ ╚════╝ ╚═╝ ╚═════╝ +Neural Rendering Special Interesting Group of SJTU + + """ + print(logo_str) + + @staticmethod + def print_start_training(): + logo_str = """ + _____ __ __ ______ _ _ + / ___/ / /_ ____ _ _____ / /_ /_ __/_____ ____ _ (_)____ (_)____ ____ _ + \__ \ / __// __ `// ___// __/ / / / ___// __ `// // __ \ / // __ \ / __ `/ + ___/ // /_ / /_/ // / / /_ / / / / / /_/ // // / / // // / / // /_/ / +/____/ \__/ \__,_//_/ \__/ /_/ /_/ \__,_//_//_/ /_//_//_/ /_/ \__, / + /____/ + """ + print(logo_str) + +if __name__=="__main__": + # logo_class.print_group_logo() + logo_class.print_start_training() \ No newline at end of file diff --git a/util/plot.py b/util/plot.py new file mode 100644 index 0000000..0da1c75 --- /dev/null +++ b/util/plot.py @@ -0,0 +1,37 @@ +import numpy as np +import math +import PIL + +def postprocess(x): + """[0,1] to uint8.""" + + x = np.clip(255 * x, 0, 255) + x = np.cast[np.uint8](x) + return x + +def tile(X, rows, cols): + """Tile images for display.""" + tiling = np.zeros((rows * X.shape[1], cols * X.shape[2], X.shape[3]), dtype = X.dtype) + for i in range(rows): + for j in range(cols): + idx = i * cols + j + if idx < X.shape[0]: + img = X[idx,...] + tiling[ + i*X.shape[1]:(i+1)*X.shape[1], + j*X.shape[2]:(j+1)*X.shape[2], + :] = img + return tiling + + +def plot_batch(X, out_path): + """Save batch of images tiled.""" + n_channels = X.shape[3] + if n_channels > 3: + X = X[:,:,:,np.random.choice(n_channels, size = 3)] + X = postprocess(X) + rc = math.sqrt(X.shape[0]) + rows = cols = math.ceil(rc) + canvas = tile(X, rows, cols) + canvas = np.squeeze(canvas) + PIL.Image.fromarray(canvas).save(out_path) \ No newline at end of file diff --git a/util/reverse2original.py b/util/reverse2original.py index 0a2f99c..83f1bdb 100644 --- a/util/reverse2original.py +++ b/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, (crop_size, crop_size)) + # 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) - - \ No newline at end of file diff --git a/util/save_heatmap.py b/util/save_heatmap.py new file mode 100644 index 0000000..71ce4c9 --- /dev/null +++ b/util/save_heatmap.py @@ -0,0 +1,57 @@ +#!/usr/bin/env python3 +# -*- coding:utf-8 -*- +############################################################# +# File: save_heatmap.py +# Created Date: Friday January 15th 2021 +# Author: Chen Xuanhong +# Email: chenxuanhongzju@outlook.com +# Last Modified: Wednesday, 19th January 2022 1:22:47 am +# Modified By: Chen Xuanhong +# Copyright (c) 2021 Shanghai Jiao Tong University +############################################################# + +import os +import shutil +import seaborn as sns +import matplotlib.pyplot as plt +import cv2 +import numpy as np + +def SaveHeatmap(heatmaps, path, row=-1, dpi=72): + """ + The input tensor must be B X 1 X H X W + """ + batch_size = heatmaps.shape[0] + temp_path = ".temp/" + if not os.path.exists(temp_path): + os.makedirs(temp_path) + final_img = None + if row < 1: + col = batch_size + row = 1 + else: + col = batch_size // row + if row * col = col: + col_i = 0 + row_i += 1 + cv2.imwrite(path,final_img) + +if __name__ == "__main__": + random_map = np.random.randn(16,1,10,10) + SaveHeatmap(random_map,"./wocao.png",1) diff --git a/util/videoswap.py b/util/videoswap.py index 82cd814..4a6c8fe 100644 --- a/util/videoswap.py +++ b/util/videoswap.py @@ -1,3 +1,11 @@ +''' +Author: Naiyuan liu +Github: https://github.com/NNNNAI +Date: 2021-11-23 17:03:58 +LastEditors: Naiyuan liu +LastEditTime: 2021-11-24 19:19:52 +Description: +''' import os import cv2 import glob @@ -11,14 +19,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 +54,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 +78,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 @@ -67,11 +87,14 @@ def video_swap(video_path, id_vetor, swap_model, detect_model, save_path, temp_r frame_align_crop_tenor = _totensor(cv2.cvtColor(frame_align_crop,cv2.COLOR_BGR2RGB))[None,...].cuda() swap_result = swap_model(None, frame_align_crop_tenor, id_vetor, None, True)[0] + cv2.imwrite(os.path.join(temp_results_dir, 'frame_{:0>7d}.jpg'.format(frame_index)), frame) 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 +118,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') diff --git a/util/videoswap_multispecific.py b/util/videoswap_multispecific.py index f5364fe..40e9488 100644 --- a/util/videoswap_multispecific.py +++ b/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() @@ -72,7 +83,7 @@ def video_swap(video_path, target_id_norm_list,source_specific_id_nonorm_list,id frame_align_crop_tenor = _totensor(cv2.cvtColor(frame_align_crop,cv2.COLOR_BGR2RGB))[None,...].cuda() frame_align_crop_tenor_arcnorm = spNorm(frame_align_crop_tenor) - frame_align_crop_tenor_arcnorm_downsample = F.interpolate(frame_align_crop_tenor_arcnorm, scale_factor=0.5) + frame_align_crop_tenor_arcnorm_downsample = F.interpolate(frame_align_crop_tenor_arcnorm, size=(112,112)) frame_align_crop_crop_id_nonorm = swap_model.netArc(frame_align_crop_tenor_arcnorm_downsample) id_compare_values.append([]) for source_specific_id_nonorm_tmp in source_specific_id_nonorm_list: @@ -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') diff --git a/util/videoswap_specific.py b/util/videoswap_specific.py index 086eb32..8177efb 100644 --- a/util/videoswap_specific.py +++ b/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() @@ -72,7 +83,7 @@ def video_swap(video_path, id_vetor,specific_person_id_nonorm,id_thres, swap_mod frame_align_crop_tenor = _totensor(cv2.cvtColor(frame_align_crop,cv2.COLOR_BGR2RGB))[None,...].cuda() frame_align_crop_tenor_arcnorm = spNorm(frame_align_crop_tenor) - frame_align_crop_tenor_arcnorm_downsample = F.interpolate(frame_align_crop_tenor_arcnorm, scale_factor=0.5) + frame_align_crop_tenor_arcnorm_downsample = F.interpolate(frame_align_crop_tenor_arcnorm, size=(112,112)) frame_align_crop_crop_id_nonorm = swap_model.netArc(frame_align_crop_tenor_arcnorm_downsample) id_compare_values.append(mse(frame_align_crop_crop_id_nonorm,specific_person_id_nonorm).detach().cpu().numpy()) @@ -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')