[
{
"path": ".gitmodules",
"content": "[submodule \"Deep3DFaceRecon_pytorch\"]\n\tpath = Deep3DFaceRecon_pytorch\n\turl = https://github.com/sicxu/Deep3DFaceRecon_pytorch\n"
},
{
"path": "DatasetHelper.md",
"content": "### Extract 3DMM Coefficients for Videos\n\nWe provide scripts for extracting 3dmm coefficients for videos by using [DeepFaceRecon_pytorch](https://github.com/sicxu/Deep3DFaceRecon_pytorch/tree/73d491102af6731bded9ae6b3cc7466c3b2e9e48).\n\n1. Follow the instructions of their repo to build the environment of DeepFaceRecon.\n\n2. Copy the provided scrips to the folder `Deep3DFaceRecon_pytorch`.\n\n ```bash\n cp scripts/face_recon_videos.py ./Deep3DFaceRecon_pytorch\n cp scripts/extract_kp_videos.py ./Deep3DFaceRecon_pytorch\n cp scripts/coeff_detector.py ./Deep3DFaceRecon_pytorch\n cp scripts/inference_options.py ./Deep3DFaceRecon_pytorch/options\n\n cd Deep3DFaceRecon_pytorch\n ```\n\n3. Extract facial landmarks from videos.\n\n ```bash\n python extract_kp_videos.py \\\n --input_dir path_to_viodes \\\n --output_dir path_to_keypoint \\\n --device_ids 0,1,2,3 \\\n --workers 12\n ```\n\n4. Extract coefficients for videos\n\n ```bash\n python face_recon_videos.py \\\n --input_dir path_to_videos \\\n --keypoint_dir path_to_keypoint \\\n --output_dir output_dir \\\n --inference_batch_size 100 \\\n --name=model_name \\\n --epoch=20 \\\n --model facerecon\n ```\n\n \n\n\n\n"
},
{
"path": "LICENSE.md",
"content": "## creative commons\n\n# Attribution-NonCommercial 4.0 International\n\nCreative Commons Corporation (“Creative Commons”) is not a law firm and does not provide legal services or legal advice. Distribution of Creative Commons public licenses does not create a lawyer-client or other relationship. Creative Commons makes its licenses and related information available on an “as-is” basis. Creative Commons gives no warranties regarding its licenses, any material licensed under their terms and conditions, or any related information. Creative Commons disclaims all liability for damages resulting from their use to the fullest extent possible.\n\n### Using Creative Commons Public Licenses\n\nCreative Commons public licenses provide a standard set of terms and conditions that creators and other rights holders may use to share original works of authorship and other material subject to copyright and certain other rights specified in the public license below. The following considerations are for informational purposes only, are not exhaustive, and do not form part of our licenses.\n\n* __Considerations for licensors:__ Our public licenses are intended for use by those authorized to give the public permission to use material in ways otherwise restricted by copyright and certain other rights. Our licenses are irrevocable. Licensors should read and understand the terms and conditions of the license they choose before applying it. Licensors should also secure all rights necessary before applying our licenses so that the public can reuse the material as expected. Licensors should clearly mark any material not subject to the license. This includes other CC-licensed material, or material used under an exception or limitation to copyright. [More considerations for licensors](http://wiki.creativecommons.org/Considerations_for_licensors_and_licensees#Considerations_for_licensors).\n\n* __Considerations for the public:__ By using one of our public licenses, a licensor grants the public permission to use the licensed material under specified terms and conditions. If the licensor’s permission is not necessary for any reason–for example, because of any applicable exception or limitation to copyright–then that use is not regulated by the license. Our licenses grant only permissions under copyright and certain other rights that a licensor has authority to grant. Use of the licensed material may still be restricted for other reasons, including because others have copyright or other rights in the material. A licensor may make special requests, such as asking that all changes be marked or described. Although not required by our licenses, you are encouraged to respect those requests where reasonable. [More considerations for the public](http://wiki.creativecommons.org/Considerations_for_licensors_and_licensees#Considerations_for_licensees).\n\n## Creative Commons Attribution-NonCommercial 4.0 International Public License\n\nBy exercising the Licensed Rights (defined below), You accept and agree to be bound by the terms and conditions of this Creative Commons Attribution-NonCommercial 4.0 International Public License (\"Public License\"). To the extent this Public License may be interpreted as a contract, You are granted the Licensed Rights in consideration of Your acceptance of these terms and conditions, and the Licensor grants You such rights in consideration of benefits the Licensor receives from making the Licensed Material available under these terms and conditions.\n\n### Section 1 – Definitions.\n\na. __Adapted Material__ means material subject to Copyright and Similar Rights that is derived from or based upon the Licensed Material and in which the Licensed Material is translated, altered, arranged, transformed, or otherwise modified in a manner requiring permission under the Copyright and Similar Rights held by the Licensor. For purposes of this Public License, where the Licensed Material is a musical work, performance, or sound recording, Adapted Material is always produced where the Licensed Material is synched in timed relation with a moving image.\n\nb. __Adapter's License__ means the license You apply to Your Copyright and Similar Rights in Your contributions to Adapted Material in accordance with the terms and conditions of this Public License.\n\nc. __Copyright and Similar Rights__ means copyright and/or similar rights closely related to copyright including, without limitation, performance, broadcast, sound recording, and Sui Generis Database Rights, without regard to how the rights are labeled or categorized. For purposes of this Public License, the rights specified in Section 2(b)(1)-(2) are not Copyright and Similar Rights.\n\nd. __Effective Technological Measures__ means those measures that, in the absence of proper authority, may not be circumvented under laws fulfilling obligations under Article 11 of the WIPO Copyright Treaty adopted on December 20, 1996, and/or similar international agreements.\n\ne. __Exceptions and Limitations__ means fair use, fair dealing, and/or any other exception or limitation to Copyright and Similar Rights that applies to Your use of the Licensed Material.\n\nf. __Licensed Material__ means the artistic or literary work, database, or other material to which the Licensor applied this Public License.\n\ng. __Licensed Rights__ means the rights granted to You subject to the terms and conditions of this Public License, which are limited to all Copyright and Similar Rights that apply to Your use of the Licensed Material and that the Licensor has authority to license.\n\nh. __Licensor__ means the individual(s) or entity(ies) granting rights under this Public License.\n\ni. __NonCommercial__ means not primarily intended for or directed towards commercial advantage or monetary compensation. For purposes of this Public License, the exchange of the Licensed Material for other material subject to Copyright and Similar Rights by digital file-sharing or similar means is NonCommercial provided there is no payment of monetary compensation in connection with the exchange.\n\nj. __Share__ means to provide material to the public by any means or process that requires permission under the Licensed Rights, such as reproduction, public display, public performance, distribution, dissemination, communication, or importation, and to make material available to the public including in ways that members of the public may access the material from a place and at a time individually chosen by them.\n\nk. __Sui Generis Database Rights__ means rights other than copyright resulting from Directive 96/9/EC of the European Parliament and of the Council of 11 March 1996 on the legal protection of databases, as amended and/or succeeded, as well as other essentially equivalent rights anywhere in the world.\n\nl. __You__ means the individual or entity exercising the Licensed Rights under this Public License. Your has a corresponding meaning.\n\n### Section 2 – Scope.\n\na. ___License grant.___\n\n 1. Subject to the terms and conditions of this Public License, the Licensor hereby grants You a worldwide, royalty-free, non-sublicensable, non-exclusive, irrevocable license to exercise the Licensed Rights in the Licensed Material to:\n\n A. reproduce and Share the Licensed Material, in whole or in part, for NonCommercial purposes only; and\n\n B. produce, reproduce, and Share Adapted Material for NonCommercial purposes only.\n\n 2. __Exceptions and Limitations.__ For the avoidance of doubt, where Exceptions and Limitations apply to Your use, this Public License does not apply, and You do not need to comply with its terms and conditions.\n\n 3. __Term.__ The term of this Public License is specified in Section 6(a).\n\n 4. __Media and formats; technical modifications allowed.__ The Licensor authorizes You to exercise the Licensed Rights in all media and formats whether now known or hereafter created, and to make technical modifications necessary to do so. The Licensor waives and/or agrees not to assert any right or authority to forbid You from making technical modifications necessary to exercise the Licensed Rights, including technical modifications necessary to circumvent Effective Technological Measures. For purposes of this Public License, simply making modifications authorized by this Section 2(a)(4) never produces Adapted Material.\n\n 5. __Downstream recipients.__\n\n A. __Offer from the Licensor – Licensed Material.__ Every recipient of the Licensed Material automatically receives an offer from the Licensor to exercise the Licensed Rights under the terms and conditions of this Public License.\n\n B. __No downstream restrictions.__ You may not offer or impose any additional or different terms or conditions on, or apply any Effective Technological Measures to, the Licensed Material if doing so restricts exercise of the Licensed Rights by any recipient of the Licensed Material.\n\n 6. __No endorsement.__ Nothing in this Public License constitutes or may be construed as permission to assert or imply that You are, or that Your use of the Licensed Material is, connected with, or sponsored, endorsed, or granted official status by, the Licensor or others designated to receive attribution as provided in Section 3(a)(1)(A)(i).\n\nb. ___Other rights.___\n\n 1. Moral rights, such as the right of integrity, are not licensed under this Public License, nor are publicity, privacy, and/or other similar personality rights; however, to the extent possible, the Licensor waives and/or agrees not to assert any such rights held by the Licensor to the limited extent necessary to allow You to exercise the Licensed Rights, but not otherwise.\n\n 2. Patent and trademark rights are not licensed under this Public License.\n\n 3. To the extent possible, the Licensor waives any right to collect royalties from You for the exercise of the Licensed Rights, whether directly or through a collecting society under any voluntary or waivable statutory or compulsory licensing scheme. In all other cases the Licensor expressly reserves any right to collect such royalties, including when the Licensed Material is used other than for NonCommercial purposes.\n\n### Section 3 – License Conditions.\n\nYour exercise of the Licensed Rights is expressly made subject to the following conditions.\n\na. ___Attribution.___\n\n 1. If You Share the Licensed Material (including in modified form), You must:\n\n A. retain the following if it is supplied by the Licensor with the Licensed Material:\n\n i. identification of the creator(s) of the Licensed Material and any others designated to receive attribution, in any reasonable manner requested by the Licensor (including by pseudonym if designated);\n\n ii. a copyright notice;\n\n iii. a notice that refers to this Public License;\n\n iv. a notice that refers to the disclaimer of warranties;\n\n v. a URI or hyperlink to the Licensed Material to the extent reasonably practicable;\n\n B. indicate if You modified the Licensed Material and retain an indication of any previous modifications; and\n\n C. indicate the Licensed Material is licensed under this Public License, and include the text of, or the URI or hyperlink to, this Public License.\n\n 2. You may satisfy the conditions in Section 3(a)(1) in any reasonable manner based on the medium, means, and context in which You Share the Licensed Material. For example, it may be reasonable to satisfy the conditions by providing a URI or hyperlink to a resource that includes the required information.\n\n 3. If requested by the Licensor, You must remove any of the information required by Section 3(a)(1)(A) to the extent reasonably practicable.\n\n 4. If You Share Adapted Material You produce, the Adapter's License You apply must not prevent recipients of the Adapted Material from complying with this Public License.\n\n### Section 4 – Sui Generis Database Rights.\n\nWhere the Licensed Rights include Sui Generis Database Rights that apply to Your use of the Licensed Material:\n\na. for the avoidance of doubt, Section 2(a)(1) grants You the right to extract, reuse, reproduce, and Share all or a substantial portion of the contents of the database for NonCommercial purposes only;\n\nb. if You include all or a substantial portion of the database contents in a database in which You have Sui Generis Database Rights, then the database in which You have Sui Generis Database Rights (but not its individual contents) is Adapted Material; and\n\nc. You must comply with the conditions in Section 3(a) if You Share all or a substantial portion of the contents of the database.\n\nFor the avoidance of doubt, this Section 4 supplements and does not replace Your obligations under this Public License where the Licensed Rights include other Copyright and Similar Rights.\n\n### Section 5 – Disclaimer of Warranties and Limitation of Liability.\n\na. __Unless otherwise separately undertaken by the Licensor, to the extent possible, the Licensor offers the Licensed Material as-is and as-available, and makes no representations or warranties of any kind concerning the Licensed Material, whether express, implied, statutory, or other. This includes, without limitation, warranties of title, merchantability, fitness for a particular purpose, non-infringement, absence of latent or other defects, accuracy, or the presence or absence of errors, whether or not known or discoverable. Where disclaimers of warranties are not allowed in full or in part, this disclaimer may not apply to You.__\n\nb. __To the extent possible, in no event will the Licensor be liable to You on any legal theory (including, without limitation, negligence) or otherwise for any direct, special, indirect, incidental, consequential, punitive, exemplary, or other losses, costs, expenses, or damages arising out of this Public License or use of the Licensed Material, even if the Licensor has been advised of the possibility of such losses, costs, expenses, or damages. Where a limitation of liability is not allowed in full or in part, this limitation may not apply to You.__\n\nc. The disclaimer of warranties and limitation of liability provided above shall be interpreted in a manner that, to the extent possible, most closely approximates an absolute disclaimer and waiver of all liability.\n\n### Section 6 – Term and Termination.\n\na. This Public License applies for the term of the Copyright and Similar Rights licensed here. However, if You fail to comply with this Public License, then Your rights under this Public License terminate automatically.\n\nb. Where Your right to use the Licensed Material has terminated under Section 6(a), it reinstates:\n\n 1. automatically as of the date the violation is cured, provided it is cured within 30 days of Your discovery of the violation; or\n\n 2. upon express reinstatement by the Licensor.\n\n For the avoidance of doubt, this Section 6(b) does not affect any right the Licensor may have to seek remedies for Your violations of this Public License.\n\nc. For the avoidance of doubt, the Licensor may also offer the Licensed Material under separate terms or conditions or stop distributing the Licensed Material at any time; however, doing so will not terminate this Public License.\n\nd. Sections 1, 5, 6, 7, and 8 survive termination of this Public License.\n\n### Section 7 – Other Terms and Conditions.\n\na. The Licensor shall not be bound by any additional or different terms or conditions communicated by You unless expressly agreed.\n\nb. Any arrangements, understandings, or agreements regarding the Licensed Material not stated herein are separate from and independent of the terms and conditions of this Public License.\n\n### Section 8 – Interpretation.\n\na. For the avoidance of doubt, this Public License does not, and shall not be interpreted to, reduce, limit, restrict, or impose conditions on any use of the Licensed Material that could lawfully be made without permission under this Public License.\n\nb. To the extent possible, if any provision of this Public License is deemed unenforceable, it shall be automatically reformed to the minimum extent necessary to make it enforceable. If the provision cannot be reformed, it shall be severed from this Public License without affecting the enforceability of the remaining terms and conditions.\n\nc. No term or condition of this Public License will be waived and no failure to comply consented to unless expressly agreed to by the Licensor.\n\nd. Nothing in this Public License constitutes or may be interpreted as a limitation upon, or waiver of, any privileges and immunities that apply to the Licensor or You, including from the legal processes of any jurisdiction or authority.\n\n> Creative Commons is not a party to its public licenses. Notwithstanding, Creative Commons may elect to apply one of its public licenses to material it publishes and in those instances will be considered the “Licensor.” Except for the limited purpose of indicating that material is shared under a Creative Commons public license or as otherwise permitted by the Creative Commons policies published at [creativecommons.org/policies](http://creativecommons.org/policies), Creative Commons does not authorize the use of the trademark “Creative Commons” or any other trademark or logo of Creative Commons without its prior written consent including, without limitation, in connection with any unauthorized modifications to any of its public licenses or any other arrangements, understandings, or agreements concerning use of licensed material. For the avoidance of doubt, this paragraph does not form part of the public licenses.\n>\n> Creative Commons may be contacted at creativecommons.org\n"
},
{
"path": "README.md",
"content": "\n \n Website\n | \n ArXiv\n | \n Get Start\n | \n Video\n \n
\n\n\n# PIRenderer\n\nThe source code of the ICCV2021 paper \"[PIRenderer: Controllable Portrait Image Generation via Semantic Neural Rendering](https://arxiv.org/abs/2109.08379)\" (ICCV2021)\n\nThe proposed **PIRenderer** can synthesis portrait images by intuitively controlling the face motions with fully disentangled 3DMM parameters. This model can be applied to tasks such as:\n\n* **Intuitive Portrait Image Editing**\n\n \n 
\n
\n \n Intuitive Portrait Image Control \n
\n \n 
\n
\n \n Pose & Expression Alignment \n
\n \n \n* **Motion Imitation**\n \n 
\n
\n \n Same & Corss-identity Reenactment \n
\n \n* **Audio-Driven Facial Reenactment**\n\n \n 
\n
\n \n Audio-Driven Reenactment \n
\n\n## News\n\n* 2021.9.20 Code for PyTorch is available!\n\n\n\n## Colab Demo\n\nComing soon\n\n\n## Get Start\n\n### 1). Installation\n\n#### Requirements\n\n* Python 3\n* PyTorch 1.7.1\n* CUDA 10.2\n\n#### Conda Installation\n\n```bash\n# 1. Create a conda virtual environment.\nconda create -n PIRenderer python=3.6\nconda activate PIRenderer\nconda install -c pytorch pytorch=1.7.1 torchvision cudatoolkit=10.2\n\n# 2. Install other dependencies\npip install -r requirements.txt\n```\n\n### 2). Dataset\n\nWe train our model using the [VoxCeleb](https://arxiv.org/abs/1706.08612). You can download the demo dataset for inference or prepare the dataset for training and testing.\n\n#### Download the demo dataset\n\nThe demo dataset contains all 514 test videos. You can download the dataset with the following code:\n\n```bash\n./scripts/download_demo_dataset.sh\n```\n\nOr you can choose to download the resources with these links: \n\n\t[Google Driven](https://drive.google.com/drive/folders/16Yn2r46b4cV6ZozOH6a8SdFz_iG7BQk1?usp=sharing) & [BaiDu Driven](https://pan.baidu.com/s/1e615bBHvM4Wz-2snk-86Xw) with extraction passwords ”p9ab“\n\nThen unzip and save the files to `./dataset`\n\n#### Prepare the dataset\n\n1. The dataset is preprocessed follow the method used in [First-Order](https://github.com/AliaksandrSiarohin/video-preprocessing). You can follow the instructions in their repo to download and crop videos for training and testing.\n\n2. After obtaining the VoxCeleb videos, we extract 3DMM parameters using [Deep3DFaceReconstruction](https://github.com/microsoft/Deep3DFaceReconstruction). \n\n The folder are with format as:\n\n ```\n ${DATASET_ROOT_FOLDER}\n └───path_to_videos\n └───train\n └───xxx.mp4\n └───xxx.mp4\n ...\n └───test\n └───xxx.mp4\n └───xxx.mp4\n ...\n └───path_to_3dmm_coeff\n └───train\n └───xxx.mat\n └───xxx.mat\n ...\n └───test\n └───xxx.mat\n └───xxx.mat\n ...\n ```\n \n **News**: We provide Scripts for extracting 3dmm coeffs from videos. Please check the [DatasetHelper](./DatasetHelper.md) for more details.\n \n3. We save the video and 3DMM parameters in a lmdb file. Please run the following code to do this \n\n ```bash\n python scripts/prepare_vox_lmdb.py \\\n --path path_to_videos \\\n --coeff_3dmm_path path_to_3dmm_coeff \\\n --out path_to_output_dir\n ```\n\n### 3). Training and Inference\n\n#### Inference\n\nThe trained weights can be downloaded by running the following code:\n\n```bash\n./scripts/download_weights.sh\n```\n\nOr you can choose to download the resources with these links: \n\n[Google Driven](https://drive.google.com/file/d/1-0xOf6g58OmtKtEWJlU3VlnfRqPN9Uq7/view?usp=sharing) & [Baidu Driven](https://pan.baidu.com/s/18B3xfKMXnm4tOqlFSB8ntg) with extraction passwards \"4sy1\".\n\nThen unzip and save the files to `./result/face`.\n\n**Reenactment**\n\nRun the demo for face reenactment:\n\n```bash\n# same identity\npython -m torch.distributed.launch --nproc_per_node=1 --master_port 12345 inference.py \\\n--config ./config/face_demo.yaml \\\n--name face \\\n--no_resume \\\n--output_dir ./vox_result/face_reenactment\n\n# cross identity\npython -m torch.distributed.launch --nproc_per_node=1 --master_port 12345 inference.py \\\n--config ./config/face_demo.yaml \\\n--name face \\\n--no_resume \\\n--output_dir ./vox_result/face_reenactment_cross \\\n--cross_id\n```\n\nThe output results are saved at `./vox_result/face_reenactment` and `./vox_result/face_reenactment_cross`\n\n**Intuitive Control**\n\nOur model can generate results by providing intuitive controlling coefficients. \nWe provide the following code for this task. Please note that you need to build the environment of [DeepFaceRecon](https://github.com/sicxu/Deep3DFaceRecon_pytorch/tree/73d491102af6731bded9ae6b3cc7466c3b2e9e48) first.\n\n```bash\n# 1. Copy the provided scrips to the folder `Deep3DFaceRecon_pytorch`.\ncp scripts/face_recon_videos.py ./Deep3DFaceRecon_pytorch\ncp scripts/extract_kp_videos.py ./Deep3DFaceRecon_pytorch\ncp scripts/coeff_detector.py ./Deep3DFaceRecon_pytorch\ncp scripts/inference_options.py ./Deep3DFaceRecon_pytorch/options\n\ncd Deep3DFaceRecon_pytorch\n\n# 2. Extracte the 3dmm coefficients of the demo images.\npython coeff_detector.py \\\n--input_dir ../demo_images \\\n--keypoint_dir ../demo_images \\\n--output_dir ../demo_images \\\n--name=model_name \\\n--epoch=20 \\\n--model facerecon \n\n# 3. control the source image with our model\ncd ..\npython -m torch.distributed.launch --nproc_per_node=1 --master_port 12345 intuitive_control.py \\\n--config ./config/face_demo.yaml \\\n--name face \\\n--no_resume \\\n--output_dir ./vox_result/face_intuitive \\\n--input_name ./demo_images\n```\n\n\n#### Train\n\nOur model can be trained with the following code\n\n```bash\npython -m torch.distributed.launch --nproc_per_node=4 --master_port 12345 train.py \\\n--config ./config/face.yaml \\\n--name face\n```\n\n\n## Citation\n\nIf you find this code is helpful, please cite our paper\n\n```tex\n@misc{ren2021pirenderer,\n title={PIRenderer: Controllable Portrait Image Generation via Semantic Neural Rendering}, \n author={Yurui Ren and Ge Li and Yuanqi Chen and Thomas H. Li and Shan Liu},\n year={2021},\n eprint={2109.08379},\n archivePrefix={arXiv},\n primaryClass={cs.CV}\n}\n```\n\n## Acknowledgement \n\nWe build our project base on [imaginaire](https://github.com/NVlabs/imaginaire). Some dataset preprocessing methods are derived from [video-preprocessing](https://github.com/AliaksandrSiarohin/video-preprocessing).\n\n"
},
{
"path": "config/face.yaml",
"content": "# How often do you want to log the training stats.\n# network_list: \n# gen: gen_optimizer\n# dis: dis_optimizer\n\ndistributed: True\nimage_to_tensorboard: True\nsnapshot_save_iter: 40000\nsnapshot_save_epoch: 20\nsnapshot_save_start_iter: 20000\nsnapshot_save_start_epoch: 10\nimage_save_iter: 1000\nmax_epoch: 200\nlogging_iter: 100\nresults_dir: ./eval_results\n\ngen_optimizer:\n type: adam\n lr: 0.0001\n adam_beta1: 0.5\n adam_beta2: 0.999\n lr_policy:\n iteration_mode: True\n type: step\n step_size: 300000\n gamma: 0.2\n\ntrainer:\n type: trainers.face_trainer::FaceTrainer\n pretrain_warp_iteration: 200000\n loss_weight:\n weight_perceptual_warp: 2.5\n weight_perceptual_final: 4\n vgg_param_warp:\n network: vgg19\n layers: ['relu_1_1', 'relu_2_1', 'relu_3_1', 'relu_4_1', 'relu_5_1']\n use_style_loss: False\n num_scales: 4\n vgg_param_final:\n network: vgg19\n layers: ['relu_1_1', 'relu_2_1', 'relu_3_1', 'relu_4_1', 'relu_5_1']\n use_style_loss: True\n num_scales: 4 \n style_to_perceptual: 250\n init:\n type: 'normal'\n gain: 0.02\ngen:\n type: generators.face_model::FaceGenerator\n param:\n mapping_net:\n coeff_nc: 73\n descriptor_nc: 256\n layer: 3\n warpping_net:\n encoder_layer: 5\n decoder_layer: 3\n base_nc: 32\n editing_net:\n layer: 3\n num_res_blocks: 2\n base_nc: 64\n common:\n image_nc: 3\n descriptor_nc: 256\n max_nc: 256\n use_spect: False\n \n\n# Data options.\ndata:\n type: data.vox_dataset::VoxDataset\n path: ./dataset/vox_lmdb\n resolution: 256\n semantic_radius: 13\n train:\n batch_size: 5\n distributed: True\n val:\n batch_size: 8\n distributed: True\n\n\n"
},
{
"path": "config/face_demo.yaml",
"content": "# How often do you want to log the training stats.\n# network_list: \n# gen: gen_optimizer\n# dis: dis_optimizer\n\ndistributed: True\nimage_to_tensorboard: True\nsnapshot_save_iter: 40000\nsnapshot_save_epoch: 20\nsnapshot_save_start_iter: 20000\nsnapshot_save_start_epoch: 10\nimage_save_iter: 1000\nmax_epoch: 200\nlogging_iter: 100\nresults_dir: ./eval_results\n\ngen_optimizer:\n type: adam\n lr: 0.0001\n adam_beta1: 0.5\n adam_beta2: 0.999\n lr_policy:\n iteration_mode: True\n type: step\n step_size: 300000\n gamma: 0.2\n\ntrainer:\n type: trainers.face_trainer::FaceTrainer\n pretrain_warp_iteration: 200000\n loss_weight:\n weight_perceptual_warp: 2.5\n weight_perceptual_final: 4\n vgg_param_warp:\n network: vgg19\n layers: ['relu_1_1', 'relu_2_1', 'relu_3_1', 'relu_4_1', 'relu_5_1']\n use_style_loss: False\n num_scales: 4\n vgg_param_final:\n network: vgg19\n layers: ['relu_1_1', 'relu_2_1', 'relu_3_1', 'relu_4_1', 'relu_5_1']\n use_style_loss: True\n num_scales: 4 \n style_to_perceptual: 250\n init:\n type: 'normal'\n gain: 0.02\ngen:\n type: generators.face_model::FaceGenerator\n param:\n mapping_net:\n coeff_nc: 73\n descriptor_nc: 256\n layer: 3\n warpping_net:\n encoder_layer: 5\n decoder_layer: 3\n base_nc: 32\n editing_net:\n layer: 3\n num_res_blocks: 2\n base_nc: 64\n common:\n image_nc: 3\n descriptor_nc: 256\n max_nc: 256\n use_spect: False\n \n\n# Data options.\ndata:\n type: data.vox_dataset::VoxDataset\n path: ./dataset/vox_lmdb_demo\n resolution: 256\n semantic_radius: 13\n train:\n batch_size: 5\n distributed: True\n val:\n batch_size: 8\n distributed: True\n\n\n"
},
{
"path": "config.py",
"content": "import collections\nimport functools\nimport os\nimport re\n\nimport yaml\nfrom util.distributed import master_only_print as print\n\n\nclass AttrDict(dict):\n \"\"\"Dict as attribute trick.\"\"\"\n\n def __init__(self, *args, **kwargs):\n super(AttrDict, self).__init__(*args, **kwargs)\n self.__dict__ = self\n for key, value in self.__dict__.items():\n if isinstance(value, dict):\n self.__dict__[key] = AttrDict(value)\n elif isinstance(value, (list, tuple)):\n if isinstance(value[0], dict):\n self.__dict__[key] = [AttrDict(item) for item in value]\n else:\n self.__dict__[key] = value\n\n def yaml(self):\n \"\"\"Convert object to yaml dict and return.\"\"\"\n yaml_dict = {}\n for key, value in self.__dict__.items():\n if isinstance(value, AttrDict):\n yaml_dict[key] = value.yaml()\n elif isinstance(value, list):\n if isinstance(value[0], AttrDict):\n new_l = []\n for item in value:\n new_l.append(item.yaml())\n yaml_dict[key] = new_l\n else:\n yaml_dict[key] = value\n else:\n yaml_dict[key] = value\n return yaml_dict\n\n def __repr__(self):\n \"\"\"Print all variables.\"\"\"\n ret_str = []\n for key, value in self.__dict__.items():\n if isinstance(value, AttrDict):\n ret_str.append('{}:'.format(key))\n child_ret_str = value.__repr__().split('\\n')\n for item in child_ret_str:\n ret_str.append(' ' + item)\n elif isinstance(value, list):\n if isinstance(value[0], AttrDict):\n ret_str.append('{}:'.format(key))\n for item in value:\n # Treat as AttrDict above.\n child_ret_str = item.__repr__().split('\\n')\n for item in child_ret_str:\n ret_str.append(' ' + item)\n else:\n ret_str.append('{}: {}'.format(key, value))\n else:\n ret_str.append('{}: {}'.format(key, value))\n return '\\n'.join(ret_str)\n\n\nclass Config(AttrDict):\n r\"\"\"Configuration class. This should include every human specifiable\n hyperparameter values for your training.\"\"\"\n\n def __init__(self, filename=None, args=None, verbose=False, is_train=True):\n super(Config, self).__init__()\n # Set default parameters.\n # Logging.\n\n large_number = 1000000000\n self.snapshot_save_iter = large_number\n self.snapshot_save_epoch = large_number\n self.snapshot_save_start_iter = 0\n self.snapshot_save_start_epoch = 0\n self.image_save_iter = large_number\n self.eval_epoch = large_number\n self.start_eval_epoch = large_number\n self.eval_epoch = large_number\n self.max_epoch = large_number\n self.max_iter = large_number\n self.logging_iter = 100\n self.image_to_tensorboard=False\n self.which_iter = args.which_iter\n self.resume = not args.no_resume\n\n\n self.checkpoints_dir = args.checkpoints_dir\n self.name = args.name\n self.phase = 'train' if is_train else 'test'\n\n # Networks.\n self.gen = AttrDict(type='generators.dummy')\n self.dis = AttrDict(type='discriminators.dummy')\n\n # Optimizers.\n self.gen_optimizer = AttrDict(type='adam',\n lr=0.0001,\n adam_beta1=0.0,\n adam_beta2=0.999,\n eps=1e-8,\n lr_policy=AttrDict(iteration_mode=False,\n type='step',\n step_size=large_number,\n gamma=1))\n self.dis_optimizer = AttrDict(type='adam',\n lr=0.0001,\n adam_beta1=0.0,\n adam_beta2=0.999,\n eps=1e-8,\n lr_policy=AttrDict(iteration_mode=False,\n type='step',\n step_size=large_number,\n gamma=1))\n # Data.\n self.data = AttrDict(name='dummy',\n type='datasets.images',\n num_workers=0)\n self.test_data = AttrDict(name='dummy',\n type='datasets.images',\n num_workers=0,\n test=AttrDict(is_lmdb=False,\n roots='',\n batch_size=1))\n self.trainer = AttrDict(\n model_average=False,\n model_average_beta=0.9999,\n model_average_start_iteration=1000,\n model_average_batch_norm_estimation_iteration=30,\n model_average_remove_sn=True,\n image_to_tensorboard=False,\n hparam_to_tensorboard=False,\n distributed_data_parallel='pytorch',\n delay_allreduce=True,\n gan_relativistic=False,\n gen_step=1,\n dis_step=1)\n\n # # Cudnn.\n self.cudnn = AttrDict(deterministic=False,\n benchmark=True)\n\n # Others.\n self.pretrained_weight = ''\n self.inference_args = AttrDict()\n\n\n # Update with given configurations.\n assert os.path.exists(filename), 'File {} not exist.'.format(filename)\n loader = yaml.SafeLoader\n loader.add_implicit_resolver(\n u'tag:yaml.org,2002:float',\n re.compile(u'''^(?:\n [-+]?(?:[0-9][0-9_]*)\\\\.[0-9_]*(?:[eE][-+]?[0-9]+)?\n |[-+]?(?:[0-9][0-9_]*)(?:[eE][-+]?[0-9]+)\n |\\\\.[0-9_]+(?:[eE][-+][0-9]+)?\n |[-+]?[0-9][0-9_]*(?::[0-5]?[0-9])+\\\\.[0-9_]*\n |[-+]?\\\\.(?:inf|Inf|INF)\n |\\\\.(?:nan|NaN|NAN))$''', re.X),\n list(u'-+0123456789.'))\n try:\n with open(filename, 'r') as f:\n cfg_dict = yaml.load(f, Loader=loader)\n except EnvironmentError:\n print('Please check the file with name of \"%s\"', filename)\n recursive_update(self, cfg_dict)\n\n # Put common opts in both gen and dis.\n if 'common' in cfg_dict:\n self.common = AttrDict(**cfg_dict['common'])\n self.gen.common = self.common\n self.dis.common = self.common\n\n\n if verbose:\n print(' config '.center(80, '-'))\n print(self.__repr__())\n print(''.center(80, '-'))\n\n\ndef rsetattr(obj, attr, val):\n \"\"\"Recursively find object and set value\"\"\"\n pre, _, post = attr.rpartition('.')\n return setattr(rgetattr(obj, pre) if pre else obj, post, val)\n\n\ndef rgetattr(obj, attr, *args):\n \"\"\"Recursively find object and return value\"\"\"\n\n def _getattr(obj, attr):\n r\"\"\"Get attribute.\"\"\"\n return getattr(obj, attr, *args)\n\n return functools.reduce(_getattr, [obj] + attr.split('.'))\n\n\ndef recursive_update(d, u):\n \"\"\"Recursively update AttrDict d with AttrDict u\"\"\"\n for key, value in u.items():\n if isinstance(value, collections.abc.Mapping):\n d.__dict__[key] = recursive_update(d.get(key, AttrDict({})), value)\n elif isinstance(value, (list, tuple)):\n if isinstance(value[0], dict):\n d.__dict__[key] = [AttrDict(item) for item in value]\n else:\n d.__dict__[key] = value\n else:\n d.__dict__[key] = value\n return d\n"
},
{
"path": "data/__init__.py",
"content": "import importlib\n\nimport torch.utils.data\nfrom util.distributed import master_only_print as print\n\ndef find_dataset_using_name(dataset_name):\n dataset_filename = dataset_name\n module, target = dataset_name.split('::')\n datasetlib = importlib.import_module(module)\n dataset = None\n for name, cls in datasetlib.__dict__.items():\n if name == target:\n dataset = cls\n \n if dataset is None:\n raise ValueError(\"In %s.py, there should be a class \"\n \"with class name that matches %s in lowercase.\" %\n (dataset_filename, target))\n\n return dataset\n\n\ndef get_option_setter(dataset_name): \n dataset_class = find_dataset_using_name(dataset_name)\n return dataset_class.modify_commandline_options\n\n\ndef create_dataloader(opt, is_inference):\n dataset = find_dataset_using_name(opt.type)\n instance = dataset(opt, is_inference)\n phase = 'val' if is_inference else 'training'\n batch_size = opt.val.batch_size if is_inference else opt.train.batch_size\n print(\"%s dataset [%s] of size %d was created\" %\n (phase, opt.type, len(instance)))\n dataloader = torch.utils.data.DataLoader(\n instance,\n batch_size=batch_size,\n sampler=data_sampler(instance, shuffle=not is_inference, distributed=opt.train.distributed),\n drop_last=not is_inference,\n num_workers=getattr(opt, 'num_workers', 0),\n ) \n\n return dataloader\n\n\ndef data_sampler(dataset, shuffle, distributed):\n if distributed:\n return torch.utils.data.distributed.DistributedSampler(dataset, shuffle=shuffle)\n if shuffle:\n return torch.utils.data.RandomSampler(dataset)\n else:\n return torch.utils.data.SequentialSampler(dataset)\n\n\ndef get_dataloader(opt, is_inference=False):\n dataset = create_dataloader(opt, is_inference=is_inference)\n return dataset\n\n\ndef get_train_val_dataloader(opt):\n val_dataset = create_dataloader(opt, is_inference=True)\n train_dataset = create_dataloader(opt, is_inference=False)\n return val_dataset, train_dataset\n"
},
{
"path": "data/image_dataset.py",
"content": "import os\nimport glob\nimport time\nimport numpy as np\nfrom PIL import Image\n\nimport torch\nimport torchvision.transforms.functional as F\n\n\n\nclass ImageDataset():\n def __init__(self, opt, input_name):\n self.opt = opt\n self.IMAGEEXT = ['png', 'jpg']\n self.input_image_list, self.coeff_list = self.obtain_inputs(input_name)\n self.index = -1\n # load image dataset opt\n self.resolution = opt.resolution\n self.semantic_radius = opt.semantic_radius\n\n def next_image(self):\n self.index += 1\n image_name = self.input_image_list[self.index]\n coeff_name = self.coeff_list[self.index]\n img = Image.open(image_name)\n input_image = self.trans_image(img)\n\n coeff_3dmm = np.loadtxt(coeff_name).astype(np.float32)\n coeff_3dmm = self.transform_semantic(coeff_3dmm)\n \n return {\n 'source_image': input_image[None],\n 'target_semantics': coeff_3dmm[None],\n 'name': os.path.splitext(os.path.basename(image_name))[0]\n }\n\n def obtain_inputs(self, root):\n filenames = list()\n\n IMAGE_EXTENSIONS_LOWERCASE = {'jpg', 'png', 'jpeg', 'webp'}\n IMAGE_EXTENSIONS = IMAGE_EXTENSIONS_LOWERCASE.union({f.upper() for f in IMAGE_EXTENSIONS_LOWERCASE})\n extensions = IMAGE_EXTENSIONS\n\n for ext in extensions:\n filenames += glob.glob(f'{root}/*.{ext}', recursive=True)\n filenames = sorted(filenames)\n coeffnames = sorted(glob.glob(f'{root}/*_3dmm_coeff.txt')) \n\n return filenames, coeffnames\n\n def transform_semantic(self, semantic):\n semantic = semantic[None].repeat(self.semantic_radius*2+1, 0)\n ex_coeff = semantic[:,80:144] #expression\n angles = semantic[:,224:227] #euler angles for pose\n translation = semantic[:,254:257] #translation\n crop = semantic[:,259:262] #crop param\n\n coeff_3dmm = np.concatenate([ex_coeff, angles, translation, crop], 1)\n return torch.Tensor(coeff_3dmm).permute(1,0) \n\n def trans_image(self, image):\n image = F.resize(\n image, size=self.resolution, interpolation=Image.BICUBIC)\n image = F.to_tensor(image)\n image = F.normalize(image, mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))\n return image\n \n def __len__(self):\n return len(self.input_image_list)\n\n \n"
},
{
"path": "data/vox_dataset.py",
"content": "import os\nimport lmdb\nimport random\nimport collections\nimport numpy as np\nfrom PIL import Image\nfrom io import BytesIO\n\nimport torch\nfrom torch.utils.data import Dataset\nfrom torchvision import transforms\n\ndef format_for_lmdb(*args):\n key_parts = []\n for arg in args:\n if isinstance(arg, int):\n arg = str(arg).zfill(7)\n key_parts.append(arg)\n return '-'.join(key_parts).encode('utf-8')\n\nclass VoxDataset(Dataset):\n def __init__(self, opt, is_inference):\n path = opt.path\n self.env = lmdb.open(\n os.path.join(path, str(opt.resolution)),\n max_readers=32,\n readonly=True,\n lock=False,\n readahead=False,\n meminit=False,\n )\n\n if not self.env:\n raise IOError('Cannot open lmdb dataset', path)\n list_file = \"test_list.txt\" if is_inference else \"train_list.txt\"\n list_file = os.path.join(path, list_file)\n with open(list_file, 'r') as f:\n lines = f.readlines()\n videos = [line.replace('\\n', '') for line in lines]\n\n self.resolution = opt.resolution\n self.semantic_radius = opt.semantic_radius\n self.video_items, self.person_ids = self.get_video_index(videos)\n self.idx_by_person_id = self.group_by_key(self.video_items, key='person_id')\n self.person_ids = self.person_ids * 100\n\n self.transform = transforms.Compose(\n [\n transforms.ToTensor(),\n transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True),\n ])\n\n def get_video_index(self, videos):\n video_items = []\n for video in videos:\n video_items.append(self.Video_Item(video))\n\n person_ids = sorted(list({video.split('#')[0] for video in videos}))\n\n return video_items, person_ids \n\n def group_by_key(self, video_list, key):\n return_dict = collections.defaultdict(list)\n for index, video_item in enumerate(video_list):\n return_dict[video_item[key]].append(index)\n return return_dict \n \n def Video_Item(self, video_name):\n video_item = {}\n video_item['video_name'] = video_name\n video_item['person_id'] = video_name.split('#')[0]\n with self.env.begin(write=False) as txn:\n key = format_for_lmdb(video_item['video_name'], 'length')\n length = int(txn.get(key).decode('utf-8'))\n video_item['num_frame'] = length\n \n return video_item\n\n def __len__(self):\n return len(self.person_ids)\n\n def __getitem__(self, index):\n data={}\n person_id = self.person_ids[index]\n video_item = self.video_items[random.choices(self.idx_by_person_id[person_id], k=1)[0]]\n frame_source, frame_target = self.random_select_frames(video_item)\n\n with self.env.begin(write=False) as txn:\n key = format_for_lmdb(video_item['video_name'], frame_source)\n img_bytes_1 = txn.get(key) \n key = format_for_lmdb(video_item['video_name'], frame_target)\n img_bytes_2 = txn.get(key) \n semantics_key = format_for_lmdb(video_item['video_name'], 'coeff_3dmm')\n semantics_numpy = np.frombuffer(txn.get(semantics_key), dtype=np.float32)\n semantics_numpy = semantics_numpy.reshape((video_item['num_frame'],-1))\n\n img1 = Image.open(BytesIO(img_bytes_1))\n data['source_image'] = self.transform(img1)\n\n img2 = Image.open(BytesIO(img_bytes_2))\n data['target_image'] = self.transform(img2) \n\n data['target_semantics'] = self.transform_semantic(semantics_numpy, frame_target)\n data['source_semantics'] = self.transform_semantic(semantics_numpy, frame_source)\n \n return data\n \n def random_select_frames(self, video_item):\n num_frame = video_item['num_frame']\n frame_idx = random.choices(list(range(num_frame)), k=2)\n return frame_idx[0], frame_idx[1]\n\n def transform_semantic(self, semantic, frame_index):\n index = self.obtain_seq_index(frame_index, semantic.shape[0])\n coeff_3dmm = semantic[index,...]\n # id_coeff = coeff_3dmm[:,:80] #identity\n ex_coeff = coeff_3dmm[:,80:144] #expression\n # tex_coeff = coeff_3dmm[:,144:224] #texture\n angles = coeff_3dmm[:,224:227] #euler angles for pose\n # gamma = coeff_3dmm[:,227:254] #lighting\n translation = coeff_3dmm[:,254:257] #translation\n crop = coeff_3dmm[:,257:260] #crop param\n\n coeff_3dmm = np.concatenate([ex_coeff, angles, translation, crop], 1)\n return torch.Tensor(coeff_3dmm).permute(1,0)\n\n def obtain_seq_index(self, index, num_frames):\n seq = list(range(index-self.semantic_radius, index+self.semantic_radius+1))\n seq = [ min(max(item, 0), num_frames-1) for item in seq ]\n return seq\n\n\n\n"
},
{
"path": "data/vox_video_dataset.py",
"content": "import os\nimport lmdb\nimport random\nimport collections\nimport numpy as np\nfrom PIL import Image\nfrom io import BytesIO\n\nimport torch\n\nfrom data.vox_dataset import VoxDataset\nfrom data.vox_dataset import format_for_lmdb\n\nclass VoxVideoDataset(VoxDataset):\n def __init__(self, opt, is_inference):\n super(VoxVideoDataset, self).__init__(opt, is_inference)\n self.video_index = -1\n self.cross_id = opt.cross_id\n # whether normalize the crop parameters when performing cross_id reenactments\n # set it as \"True\" always brings better performance\n self.norm_crop_param = True\n\n def __len__(self):\n return len(self.video_items)\n\n def load_next_video(self):\n data={}\n self.video_index += 1\n video_item = self.video_items[self.video_index]\n source_video_item = self.random_video(video_item) if self.cross_id else video_item \n\n with self.env.begin(write=False) as txn:\n key = format_for_lmdb(source_video_item['video_name'], 0)\n img_bytes_1 = txn.get(key) \n img1 = Image.open(BytesIO(img_bytes_1))\n data['source_image'] = self.transform(img1)\n\n semantics_key = format_for_lmdb(video_item['video_name'], 'coeff_3dmm')\n semantics_numpy = np.frombuffer(txn.get(semantics_key), dtype=np.float32)\n semantics_numpy = semantics_numpy.reshape((video_item['num_frame'],-1))\n if self.cross_id and self.norm_crop_param:\n semantics_source_key = format_for_lmdb(source_video_item['video_name'], 'coeff_3dmm')\n semantics_source_numpy = np.frombuffer(txn.get(semantics_source_key), dtype=np.float32)\n semantic_source_numpy = semantics_source_numpy.reshape((source_video_item['num_frame'],-1))[0:1]\n crop_norm_ratio = self.find_crop_norm_ratio(semantic_source_numpy, semantics_numpy)\n else:\n crop_norm_ratio = None \n\n data['target_image'], data['target_semantics'] = [], []\n for frame_index in range(video_item['num_frame']):\n key = format_for_lmdb(video_item['video_name'], frame_index)\n img_bytes_1 = txn.get(key) \n img1 = Image.open(BytesIO(img_bytes_1))\n data['target_image'].append(self.transform(img1))\n data['target_semantics'].append(\n self.transform_semantic(semantics_numpy, frame_index, crop_norm_ratio)\n )\n data['video_name'] = self.obtain_name(video_item['video_name'], source_video_item['video_name'])\n return data \n \n def random_video(self, target_video_item):\n target_person_id = target_video_item['person_id']\n assert len(self.person_ids) > 1 \n source_person_id = np.random.choice(self.person_ids)\n if source_person_id == target_person_id:\n source_person_id = np.random.choice(self.person_ids)\n source_video_index = np.random.choice(self.idx_by_person_id[source_person_id])\n source_video_item = self.video_items[source_video_index]\n return source_video_item\n\n def find_crop_norm_ratio(self, source_coeff, target_coeffs):\n alpha = 0.3\n exp_diff = np.mean(np.abs(target_coeffs[:,80:144] - source_coeff[:,80:144]), 1)\n angle_diff = np.mean(np.abs(target_coeffs[:,224:227] - source_coeff[:,224:227]), 1)\n index = np.argmin(alpha*exp_diff + (1-alpha)*angle_diff)\n crop_norm_ratio = source_coeff[:,-3] / target_coeffs[index:index+1, -3]\n return crop_norm_ratio\n \n def transform_semantic(self, semantic, frame_index, crop_norm_ratio):\n index = self.obtain_seq_index(frame_index, semantic.shape[0])\n coeff_3dmm = semantic[index,...]\n # id_coeff = coeff_3dmm[:,:80] #identity\n ex_coeff = coeff_3dmm[:,80:144] #expression\n # tex_coeff = coeff_3dmm[:,144:224] #texture\n angles = coeff_3dmm[:,224:227] #euler angles for pose\n # gamma = coeff_3dmm[:,227:254] #lighting\n translation = coeff_3dmm[:,254:257] #translation\n crop = coeff_3dmm[:,257:300] #crop param\n\n if self.cross_id and self.norm_crop_param:\n crop[:, -3] = crop[:, -3] * crop_norm_ratio\n\n coeff_3dmm = np.concatenate([ex_coeff, angles, translation, crop], 1)\n return torch.Tensor(coeff_3dmm).permute(1,0) \n\n def obtain_name(self, target_name, source_name):\n if not self.cross_id:\n return target_name\n else:\n source_name = os.path.splitext(os.path.basename(source_name))[0]\n return source_name+'_to_'+target_name"
},
{
"path": "generators/base_function.py",
"content": "import sys\nimport math\n\nimport torch\nfrom torch import nn\nfrom torch.nn import functional as F\nfrom torch.autograd import Function\nfrom torch.nn.utils.spectral_norm import spectral_norm as SpectralNorm\n\n\nclass LayerNorm2d(nn.Module):\n def __init__(self, n_out, affine=True):\n super(LayerNorm2d, self).__init__()\n self.n_out = n_out\n self.affine = affine\n\n if self.affine:\n self.weight = nn.Parameter(torch.ones(n_out, 1, 1))\n self.bias = nn.Parameter(torch.zeros(n_out, 1, 1))\n\n def forward(self, x):\n normalized_shape = x.size()[1:]\n if self.affine:\n return F.layer_norm(x, normalized_shape, \\\n self.weight.expand(normalized_shape), \n self.bias.expand(normalized_shape))\n \n else:\n return F.layer_norm(x, normalized_shape) \n\nclass ADAINHourglass(nn.Module):\n def __init__(self, image_nc, pose_nc, ngf, img_f, encoder_layers, decoder_layers, nonlinearity, use_spect):\n super(ADAINHourglass, self).__init__()\n self.encoder = ADAINEncoder(image_nc, pose_nc, ngf, img_f, encoder_layers, nonlinearity, use_spect)\n self.decoder = ADAINDecoder(pose_nc, ngf, img_f, encoder_layers, decoder_layers, True, nonlinearity, use_spect)\n self.output_nc = self.decoder.output_nc\n\n def forward(self, x, z):\n return self.decoder(self.encoder(x, z), z) \n\n\n\nclass ADAINEncoder(nn.Module):\n def __init__(self, image_nc, pose_nc, ngf, img_f, layers, nonlinearity=nn.LeakyReLU(), use_spect=False):\n super(ADAINEncoder, self).__init__()\n self.layers = layers\n self.input_layer = nn.Conv2d(image_nc, ngf, kernel_size=7, stride=1, padding=3)\n for i in range(layers):\n in_channels = min(ngf * (2**i), img_f)\n out_channels = min(ngf *(2**(i+1)), img_f)\n model = ADAINEncoderBlock(in_channels, out_channels, pose_nc, nonlinearity, use_spect)\n setattr(self, 'encoder' + str(i), model)\n self.output_nc = out_channels\n \n def forward(self, x, z):\n out = self.input_layer(x)\n out_list = [out]\n for i in range(self.layers):\n model = getattr(self, 'encoder' + str(i))\n out = model(out, z)\n out_list.append(out)\n return out_list\n \nclass ADAINDecoder(nn.Module):\n \"\"\"docstring for ADAINDecoder\"\"\"\n def __init__(self, pose_nc, ngf, img_f, encoder_layers, decoder_layers, skip_connect=True, \n nonlinearity=nn.LeakyReLU(), use_spect=False):\n\n super(ADAINDecoder, self).__init__()\n self.encoder_layers = encoder_layers\n self.decoder_layers = decoder_layers\n self.skip_connect = skip_connect\n use_transpose = True\n\n for i in range(encoder_layers-decoder_layers, encoder_layers)[::-1]:\n in_channels = min(ngf * (2**(i+1)), img_f)\n in_channels = in_channels*2 if i != (encoder_layers-1) and self.skip_connect else in_channels\n out_channels = min(ngf * (2**i), img_f)\n model = ADAINDecoderBlock(in_channels, out_channels, out_channels, pose_nc, use_transpose, nonlinearity, use_spect)\n setattr(self, 'decoder' + str(i), model)\n\n self.output_nc = out_channels*2 if self.skip_connect else out_channels\n\n def forward(self, x, z):\n out = x.pop() if self.skip_connect else x\n for i in range(self.encoder_layers-self.decoder_layers, self.encoder_layers)[::-1]:\n model = getattr(self, 'decoder' + str(i))\n out = model(out, z)\n out = torch.cat([out, x.pop()], 1) if self.skip_connect else out\n return out\n\nclass ADAINEncoderBlock(nn.Module): \n def __init__(self, input_nc, output_nc, feature_nc, nonlinearity=nn.LeakyReLU(), use_spect=False):\n super(ADAINEncoderBlock, self).__init__()\n kwargs_down = {'kernel_size': 4, 'stride': 2, 'padding': 1}\n kwargs_fine = {'kernel_size': 3, 'stride': 1, 'padding': 1}\n\n self.conv_0 = spectral_norm(nn.Conv2d(input_nc, output_nc, **kwargs_down), use_spect)\n self.conv_1 = spectral_norm(nn.Conv2d(output_nc, output_nc, **kwargs_fine), use_spect)\n\n\n self.norm_0 = ADAIN(input_nc, feature_nc)\n self.norm_1 = ADAIN(output_nc, feature_nc)\n self.actvn = nonlinearity\n\n def forward(self, x, z):\n x = self.conv_0(self.actvn(self.norm_0(x, z)))\n x = self.conv_1(self.actvn(self.norm_1(x, z)))\n return x\n\nclass ADAINDecoderBlock(nn.Module):\n def __init__(self, input_nc, output_nc, hidden_nc, feature_nc, use_transpose=True, nonlinearity=nn.LeakyReLU(), use_spect=False):\n super(ADAINDecoderBlock, self).__init__() \n # Attributes\n self.actvn = nonlinearity\n hidden_nc = min(input_nc, output_nc) if hidden_nc is None else hidden_nc\n\n kwargs_fine = {'kernel_size':3, 'stride':1, 'padding':1}\n if use_transpose:\n kwargs_up = {'kernel_size':3, 'stride':2, 'padding':1, 'output_padding':1}\n else:\n kwargs_up = {'kernel_size':3, 'stride':1, 'padding':1}\n\n # create conv layers\n self.conv_0 = spectral_norm(nn.Conv2d(input_nc, hidden_nc, **kwargs_fine), use_spect)\n if use_transpose:\n self.conv_1 = spectral_norm(nn.ConvTranspose2d(hidden_nc, output_nc, **kwargs_up), use_spect)\n self.conv_s = spectral_norm(nn.ConvTranspose2d(input_nc, output_nc, **kwargs_up), use_spect)\n else:\n self.conv_1 = nn.Sequential(spectral_norm(nn.Conv2d(hidden_nc, output_nc, **kwargs_up), use_spect),\n nn.Upsample(scale_factor=2))\n self.conv_s = nn.Sequential(spectral_norm(nn.Conv2d(input_nc, output_nc, **kwargs_up), use_spect),\n nn.Upsample(scale_factor=2))\n # define normalization layers\n self.norm_0 = ADAIN(input_nc, feature_nc)\n self.norm_1 = ADAIN(hidden_nc, feature_nc)\n self.norm_s = ADAIN(input_nc, feature_nc)\n \n def forward(self, x, z):\n x_s = self.shortcut(x, z)\n dx = self.conv_0(self.actvn(self.norm_0(x, z)))\n dx = self.conv_1(self.actvn(self.norm_1(dx, z)))\n out = x_s + dx\n return out\n\n def shortcut(self, x, z):\n x_s = self.conv_s(self.actvn(self.norm_s(x, z)))\n return x_s \n\n\ndef spectral_norm(module, use_spect=True):\n \"\"\"use spectral normal layer to stable the training process\"\"\"\n if use_spect:\n return SpectralNorm(module)\n else:\n return module\n\n\nclass ADAIN(nn.Module):\n def __init__(self, norm_nc, feature_nc):\n super().__init__()\n\n self.param_free_norm = nn.InstanceNorm2d(norm_nc, affine=False)\n\n nhidden = 128\n use_bias=True\n\n self.mlp_shared = nn.Sequential(\n nn.Linear(feature_nc, nhidden, bias=use_bias), \n nn.ReLU()\n )\n self.mlp_gamma = nn.Linear(nhidden, norm_nc, bias=use_bias) \n self.mlp_beta = nn.Linear(nhidden, norm_nc, bias=use_bias) \n\n def forward(self, x, feature):\n\n # Part 1. generate parameter-free normalized activations\n normalized = self.param_free_norm(x)\n\n # Part 2. produce scaling and bias conditioned on feature\n feature = feature.view(feature.size(0), -1)\n actv = self.mlp_shared(feature)\n gamma = self.mlp_gamma(actv)\n beta = self.mlp_beta(actv)\n\n # apply scale and bias\n gamma = gamma.view(*gamma.size()[:2], 1,1)\n beta = beta.view(*beta.size()[:2], 1,1)\n out = normalized * (1 + gamma) + beta\n return out\n\n\nclass FineEncoder(nn.Module):\n \"\"\"docstring for Encoder\"\"\"\n def __init__(self, image_nc, ngf, img_f, layers, norm_layer=nn.BatchNorm2d, nonlinearity=nn.LeakyReLU(), use_spect=False):\n super(FineEncoder, self).__init__()\n self.layers = layers\n self.first = FirstBlock2d(image_nc, ngf, norm_layer, nonlinearity, use_spect)\n for i in range(layers):\n in_channels = min(ngf*(2**i), img_f)\n out_channels = min(ngf*(2**(i+1)), img_f)\n model = DownBlock2d(in_channels, out_channels, norm_layer, nonlinearity, use_spect)\n setattr(self, 'down' + str(i), model)\n self.output_nc = out_channels\n\n def forward(self, x):\n x = self.first(x)\n out=[x]\n for i in range(self.layers):\n model = getattr(self, 'down'+str(i))\n x = model(x)\n out.append(x)\n return out\n\nclass FineDecoder(nn.Module):\n \"\"\"docstring for FineDecoder\"\"\"\n def __init__(self, image_nc, feature_nc, ngf, img_f, layers, num_block, norm_layer=nn.BatchNorm2d, nonlinearity=nn.LeakyReLU(), use_spect=False):\n super(FineDecoder, self).__init__()\n self.layers = layers\n for i in range(layers)[::-1]:\n in_channels = min(ngf*(2**(i+1)), img_f)\n out_channels = min(ngf*(2**i), img_f)\n up = UpBlock2d(in_channels, out_channels, norm_layer, nonlinearity, use_spect)\n res = FineADAINResBlocks(num_block, in_channels, feature_nc, norm_layer, nonlinearity, use_spect)\n jump = Jump(out_channels, norm_layer, nonlinearity, use_spect)\n\n setattr(self, 'up' + str(i), up)\n setattr(self, 'res' + str(i), res) \n setattr(self, 'jump' + str(i), jump)\n\n self.final = FinalBlock2d(out_channels, image_nc, use_spect, 'tanh')\n\n self.output_nc = out_channels\n\n def forward(self, x, z):\n out = x.pop()\n for i in range(self.layers)[::-1]:\n res_model = getattr(self, 'res' + str(i))\n up_model = getattr(self, 'up' + str(i))\n jump_model = getattr(self, 'jump' + str(i))\n out = res_model(out, z)\n out = up_model(out)\n out = jump_model(x.pop()) + out\n out_image = self.final(out)\n return out_image\n\nclass FirstBlock2d(nn.Module):\n \"\"\"\n Downsampling block for use in encoder.\n \"\"\"\n def __init__(self, input_nc, output_nc, norm_layer=nn.BatchNorm2d, nonlinearity=nn.LeakyReLU(), use_spect=False):\n super(FirstBlock2d, self).__init__()\n kwargs = {'kernel_size': 7, 'stride': 1, 'padding': 3}\n conv = spectral_norm(nn.Conv2d(input_nc, output_nc, **kwargs), use_spect)\n\n if type(norm_layer) == type(None):\n self.model = nn.Sequential(conv, nonlinearity)\n else:\n self.model = nn.Sequential(conv, norm_layer(output_nc), nonlinearity)\n\n\n def forward(self, x):\n out = self.model(x)\n return out \n\nclass DownBlock2d(nn.Module):\n def __init__(self, input_nc, output_nc, norm_layer=nn.BatchNorm2d, nonlinearity=nn.LeakyReLU(), use_spect=False):\n super(DownBlock2d, self).__init__()\n\n\n kwargs = {'kernel_size': 3, 'stride': 1, 'padding': 1}\n conv = spectral_norm(nn.Conv2d(input_nc, output_nc, **kwargs), use_spect)\n pool = nn.AvgPool2d(kernel_size=(2, 2))\n\n if type(norm_layer) == type(None):\n self.model = nn.Sequential(conv, nonlinearity, pool)\n else:\n self.model = nn.Sequential(conv, norm_layer(output_nc), nonlinearity, pool)\n\n def forward(self, x):\n out = self.model(x)\n return out \n\nclass UpBlock2d(nn.Module):\n def __init__(self, input_nc, output_nc, norm_layer=nn.BatchNorm2d, nonlinearity=nn.LeakyReLU(), use_spect=False):\n super(UpBlock2d, self).__init__()\n kwargs = {'kernel_size': 3, 'stride': 1, 'padding': 1}\n conv = spectral_norm(nn.Conv2d(input_nc, output_nc, **kwargs), use_spect)\n if type(norm_layer) == type(None):\n self.model = nn.Sequential(conv, nonlinearity)\n else:\n self.model = nn.Sequential(conv, norm_layer(output_nc), nonlinearity)\n\n def forward(self, x):\n out = self.model(F.interpolate(x, scale_factor=2))\n return out\n\nclass FineADAINResBlocks(nn.Module):\n def __init__(self, num_block, input_nc, feature_nc, norm_layer=nn.BatchNorm2d, nonlinearity=nn.LeakyReLU(), use_spect=False):\n super(FineADAINResBlocks, self).__init__() \n self.num_block = num_block\n for i in range(num_block):\n model = FineADAINResBlock2d(input_nc, feature_nc, norm_layer, nonlinearity, use_spect)\n setattr(self, 'res'+str(i), model)\n\n def forward(self, x, z):\n for i in range(self.num_block):\n model = getattr(self, 'res'+str(i))\n x = model(x, z)\n return x \n\nclass Jump(nn.Module):\n def __init__(self, input_nc, norm_layer=nn.BatchNorm2d, nonlinearity=nn.LeakyReLU(), use_spect=False):\n super(Jump, self).__init__()\n kwargs = {'kernel_size': 3, 'stride': 1, 'padding': 1}\n conv = spectral_norm(nn.Conv2d(input_nc, input_nc, **kwargs), use_spect)\n\n if type(norm_layer) == type(None):\n self.model = nn.Sequential(conv, nonlinearity)\n else:\n self.model = nn.Sequential(conv, norm_layer(input_nc), nonlinearity)\n\n def forward(self, x):\n out = self.model(x)\n return out \n\nclass FineADAINResBlock2d(nn.Module):\n \"\"\"\n Define an Residual block for different types\n \"\"\"\n def __init__(self, input_nc, feature_nc, norm_layer=nn.BatchNorm2d, nonlinearity=nn.LeakyReLU(), use_spect=False):\n super(FineADAINResBlock2d, self).__init__()\n\n kwargs = {'kernel_size': 3, 'stride': 1, 'padding': 1}\n\n self.conv1 = spectral_norm(nn.Conv2d(input_nc, input_nc, **kwargs), use_spect)\n self.conv2 = spectral_norm(nn.Conv2d(input_nc, input_nc, **kwargs), use_spect)\n self.norm1 = ADAIN(input_nc, feature_nc)\n self.norm2 = ADAIN(input_nc, feature_nc)\n\n self.actvn = nonlinearity\n\n\n def forward(self, x, z):\n dx = self.actvn(self.norm1(self.conv1(x), z))\n dx = self.norm2(self.conv2(x), z)\n out = dx + x\n return out \n\nclass FinalBlock2d(nn.Module):\n \"\"\"\n Define the output layer\n \"\"\"\n def __init__(self, input_nc, output_nc, use_spect=False, tanh_or_sigmoid='tanh'):\n super(FinalBlock2d, self).__init__()\n\n kwargs = {'kernel_size': 7, 'stride': 1, 'padding':3}\n conv = spectral_norm(nn.Conv2d(input_nc, output_nc, **kwargs), use_spect)\n\n if tanh_or_sigmoid == 'sigmoid':\n out_nonlinearity = nn.Sigmoid()\n else:\n out_nonlinearity = nn.Tanh() \n\n self.model = nn.Sequential(conv, out_nonlinearity)\n def forward(self, x):\n out = self.model(x)\n return out "
},
{
"path": "generators/face_model.py",
"content": "import functools\nimport numpy as np\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom util import flow_util\nfrom generators.base_function import LayerNorm2d, ADAINHourglass, FineEncoder, FineDecoder\n\nclass FaceGenerator(nn.Module):\n def __init__(\n self, \n mapping_net, \n warpping_net, \n editing_net, \n common\n ): \n super(FaceGenerator, self).__init__()\n self.mapping_net = MappingNet(**mapping_net)\n self.warpping_net = WarpingNet(**warpping_net, **common)\n self.editing_net = EditingNet(**editing_net, **common)\n \n def forward(\n self, \n input_image, \n driving_source, \n stage=None\n ):\n if stage == 'warp':\n descriptor = self.mapping_net(driving_source)\n output = self.warpping_net(input_image, descriptor)\n else:\n descriptor = self.mapping_net(driving_source)\n output = self.warpping_net(input_image, descriptor)\n output['fake_image'] = self.editing_net(input_image, output['warp_image'], descriptor)\n return output\n\nclass MappingNet(nn.Module):\n def __init__(self, coeff_nc, descriptor_nc, layer):\n super( MappingNet, self).__init__()\n\n self.layer = layer\n nonlinearity = nn.LeakyReLU(0.1)\n\n self.first = nn.Sequential(\n torch.nn.Conv1d(coeff_nc, descriptor_nc, kernel_size=7, padding=0, bias=True))\n\n for i in range(layer):\n net = nn.Sequential(nonlinearity,\n torch.nn.Conv1d(descriptor_nc, descriptor_nc, kernel_size=3, padding=0, dilation=3))\n setattr(self, 'encoder' + str(i), net) \n\n self.pooling = nn.AdaptiveAvgPool1d(1)\n self.output_nc = descriptor_nc\n\n def forward(self, input_3dmm):\n out = self.first(input_3dmm)\n for i in range(self.layer):\n model = getattr(self, 'encoder' + str(i))\n out = model(out) + out[:,:,3:-3]\n out = self.pooling(out)\n return out \n\nclass WarpingNet(nn.Module):\n def __init__(\n self, \n image_nc, \n descriptor_nc, \n base_nc, \n max_nc, \n encoder_layer, \n decoder_layer, \n use_spect\n ):\n super( WarpingNet, self).__init__()\n\n nonlinearity = nn.LeakyReLU(0.1)\n norm_layer = functools.partial(LayerNorm2d, affine=True) \n kwargs = {'nonlinearity':nonlinearity, 'use_spect':use_spect}\n\n self.descriptor_nc = descriptor_nc \n self.hourglass = ADAINHourglass(image_nc, self.descriptor_nc, base_nc,\n max_nc, encoder_layer, decoder_layer, **kwargs)\n\n self.flow_out = nn.Sequential(norm_layer(self.hourglass.output_nc), \n nonlinearity,\n nn.Conv2d(self.hourglass.output_nc, 2, kernel_size=7, stride=1, padding=3))\n\n self.pool = nn.AdaptiveAvgPool2d(1)\n\n def forward(self, input_image, descriptor):\n final_output={}\n output = self.hourglass(input_image, descriptor)\n final_output['flow_field'] = self.flow_out(output)\n\n deformation = flow_util.convert_flow_to_deformation(final_output['flow_field'])\n final_output['warp_image'] = flow_util.warp_image(input_image, deformation)\n return final_output\n\n\nclass EditingNet(nn.Module):\n def __init__(\n self, \n image_nc, \n descriptor_nc, \n layer, \n base_nc, \n max_nc, \n num_res_blocks, \n use_spect): \n super(EditingNet, self).__init__()\n\n nonlinearity = nn.LeakyReLU(0.1)\n norm_layer = functools.partial(LayerNorm2d, affine=True) \n kwargs = {'norm_layer':norm_layer, 'nonlinearity':nonlinearity, 'use_spect':use_spect}\n self.descriptor_nc = descriptor_nc\n\n # encoder part\n self.encoder = FineEncoder(image_nc*2, base_nc, max_nc, layer, **kwargs)\n self.decoder = FineDecoder(image_nc, self.descriptor_nc, base_nc, max_nc, layer, num_res_blocks, **kwargs)\n\n def forward(self, input_image, warp_image, descriptor):\n x = torch.cat([input_image, warp_image], 1)\n x = self.encoder(x)\n gen_image = self.decoder(x, descriptor)\n return gen_image\n"
},
{
"path": "inference.py",
"content": "import os\nimport cv2 \nimport lmdb\nimport math\nimport argparse\nimport numpy as np\nfrom io import BytesIO\nfrom PIL import Image\n\nimport torch\nimport torchvision.transforms.functional as F\nimport torchvision.transforms as transforms\n\nfrom util.logging import init_logging, make_logging_dir\nfrom util.distributed import init_dist\nfrom util.trainer import get_model_optimizer_and_scheduler, set_random_seed, get_trainer\nfrom util.distributed import master_only_print as print\nfrom data.vox_video_dataset import VoxVideoDataset\nfrom config import Config\n\n\ndef parse_args():\n parser = argparse.ArgumentParser(description='Training')\n parser.add_argument('--config', default='./config/face.yaml')\n parser.add_argument('--name', default=None)\n parser.add_argument('--checkpoints_dir', default='result',\n help='Dir for saving logs and models.')\n parser.add_argument('--seed', type=int, default=0, help='Random seed.')\n parser.add_argument('--cross_id', action='store_true')\n parser.add_argument('--which_iter', type=int, default=None)\n parser.add_argument('--no_resume', action='store_true')\n parser.add_argument('--local_rank', type=int, default=0)\n parser.add_argument('--single_gpu', action='store_true')\n parser.add_argument('--output_dir', type=str)\n\n\n args = parser.parse_args()\n return args\n\ndef write2video(results_dir, *video_list):\n cat_video=None\n\n for video in video_list:\n video_numpy = video[:,:3,:,:].cpu().float().detach().numpy()\n video_numpy = (np.transpose(video_numpy, (0, 2, 3, 1)) + 1) / 2.0 * 255.0\n video_numpy = video_numpy.astype(np.uint8)\n cat_video = np.concatenate([cat_video, video_numpy], 2) if cat_video is not None else video_numpy\n\n image_array=[]\n for i in range(cat_video.shape[0]):\n image_array.append(cat_video[i]) \n\n out_name = results_dir+'.mp4' \n _, height, width, layers = cat_video.shape\n size = (width,height)\n out = cv2.VideoWriter(out_name, cv2.VideoWriter_fourcc(*'mp4v'), 15, size)\n\n for i in range(len(image_array)):\n out.write(image_array[i][:,:,::-1])\n out.release() \n\nif __name__ == '__main__':\n args = parse_args()\n set_random_seed(args.seed)\n opt = Config(args.config, args, is_train=False)\n\n if not args.single_gpu:\n opt.local_rank = args.local_rank\n init_dist(opt.local_rank) \n opt.device = torch.cuda.current_device()\n # create a visualizer\n date_uid, logdir = init_logging(opt)\n opt.logdir = logdir\n make_logging_dir(logdir, date_uid)\n\n # create a model\n net_G, net_G_ema, opt_G, sch_G \\\n = get_model_optimizer_and_scheduler(opt)\n\n trainer = get_trainer(opt, net_G, net_G_ema, \\\n opt_G, sch_G, None)\n\n current_epoch, current_iteration = trainer.load_checkpoint(\n opt, args.which_iter) \n net_G = trainer.net_G_ema.eval()\n\n output_dir = os.path.join(\n args.output_dir, \n 'epoch_{:05}_iteration_{:09}'.format(current_epoch, current_iteration)\n )\n os.makedirs(output_dir, exist_ok=True)\n opt.data.cross_id = args.cross_id\n dataset = VoxVideoDataset(opt.data, is_inference=True)\n with torch.no_grad():\n for video_index in range(dataset.__len__()):\n data = dataset.load_next_video()\n input_source = data['source_image'][None].cuda()\n name = data['video_name']\n\n output_images, gt_images, warp_images = [],[],[]\n for frame_index in range(len(data['target_semantics'])):\n target_semantic = data['target_semantics'][frame_index][None].cuda()\n output_dict = net_G(input_source, target_semantic)\n output_images.append(\n output_dict['fake_image'].cpu().clamp_(-1, 1)\n )\n warp_images.append(\n output_dict['warp_image'].cpu().clamp_(-1, 1)\n ) \n gt_images.append(\n data['target_image'][frame_index][None]\n )\n \n gen_images = torch.cat(output_images, 0)\n gt_images = torch.cat(gt_images, 0)\n warp_images = torch.cat(warp_images, 0)\n\n write2video(\"{}/{}\".format(output_dir, name), gt_images, warp_images, gen_images)\n print(\"write results to video {}/{}\".format(output_dir, name))\n\n"
},
{
"path": "intuitive_control.py",
"content": "import os\nimport math\nimport argparse\nimport numpy as np\nfrom scipy.io import savemat,loadmat\n\nimport torch\nimport torchvision.transforms.functional as F\nimport torchvision.transforms as transforms\n\nfrom config import Config\nfrom util.logging import init_logging, make_logging_dir\nfrom util.distributed import init_dist\nfrom util.trainer import get_model_optimizer_and_scheduler, set_random_seed, get_trainer\nfrom util.distributed import master_only_print as print\nfrom data.image_dataset import ImageDataset\nfrom inference import write2video\n\n\ndef parse_args():\n parser = argparse.ArgumentParser(description='Training')\n parser.add_argument('--config', default='./config/face.yaml')\n parser.add_argument('--name', default=None)\n parser.add_argument('--checkpoints_dir', default='result',\n help='Dir for saving logs and models.')\n parser.add_argument('--seed', type=int, default=0, help='Random seed.')\n parser.add_argument('--which_iter', type=int, default=None)\n parser.add_argument('--no_resume', action='store_true')\n parser.add_argument('--input_name', type=str)\n parser.add_argument('--local_rank', type=int, default=0)\n parser.add_argument('--single_gpu', action='store_true')\n parser.add_argument('--output_dir', type=str)\n\n args = parser.parse_args()\n return args\n\ndef get_control(input_name):\n control_dict = {}\n control_dict['rotation_center'] = torch.tensor([0,0,0,0,0,0.45])\n control_dict['rotation_left_x'] = torch.tensor([0,0,math.pi/10,0,0,0.45])\n control_dict['rotation_right_x'] = torch.tensor([0,0,-math.pi/10,0,0,0.45])\n\n control_dict['rotation_left_y'] = torch.tensor([math.pi/10,0,0,0,0,0.45])\n control_dict['rotation_right_y'] = torch.tensor([-math.pi/10,0,0,0,0,0.45]) \n\n control_dict['rotation_left_z'] = torch.tensor([0,math.pi/8,0,0,0,0.45])\n control_dict['rotation_right_z'] = torch.tensor([0,-math.pi/8,0,0,0,0.45]) \n\n expession = loadmat('{}/expression.mat'.format(input_name))\n\n for item in ['expression_center', 'expression_mouth', 'expression_eyebrow', 'expression_eyes']:\n control_dict[item] = torch.tensor(expession[item])[0]\n\n sort_rot_control = [\n 'rotation_left_x', 'rotation_center', \n 'rotation_right_x', 'rotation_center',\n 'rotation_left_y', 'rotation_center',\n 'rotation_right_y', 'rotation_center',\n 'rotation_left_z', 'rotation_center',\n 'rotation_right_z', 'rotation_center'\n ]\n \n sort_exp_control = [\n 'expression_center', 'expression_mouth',\n 'expression_center', 'expression_eyebrow',\n 'expression_center', 'expression_eyes',\n ]\n return control_dict, sort_rot_control, sort_exp_control\n\nif __name__ == '__main__':\n args = parse_args()\n set_random_seed(args.seed)\n opt = Config(args.config, args, is_train=False)\n\n if not args.single_gpu:\n opt.local_rank = args.local_rank\n init_dist(opt.local_rank) \n opt.device = torch.cuda.current_device()\n\n # create a visualizer\n date_uid, logdir = init_logging(opt)\n opt.logdir = logdir\n make_logging_dir(logdir, date_uid)\n\n # create a model\n net_G, net_G_ema, opt_G, sch_G \\\n = get_model_optimizer_and_scheduler(opt)\n\n trainer = get_trainer(opt, net_G, net_G_ema, \\\n opt_G, sch_G, None)\n\n current_epoch, current_iteration = trainer.load_checkpoint(\n opt, args.which_iter) \n net_G = trainer.net_G_ema.eval()\n\n output_dir = os.path.join(\n args.output_dir, \n 'epoch_{:05}_iteration_{:09}'.format(current_epoch, current_iteration)\n )\n\n os.makedirs(output_dir, exist_ok=True)\n image_dataset = ImageDataset(opt.data, args.input_name)\n\n control_dict, sort_rot_control, sort_exp_control = get_control(args.input_name)\n for _ in range(image_dataset.__len__()):\n with torch.no_grad():\n data = image_dataset.next_image()\n num = 10\n output_images = [] \n # rotation control\n current = control_dict['rotation_center']\n for control in sort_rot_control: \n for i in range(num):\n rotation = (control_dict[control]-current)*i/(num-1)+current\n data['target_semantics'][:, 64:70, :] = rotation[None, :, None]\n output_dict = net_G(data['source_image'].cuda(), data['target_semantics'].cuda())\n output_images.append(\n output_dict['fake_image'].cpu().clamp_(-1, 1)\n ) \n current = rotation\n\n # expression control\n current = data['target_semantics'][0, :64, 0]\n for control in sort_exp_control: \n for i in range(num):\n expression = (control_dict[control]-current)*i/(num-1)+current\n data['target_semantics'][:, :64, :] = expression[None, :, None]\n output_dict = net_G(data['source_image'].cuda(), data['target_semantics'].cuda())\n output_images.append(\n output_dict['fake_image'].cpu().clamp_(-1, 1)\n ) \n current = expression\n output_images = torch.cat(output_images, 0) \n print('write results to file {}/{}'.format(output_dir, data['name']))\n write2video('{}/{}'.format(output_dir, data['name']), output_images)\n\n"
},
{
"path": "loss/perceptual.py",
"content": "import torch\nimport torch.nn.functional as F\nimport torchvision\nfrom torch import nn\n\nfrom util.distributed import master_only_print as print\n\ndef apply_imagenet_normalization(input):\n r\"\"\"Normalize using ImageNet mean and std.\n\n Args:\n input (4D tensor NxCxHxW): The input images, assuming to be [-1, 1].\n\n Returns:\n Normalized inputs using the ImageNet normalization.\n \"\"\"\n # normalize the input back to [0, 1]\n normalized_input = (input + 1) / 2\n # normalize the input using the ImageNet mean and std\n mean = normalized_input.new_tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)\n std = normalized_input.new_tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)\n output = (normalized_input - mean) / std\n return output\n\nclass PerceptualLoss(nn.Module):\n r\"\"\"Perceptual loss initialization.\n\n Args:\n network (str) : The name of the loss network: 'vgg16' | 'vgg19'.\n layers (str or list of str) : The layers used to compute the loss.\n weights (float or list of float : The loss weights of each layer.\n criterion (str): The type of distance function: 'l1' | 'l2'.\n resize (bool) : If ``True``, resize the input images to 224x224.\n resize_mode (str): Algorithm used for resizing.\n instance_normalized (bool): If ``True``, applies instance normalization\n to the feature maps before computing the distance.\n num_scales (int): The loss will be evaluated at original size and\n this many times downsampled sizes.\n \"\"\"\n\n def __init__(self, network='vgg19', layers='relu_4_1', weights=None,\n criterion='l1', resize=False, resize_mode='bilinear',\n instance_normalized=False, num_scales=1, \n use_style_loss=False, weight_style_to_perceptual=0):\n super().__init__()\n if isinstance(layers, str):\n layers = [layers]\n if weights is None:\n weights = [1.] * len(layers)\n elif isinstance(layers, float) or isinstance(layers, int):\n weights = [weights]\n\n assert len(layers) == len(weights), \\\n 'The number of layers (%s) must be equal to ' \\\n 'the number of weights (%s).' % (len(layers), len(weights))\n if network == 'vgg19':\n self.model = _vgg19(layers)\n elif network == 'vgg16':\n self.model = _vgg16(layers)\n elif network == 'alexnet':\n self.model = _alexnet(layers)\n elif network == 'inception_v3':\n self.model = _inception_v3(layers)\n elif network == 'resnet50':\n self.model = _resnet50(layers)\n elif network == 'robust_resnet50':\n self.model = _robust_resnet50(layers)\n elif network == 'vgg_face_dag':\n self.model = _vgg_face_dag(layers)\n else:\n raise ValueError('Network %s is not recognized' % network)\n\n self.num_scales = num_scales\n self.layers = layers\n self.weights = weights\n if criterion == 'l1':\n self.criterion = nn.L1Loss()\n elif criterion == 'l2' or criterion == 'mse':\n self.criterion = nn.MSELoss()\n else:\n raise ValueError('Criterion %s is not recognized' % criterion)\n self.resize = resize\n self.resize_mode = resize_mode\n self.instance_normalized = instance_normalized\n self.use_style_loss = use_style_loss\n self.weight_style = weight_style_to_perceptual\n\n print('Perceptual loss:')\n print('\\tMode: {}'.format(network))\n\n def forward(self, inp, target, mask=None):\n r\"\"\"Perceptual loss forward.\n\n Args:\n inp (4D tensor) : Input tensor.\n target (4D tensor) : Ground truth tensor, same shape as the input.\n\n Returns:\n (scalar tensor) : The perceptual loss.\n \"\"\"\n # Perceptual loss should operate in eval mode by default.\n self.model.eval()\n inp, target = \\\n apply_imagenet_normalization(inp), \\\n apply_imagenet_normalization(target)\n if self.resize:\n inp = F.interpolate(\n inp, mode=self.resize_mode, size=(224, 224),\n align_corners=False)\n target = F.interpolate(\n target, mode=self.resize_mode, size=(224, 224),\n align_corners=False)\n\n # Evaluate perceptual loss at each scale.\n loss = 0\n style_loss=0\n for scale in range(self.num_scales):\n input_features, target_features = \\\n self.model(inp), self.model(target)\n for layer, weight in zip(self.layers, self.weights):\n # Example per-layer VGG19 loss values after applying\n # [0.03125, 0.0625, 0.125, 0.25, 1.0] weighting.\n # relu_1_1, 0.014698\n # relu_2_1, 0.085817\n # relu_3_1, 0.349977\n # relu_4_1, 0.544188\n # relu_5_1, 0.906261\n input_feature = input_features[layer]\n target_feature = target_features[layer].detach()\n if self.instance_normalized:\n input_feature = F.instance_norm(input_feature)\n target_feature = F.instance_norm(target_feature)\n\n if mask is not None:\n mask_ = F.interpolate(mask, input_feature.shape[2:],\n mode='bilinear',\n align_corners=False)\n input_feature = input_feature * mask_ \n target_feature = target_feature * mask_ \n # print('mask',mask_.shape) \n\n\n loss += weight * self.criterion(input_feature,\n target_feature)\n if self.use_style_loss and scale==0:\n style_loss += self.criterion(self.compute_gram(input_feature),\n self.compute_gram(target_feature))\n\n # Downsample the input and target.\n if scale != self.num_scales - 1:\n inp = F.interpolate(\n inp, mode=self.resize_mode, scale_factor=0.5,\n align_corners=False, recompute_scale_factor=True)\n target = F.interpolate(\n target, mode=self.resize_mode, scale_factor=0.5,\n align_corners=False, recompute_scale_factor=True)\n\n if self.use_style_loss:\n return loss + style_loss*self.weight_style\n else:\n return loss\n\n\n def compute_gram(self, x):\n b, ch, h, w = x.size()\n f = x.view(b, ch, w * h)\n f_T = f.transpose(1, 2)\n G = f.bmm(f_T) / (h * w * ch)\n return G\n\nclass _PerceptualNetwork(nn.Module):\n r\"\"\"The network that extracts features to compute the perceptual loss.\n\n Args:\n network (nn.Sequential) : The network that extracts features.\n layer_name_mapping (dict) : The dictionary that\n maps a layer's index to its name.\n layers (list of str): The list of layer names that we are using.\n \"\"\"\n\n def __init__(self, network, layer_name_mapping, layers):\n super().__init__()\n assert isinstance(network, nn.Sequential), \\\n 'The network needs to be of type \"nn.Sequential\".'\n self.network = network\n self.layer_name_mapping = layer_name_mapping\n self.layers = layers\n for param in self.parameters():\n param.requires_grad = False\n\n def forward(self, x):\n r\"\"\"Extract perceptual features.\"\"\"\n output = {}\n for i, layer in enumerate(self.network):\n x = layer(x)\n layer_name = self.layer_name_mapping.get(i, None)\n if layer_name in self.layers:\n # If the current layer is used by the perceptual loss.\n output[layer_name] = x\n return output\n\n\ndef _vgg19(layers):\n r\"\"\"Get vgg19 layers\"\"\"\n network = torchvision.models.vgg19(pretrained=True).features\n layer_name_mapping = {1: 'relu_1_1',\n 3: 'relu_1_2',\n 6: 'relu_2_1',\n 8: 'relu_2_2',\n 11: 'relu_3_1',\n 13: 'relu_3_2',\n 15: 'relu_3_3',\n 17: 'relu_3_4',\n 20: 'relu_4_1',\n 22: 'relu_4_2',\n 24: 'relu_4_3',\n 26: 'relu_4_4',\n 29: 'relu_5_1'}\n return _PerceptualNetwork(network, layer_name_mapping, layers)\n\n\ndef _vgg16(layers):\n r\"\"\"Get vgg16 layers\"\"\"\n network = torchvision.models.vgg16(pretrained=True).features\n layer_name_mapping = {1: 'relu_1_1',\n 3: 'relu_1_2',\n 6: 'relu_2_1',\n 8: 'relu_2_2',\n 11: 'relu_3_1',\n 13: 'relu_3_2',\n 15: 'relu_3_3',\n 18: 'relu_4_1',\n 20: 'relu_4_2',\n 22: 'relu_4_3',\n 25: 'relu_5_1'}\n return _PerceptualNetwork(network, layer_name_mapping, layers)\n\n\ndef _alexnet(layers):\n r\"\"\"Get alexnet layers\"\"\"\n network = torchvision.models.alexnet(pretrained=True).features\n layer_name_mapping = {0: 'conv_1',\n 1: 'relu_1',\n 3: 'conv_2',\n 4: 'relu_2',\n 6: 'conv_3',\n 7: 'relu_3',\n 8: 'conv_4',\n 9: 'relu_4',\n 10: 'conv_5',\n 11: 'relu_5'}\n return _PerceptualNetwork(network, layer_name_mapping, layers)\n\n\ndef _inception_v3(layers):\n r\"\"\"Get inception v3 layers\"\"\"\n inception = torchvision.models.inception_v3(pretrained=True)\n network = nn.Sequential(inception.Conv2d_1a_3x3,\n inception.Conv2d_2a_3x3,\n inception.Conv2d_2b_3x3,\n nn.MaxPool2d(kernel_size=3, stride=2),\n inception.Conv2d_3b_1x1,\n inception.Conv2d_4a_3x3,\n nn.MaxPool2d(kernel_size=3, stride=2),\n inception.Mixed_5b,\n inception.Mixed_5c,\n inception.Mixed_5d,\n inception.Mixed_6a,\n inception.Mixed_6b,\n inception.Mixed_6c,\n inception.Mixed_6d,\n inception.Mixed_6e,\n inception.Mixed_7a,\n inception.Mixed_7b,\n inception.Mixed_7c,\n nn.AdaptiveAvgPool2d(output_size=(1, 1)))\n layer_name_mapping = {3: 'pool_1',\n 6: 'pool_2',\n 14: 'mixed_6e',\n 18: 'pool_3'}\n return _PerceptualNetwork(network, layer_name_mapping, layers)\n\n\ndef _resnet50(layers):\n r\"\"\"Get resnet50 layers\"\"\"\n resnet50 = torchvision.models.resnet50(pretrained=True)\n network = nn.Sequential(resnet50.conv1,\n resnet50.bn1,\n resnet50.relu,\n resnet50.maxpool,\n resnet50.layer1,\n resnet50.layer2,\n resnet50.layer3,\n resnet50.layer4,\n resnet50.avgpool)\n layer_name_mapping = {4: 'layer_1',\n 5: 'layer_2',\n 6: 'layer_3',\n 7: 'layer_4'}\n return _PerceptualNetwork(network, layer_name_mapping, layers)\n\n\ndef _robust_resnet50(layers):\n r\"\"\"Get robust resnet50 layers\"\"\"\n resnet50 = torchvision.models.resnet50(pretrained=False)\n state_dict = torch.utils.model_zoo.load_url(\n 'http://andrewilyas.com/ImageNet.pt')\n new_state_dict = {}\n for k, v in state_dict['model'].items():\n if k.startswith('module.model.'):\n new_state_dict[k[13:]] = v\n resnet50.load_state_dict(new_state_dict)\n network = nn.Sequential(resnet50.conv1,\n resnet50.bn1,\n resnet50.relu,\n resnet50.maxpool,\n resnet50.layer1,\n resnet50.layer2,\n resnet50.layer3,\n resnet50.layer4,\n resnet50.avgpool)\n layer_name_mapping = {4: 'layer_1',\n 5: 'layer_2',\n 6: 'layer_3',\n 7: 'layer_4'}\n return _PerceptualNetwork(network, layer_name_mapping, layers)\n\n\ndef _vgg_face_dag(layers):\n r\"\"\"Get vgg face layers\"\"\"\n network = torchvision.models.vgg16(num_classes=2622)\n state_dict = torch.utils.model_zoo.load_url(\n 'http://www.robots.ox.ac.uk/~albanie/models/pytorch-mcn/'\n 'vgg_face_dag.pth')\n feature_layer_name_mapping = {\n 0: 'conv1_1',\n 2: 'conv1_2',\n 5: 'conv2_1',\n 7: 'conv2_2',\n 10: 'conv3_1',\n 12: 'conv3_2',\n 14: 'conv3_3',\n 17: 'conv4_1',\n 19: 'conv4_2',\n 21: 'conv4_3',\n 24: 'conv5_1',\n 26: 'conv5_2',\n 28: 'conv5_3'}\n new_state_dict = {}\n for k, v in feature_layer_name_mapping.items():\n new_state_dict['features.' + str(k) + '.weight'] =\\\n state_dict[v + '.weight']\n new_state_dict['features.' + str(k) + '.bias'] = \\\n state_dict[v + '.bias']\n\n classifier_layer_name_mapping = {\n 0: 'fc6',\n 3: 'fc7',\n 6: 'fc8'}\n for k, v in classifier_layer_name_mapping.items():\n new_state_dict['classifier.' + str(k) + '.weight'] = \\\n state_dict[v + '.weight']\n new_state_dict['classifier.' + str(k) + '.bias'] = \\\n state_dict[v + '.bias']\n\n network.load_state_dict(new_state_dict)\n\n class Flatten(nn.Module):\n r\"\"\"Flatten the tensor\"\"\"\n\n def forward(self, x):\n r\"\"\"Flatten it\"\"\"\n return x.view(x.shape[0], -1)\n\n layer_name_mapping = {\n 1: 'avgpool',\n 3: 'fc6',\n 4: 'relu_6',\n 6: 'fc7',\n 7: 'relu_7',\n 9: 'fc8'}\n seq_layers = [network.features, network.avgpool, Flatten()]\n for i in range(7):\n seq_layers += [network.classifier[i]]\n network = nn.Sequential(*seq_layers)\n return _PerceptualNetwork(network, layer_name_mapping, layers)\n"
},
{
"path": "requirements.txt",
"content": "absl-py==0.13.0\nbackcall==0.2.0\ncachetools==4.2.2\ncertifi==2021.5.30\ncharset-normalizer==2.0.6\ncycler==0.10.0\ndataclasses==0.8\ndecorator==4.4.2\nfilelock==3.0.12\ngdown==3.13.1\ngoogle-auth==1.35.0\ngoogle-auth-oauthlib==0.4.6\ngrpcio==1.40.0\nidna==3.2\nimageio==2.9.0\nimportlib-metadata==4.8.1\nipython==7.16.1\nipython-genutils==0.2.0\njedi==0.18.0\nkiwisolver==1.3.1\nlmdb==1.2.1\nMarkdown==3.3.4\nmatplotlib==3.3.4\nmkl-fft==1.3.0\nmkl-random==1.1.1\nmkl-service==2.3.0\nnetworkx==2.5.1\nnumpy==1.19.2\noauthlib==3.1.1\nolefile==0.46\nopencv-python==4.5.3.56\nparso==0.8.2\npexpect==4.8.0\npickleshare==0.7.5\nPillow==8.3.1\npip==21.2.2\nprompt-toolkit==3.0.20\nprotobuf==3.18.0\nptyprocess==0.7.0\npyasn1==0.4.8\npyasn1-modules==0.2.8\nPygments==2.10.0\npyparsing==2.4.7\nPySocks==1.7.1\npython-dateutil==2.8.2\nPyWavelets==1.1.1\nPyYAML==5.4.1\nrequests==2.26.0\nrequests-oauthlib==1.3.0\nrsa==4.7.2\nscikit-image==0.17.2\nscipy==1.5.4\nsetuptools==58.0.4\nsix==1.16.0\ntensorboard==2.6.0\ntensorboard-data-server==0.6.1\ntensorboard-plugin-wit==1.8.0\ntifffile==2020.9.3\ntorch==1.7.1\ntorchvision==0.8.2\ntqdm==4.62.2\ntraitlets==4.3.3\ntyping-extensions==3.10.0.2\nurllib3==1.26.6\nwcwidth==0.2.5\nWerkzeug==2.0.1\nwheel==0.37.0\nzipp==3.5.0\n"
},
{
"path": "scripts/coeff_detector.py",
"content": "import os\nimport glob\nimport numpy as np\nfrom os import makedirs, name\nfrom PIL import Image\nfrom tqdm import tqdm\n\nimport torch\nimport torch.nn as nn\n\nfrom options.inference_options import InferenceOptions\nfrom models import create_model\nfrom util.preprocess import align_img\nfrom util.load_mats import load_lm3d\nfrom extract_kp_videos import KeypointExtractor\n\n\nclass CoeffDetector(nn.Module):\n def __init__(self, opt):\n super().__init__()\n\n self.model = create_model(opt)\n self.model.setup(opt)\n self.model.device = 'cuda'\n self.model.parallelize()\n self.model.eval()\n\n self.lm3d_std = load_lm3d(opt.bfm_folder) \n\n def forward(self, img, lm):\n \n img, trans_params = self.image_transform(img, lm)\n\n data_input = { \n 'imgs': img[None],\n } \n self.model.set_input(data_input) \n self.model.test()\n pred_coeff = {key:self.model.pred_coeffs_dict[key].cpu().numpy() for key in self.model.pred_coeffs_dict}\n pred_coeff = np.concatenate([\n pred_coeff['id'], \n pred_coeff['exp'], \n pred_coeff['tex'], \n pred_coeff['angle'],\n pred_coeff['gamma'],\n pred_coeff['trans'],\n trans_params[None],\n ], 1)\n \n return {'coeff_3dmm':pred_coeff, \n 'crop_img': Image.fromarray((img.cpu().permute(1, 2, 0).numpy()*255).astype(np.uint8))}\n\n def image_transform(self, images, lm):\n \"\"\"\n param:\n images: -- PIL image \n lm: -- numpy array\n \"\"\"\n W,H = images.size\n if np.mean(lm) == -1:\n lm = (self.lm3d_std[:, :2]+1)/2.\n lm = np.concatenate(\n [lm[:, :1]*W, lm[:, 1:2]*H], 1\n )\n else:\n lm[:, -1] = H - 1 - lm[:, -1]\n\n trans_params, img, lm, _ = align_img(images, lm, self.lm3d_std) \n img = torch.tensor(np.array(img)/255., dtype=torch.float32).permute(2, 0, 1)\n trans_params = np.array([float(item) for item in np.hsplit(trans_params, 5)])\n trans_params = torch.tensor(trans_params.astype(np.float32))\n return img, trans_params \n\ndef get_data_path(root, keypoint_root):\n filenames = list()\n keypoint_filenames = list()\n\n IMAGE_EXTENSIONS_LOWERCASE = {'jpg', 'png', 'jpeg', 'webp'}\n IMAGE_EXTENSIONS = IMAGE_EXTENSIONS_LOWERCASE.union({f.upper() for f in IMAGE_EXTENSIONS_LOWERCASE})\n extensions = IMAGE_EXTENSIONS\n\n for ext in extensions:\n filenames += glob.glob(f'{root}/*.{ext}', recursive=True)\n filenames = sorted(filenames)\n for filename in filenames:\n name = os.path.splitext(os.path.basename(filename))[0]\n keypoint_filenames.append(\n os.path.join(keypoint_root, name + '.txt')\n )\n return filenames, keypoint_filenames\n\n\nif __name__ == \"__main__\":\n opt = InferenceOptions().parse() \n coeff_detector = CoeffDetector(opt)\n kp_extractor = KeypointExtractor()\n image_names, keypoint_names = get_data_path(opt.input_dir, opt.keypoint_dir)\n makedirs(opt.keypoint_dir, exist_ok=True)\n makedirs(opt.output_dir, exist_ok=True)\n\n for image_name, keypoint_name in tqdm(zip(image_names, keypoint_names)):\n image = Image.open(image_name)\n if not os.path.isfile(keypoint_name):\n lm = kp_extractor.extract_keypoint(image, keypoint_name)\n else:\n lm = np.loadtxt(keypoint_name).astype(np.float32)\n lm = lm.reshape([-1, 2]) \n predicted = coeff_detector(image, lm)\n name = os.path.splitext(os.path.basename(image_name))[0]\n np.savetxt(\n \"{}/{}_3dmm_coeff.txt\".format(opt.output_dir, name), \n predicted['coeff_3dmm'].reshape(-1))\n\n \n\n\n\n "
},
{
"path": "scripts/download_demo_dataset.sh",
"content": "gdown https://drive.google.com/uc?id=1ruuLw5-0fpm6EREexPn3I_UQPmkrBoq9\nunzip -x ./vox_lmdb_demo.zip\nmkdir ./dataset\nmv vox_lmdb_demo ./dataset\n"
},
{
"path": "scripts/download_weights.sh",
"content": "gdown https://drive.google.com/uc?id=1-0xOf6g58OmtKtEWJlU3VlnfRqPN9Uq7\nunzip -x ./face.zip\nmkdir ./result\nmv face ./result\nrm face.zip\n"
},
{
"path": "scripts/extract_kp_videos.py",
"content": "import os\nimport cv2\nimport time\nimport glob\nimport argparse\nimport face_alignment\nimport numpy as np\nfrom PIL import Image\nfrom tqdm import tqdm\nfrom itertools import cycle\n\nfrom torch.multiprocessing import Pool, Process, set_start_method\n\nclass KeypointExtractor():\n def __init__(self):\n self.detector = face_alignment.FaceAlignment(face_alignment.LandmarksType._2D) \n\n def extract_keypoint(self, images, name=None):\n if isinstance(images, list):\n keypoints = []\n for image in images:\n current_kp = self.extract_keypoint(image)\n if np.mean(current_kp) == -1 and keypoints:\n keypoints.append(keypoints[-1])\n else:\n keypoints.append(current_kp[None])\n\n keypoints = np.concatenate(keypoints, 0)\n np.savetxt(os.path.splitext(name)[0]+'.txt', keypoints.reshape(-1))\n return keypoints\n else:\n while True:\n try:\n keypoints = self.detector.get_landmarks_from_image(np.array(images))[0]\n break\n except RuntimeError as e:\n if str(e).startswith('CUDA'):\n print(\"Warning: out of memory, sleep for 1s\")\n time.sleep(1)\n else:\n print(e)\n break \n except TypeError:\n print('No face detected in this image')\n shape = [68, 2]\n keypoints = -1. * np.ones(shape) \n break\n if name is not None:\n np.savetxt(os.path.splitext(name)[0]+'.txt', keypoints.reshape(-1))\n return keypoints\n\ndef read_video(filename):\n frames = []\n cap = cv2.VideoCapture(filename)\n while cap.isOpened():\n ret, frame = cap.read()\n if ret:\n frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)\n frame = Image.fromarray(frame)\n frames.append(frame)\n else:\n break\n cap.release()\n return frames\n\ndef run(data):\n filename, opt, device = data\n os.environ['CUDA_VISIBLE_DEVICES'] = device\n kp_extractor = KeypointExtractor()\n images = read_video(filename)\n name = filename.split('/')[-2:]\n os.makedirs(os.path.join(opt.output_dir, name[-2]), exist_ok=True)\n kp_extractor.extract_keypoint(\n images, \n name=os.path.join(opt.output_dir, name[-2], name[-1])\n )\n\nif __name__ == '__main__':\n set_start_method('spawn')\n parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)\n parser.add_argument('--input_dir', type=str, help='the folder of the input files')\n parser.add_argument('--output_dir', type=str, help='the folder of the output files')\n parser.add_argument('--device_ids', type=str, default='0,1')\n parser.add_argument('--workers', type=int, default=4)\n\n opt = parser.parse_args()\n filenames = list()\n VIDEO_EXTENSIONS_LOWERCASE = {'mp4'}\n VIDEO_EXTENSIONS = VIDEO_EXTENSIONS_LOWERCASE.union({f.upper() for f in VIDEO_EXTENSIONS_LOWERCASE})\n extensions = VIDEO_EXTENSIONS\n for ext in extensions:\n filenames = sorted(glob.glob(f'{opt.input_dir}/**/*.{ext}'))\n print('Total number of videos:', len(filenames))\n pool = Pool(opt.workers)\n args_list = cycle([opt])\n device_ids = opt.device_ids.split(\",\")\n device_ids = cycle(device_ids)\n for data in tqdm(pool.imap_unordered(run, zip(filenames, args_list, device_ids))):\n None\n"
},
{
"path": "scripts/face_recon_images.py",
"content": "import os\nimport glob\nimport numpy as np\nfrom PIL import Image\nfrom tqdm import tqdm\nfrom scipy.io import savemat\n\nimport torch \n\nfrom models import create_model\nfrom options.inference_options import InferenceOptions\nfrom util.preprocess import align_img\nfrom util.load_mats import load_lm3d\nfrom util.util import tensor2im, save_image\n\n\ndef get_data_path(root, keypoint_root):\n filenames = list()\n keypoint_filenames = list()\n\n IMAGE_EXTENSIONS_LOWERCASE = {'jpg', 'png', 'jpeg', 'webp'}\n IMAGE_EXTENSIONS = IMAGE_EXTENSIONS_LOWERCASE.union({f.upper() for f in IMAGE_EXTENSIONS_LOWERCASE})\n extensions = IMAGE_EXTENSIONS\n\n for ext in extensions:\n filenames += glob.glob(f'{root}/*.{ext}', recursive=True)\n filenames = sorted(filenames)\n for filename in filenames:\n name = os.path.splitext(os.path.basename(filename))[0]\n keypoint_filenames.append(\n os.path.join(keypoint_root, name + '.txt')\n )\n return filenames, keypoint_filenames\n\n\nclass ImagePathDataset(torch.utils.data.Dataset):\n def __init__(self, filenames, txt_filenames, bfm_folder):\n self.filenames = filenames\n self.txt_filenames = txt_filenames\n self.lm3d_std = load_lm3d(bfm_folder) \n\n def __len__(self):\n return len(self.filenames)\n\n def __getitem__(self, i):\n filename = self.filenames[i]\n txt_filename = self.txt_filenames[i]\n imgs, _, trans_params = self.read_data(filename, txt_filename)\n return {\n 'imgs':imgs,\n 'trans_param':trans_params,\n 'filename': filename\n }\n\n def image_transform(self, images, lm):\n W,H = images.size\n if np.mean(lm) == -1:\n lm = (self.lm3d_std[:, :2]+1)/2.\n lm = np.concatenate(\n [lm[:, :1]*W, lm[:, 1:2]*H], 1\n )\n else:\n lm[:, -1] = H - 1 - lm[:, -1]\n\n trans_params, img, lm, _ = align_img(images, lm, self.lm3d_std) \n img = torch.tensor(np.array(img)/255., dtype=torch.float32).permute(2, 0, 1)\n lm = torch.tensor(lm)\n trans_params = np.array([float(item) for item in np.hsplit(trans_params, 5)])\n trans_params = torch.tensor(trans_params.astype(np.float32))\n return img, lm, trans_params \n\n def read_data(self, filename, txt_filename):\n images = Image.open(filename).convert('RGB')\n lm = np.loadtxt(txt_filename).astype(np.float32)\n lm = lm.reshape([-1, 2]) \n imgs, lms, trans_params = self.image_transform(images, lm)\n return imgs, lms, trans_params\n\n\ndef main(opt, model):\n import torch.multiprocessing\n torch.multiprocessing.set_sharing_strategy('file_system')\n filenames, keypoint_filenames = get_data_path(opt.input_dir, opt.keypoint_dir)\n \n dataset = ImagePathDataset(filenames, keypoint_filenames, opt.bfm_folder)\n dataloader = torch.utils.data.DataLoader(\n dataset,\n batch_size=opt.inference_batch_size,\n shuffle=False,\n drop_last=False,\n num_workers=8,\n ) \n pred_coeffs, pred_trans_params = [], []\n print('nums of images:', dataset.__len__())\n for iteration, data in tqdm(enumerate(dataloader)):\n data_input = { \n 'imgs': data['imgs'],\n }\n \n model.set_input(data_input) \n model.test()\n pred_coeff = {key:model.pred_coeffs_dict[key].cpu().numpy() for key in model.pred_coeffs_dict}\n pred_coeff = np.concatenate([\n pred_coeff['id'], \n pred_coeff['exp'], \n pred_coeff['tex'], \n pred_coeff['angle'],\n pred_coeff['gamma'],\n pred_coeff['trans']], 1)\n pred_coeffs.append(pred_coeff) \n trans_param = data['trans_param'].cpu().numpy()\n pred_trans_params.append(trans_param)\n if opt.save_split_files:\n for index, filename in enumerate(data['filename']):\n basename = os.path.splitext(os.path.basename(filename))[0]\n output_path = os.path.join(opt.output_dir, basename+'.mat')\n savemat(\n output_path, \n {'coeff':pred_coeff[index], \n 'transform_params':trans_param[index]}\n )\n # visuals = model.get_current_visuals() # get image results\n # for name in visuals:\n # images = visuals[name]\n # for i in range(images.shape[0]):\n # image_numpy = tensor2im(images[i])\n # save_image(image_numpy, os.path.basename(data['filename'][i])+'.png') \n\n pred_coeffs = np.concatenate(pred_coeffs, 0)\n pred_trans_params = np.concatenate(pred_trans_params, 0)\n savemat(os.path.join(opt.output_dir, 'ffhq.mat'), {'coeff':pred_coeffs, 'transform_params':pred_trans_params})\n\n\nif __name__ == '__main__':\n opt = InferenceOptions().parse() # get test options\n model = create_model(opt)\n model.setup(opt)\n model.device = 'cuda:0'\n model.parallelize()\n model.eval()\n lm3d_std = load_lm3d(opt.bfm_folder) \n main(opt, model)\n\n\n"
},
{
"path": "scripts/face_recon_videos.py",
"content": "import os\nimport cv2\nimport glob\nimport numpy as np\nfrom PIL import Image\nfrom tqdm import tqdm\nfrom scipy.io import savemat\n\nimport torch \n\nfrom models import create_model\nfrom options.inference_options import InferenceOptions\nfrom util.preprocess import align_img\nfrom util.load_mats import load_lm3d\nfrom util.util import mkdirs, tensor2im, save_image\n\n\ndef get_data_path(root, keypoint_root):\n filenames = list()\n keypoint_filenames = list()\n\n VIDEO_EXTENSIONS_LOWERCASE = {'mp4'}\n VIDEO_EXTENSIONS = VIDEO_EXTENSIONS_LOWERCASE.union({f.upper() for f in VIDEO_EXTENSIONS_LOWERCASE})\n extensions = VIDEO_EXTENSIONS\n\n for ext in extensions:\n filenames += glob.glob(f'{root}/**/*.{ext}', recursive=True)\n filenames = sorted(filenames)\n keypoint_filenames = sorted(glob.glob(f'{keypoint_root}/**/*.txt', recursive=True))\n assert len(filenames) == len(keypoint_filenames)\n\n return filenames, keypoint_filenames\n\nclass VideoPathDataset(torch.utils.data.Dataset):\n def __init__(self, filenames, txt_filenames, bfm_folder):\n self.filenames = filenames\n self.txt_filenames = txt_filenames\n self.lm3d_std = load_lm3d(bfm_folder) \n\n def __len__(self):\n return len(self.filenames)\n\n def __getitem__(self, index):\n filename = self.filenames[index]\n txt_filename = self.txt_filenames[index]\n frames = self.read_video(filename)\n lm = np.loadtxt(txt_filename).astype(np.float32)\n lm = lm.reshape([len(frames), -1, 2]) \n out_images, out_trans_params = list(), list()\n for i in range(len(frames)):\n out_img, _, out_trans_param \\\n = self.image_transform(frames[i], lm[i])\n out_images.append(out_img[None])\n out_trans_params.append(out_trans_param[None])\n return {\n 'imgs': torch.cat(out_images, 0),\n 'trans_param':torch.cat(out_trans_params, 0),\n 'filename': filename\n }\n \n def read_video(self, filename):\n frames = list()\n cap = cv2.VideoCapture(filename)\n while cap.isOpened():\n ret, frame = cap.read()\n if ret:\n frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)\n frame = Image.fromarray(frame)\n frames.append(frame)\n else:\n break\n cap.release()\n return frames\n\n def image_transform(self, images, lm):\n W,H = images.size\n if np.mean(lm) == -1:\n lm = (self.lm3d_std[:, :2]+1)/2.\n lm = np.concatenate(\n [lm[:, :1]*W, lm[:, 1:2]*H], 1\n )\n else:\n lm[:, -1] = H - 1 - lm[:, -1]\n\n trans_params, img, lm, _ = align_img(images, lm, self.lm3d_std) \n img = torch.tensor(np.array(img)/255., dtype=torch.float32).permute(2, 0, 1)\n lm = torch.tensor(lm)\n trans_params = np.array([float(item) for item in np.hsplit(trans_params, 5)])\n trans_params = torch.tensor(trans_params.astype(np.float32))\n return img, lm, trans_params \n\ndef main(opt, model):\n import torch.multiprocessing\n torch.multiprocessing.set_sharing_strategy('file_system')\n filenames, keypoint_filenames = get_data_path(opt.input_dir, opt.keypoint_dir)\n dataset = VideoPathDataset(filenames, keypoint_filenames, opt.bfm_folder)\n dataloader = torch.utils.data.DataLoader(\n dataset,\n batch_size=1, # can noly set to one here!\n shuffle=False,\n drop_last=False,\n num_workers=8,\n ) \n batch_size = opt.inference_batch_size\n for data in tqdm(dataloader):\n num_batch = data['imgs'][0].shape[0] // batch_size + 1\n pred_coeffs = list()\n for index in range(num_batch):\n data_input = { \n 'imgs': data['imgs'][0,index*batch_size:(index+1)*batch_size],\n }\n model.set_input(data_input) \n model.test()\n pred_coeff = {key:model.pred_coeffs_dict[key].cpu().numpy() for key in model.pred_coeffs_dict}\n pred_coeff = np.concatenate([\n pred_coeff['id'], \n pred_coeff['exp'], \n pred_coeff['tex'], \n pred_coeff['angle'],\n pred_coeff['gamma'],\n pred_coeff['trans']], 1)\n pred_coeffs.append(pred_coeff) \n visuals = model.get_current_visuals() # get image results\n if False: # debug\n for name in visuals:\n images = visuals[name]\n for i in range(images.shape[0]):\n image_numpy = tensor2im(images[i])\n save_image(\n image_numpy, \n os.path.join(\n opt.output_dir,\n os.path.basename(data['filename'][0])+str(i).zfill(5)+'.jpg')\n )\n exit()\n\n pred_coeffs = np.concatenate(pred_coeffs, 0)\n pred_trans_params = data['trans_param'][0].cpu().numpy()\n name = data['filename'][0].split('/')[-2:]\n name[-1] = os.path.splitext(name[-1])[0] + '.mat'\n os.makedirs(os.path.join(opt.output_dir, name[-2]), exist_ok=True)\n savemat(\n os.path.join(opt.output_dir, name[-2], name[-1]), \n {'coeff':pred_coeffs, 'transform_params':pred_trans_params}\n )\n\nif __name__ == '__main__':\n opt = InferenceOptions().parse() # get test options\n model = create_model(opt)\n model.setup(opt)\n model.device = 'cuda:0'\n model.parallelize()\n model.eval()\n\n main(opt, model)\n\n\n"
},
{
"path": "scripts/inference_options.py",
"content": "from .base_options import BaseOptions\n\n\nclass InferenceOptions(BaseOptions):\n \"\"\"This class includes test options.\n\n It also includes shared options defined in BaseOptions.\n \"\"\"\n\n def initialize(self, parser):\n parser = BaseOptions.initialize(self, parser) # define shared options\n parser.add_argument('--phase', type=str, default='test', help='train, val, test, etc')\n parser.add_argument('--dataset_mode', type=str, default=None, help='chooses how datasets are loaded. [None | flist]')\n\n parser.add_argument('--input_dir', type=str, help='the folder of the input files')\n parser.add_argument('--keypoint_dir', type=str, help='the folder of the keypoint files')\n parser.add_argument('--output_dir', type=str, default='mp4', help='the output dir to save the extracted coefficients')\n parser.add_argument('--save_split_files', action='store_true', help='save split files or not')\n parser.add_argument('--inference_batch_size', type=int, default=8)\n \n # Dropout and Batchnorm has different behavior during training and test.\n self.isTrain = False\n return parser\n"
},
{
"path": "scripts/prepare_vox_lmdb.py",
"content": "import os\nimport cv2\nimport lmdb\nimport argparse\nimport multiprocessing\nimport numpy as np\n\nfrom glob import glob\nfrom io import BytesIO\nfrom tqdm import tqdm\nfrom PIL import Image\nfrom scipy.io import loadmat\nfrom torchvision.transforms import functional as trans_fn\n\ndef format_for_lmdb(*args):\n key_parts = []\n for arg in args:\n if isinstance(arg, int):\n arg = str(arg).zfill(7)\n key_parts.append(arg)\n return '-'.join(key_parts).encode('utf-8')\n\nclass Resizer:\n def __init__(self, size, kp_root, coeff_3dmm_root, img_format):\n self.size = size\n self.kp_root = kp_root\n self.coeff_3dmm_root = coeff_3dmm_root\n self.img_format = img_format\n\n def get_resized_bytes(self, img, img_format='jpeg'):\n img = trans_fn.resize(img, (self.size, self.size), interpolation=Image.BICUBIC)\n buf = BytesIO()\n img.save(buf, format=img_format)\n img_bytes = buf.getvalue()\n return img_bytes\n\n def prepare(self, filename):\n frames = {'img':[], 'kp':None, 'coeff_3dmm':None}\n cap = cv2.VideoCapture(filename)\n while cap.isOpened():\n ret, frame = cap.read()\n if ret:\n frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)\n img_pil = Image.fromarray(frame)\n img_bytes = self.get_resized_bytes(img_pil, self.img_format)\n frames['img'].append(img_bytes)\n else:\n break\n cap.release()\n video_name = os.path.splitext(os.path.basename(filename))[0]\n keypoint_byte = get_others(self.kp_root, video_name, 'keypoint')\n coeff_3dmm_byte = get_others(self.coeff_3dmm_root, video_name, 'coeff_3dmm')\n frames['kp'] = keypoint_byte\n frames['coeff_3dmm'] = coeff_3dmm_byte\n return frames\n\n def __call__(self, index_filename):\n index, filename = index_filename\n result = self.prepare(filename)\n return index, result, filename\n\ndef get_others(root, video_name, data_type):\n if root is None:\n return\n else:\n assert data_type in ('keypoint', 'coeff_3dmm')\n if os.path.isfile(os.path.join(root, 'train', video_name+'.mat')):\n file_path = os.path.join(root, 'train', video_name+'.mat')\n else:\n file_path = os.path.join(root, 'test', video_name+'.mat')\n \n if data_type == 'keypoint':\n return_byte = convert_kp(file_path)\n else:\n return_byte = convert_3dmm(file_path)\n return return_byte\n\ndef convert_kp(file_path):\n file_mat = loadmat(file_path)\n kp_byte = file_mat['landmark'].tobytes()\n return kp_byte\n\ndef convert_3dmm(file_path):\n file_mat = loadmat(file_path)\n coeff_3dmm = file_mat['coeff']\n crop_param = file_mat['transform_params']\n _, _, ratio, t0, t1 = np.hsplit(crop_param.astype(np.float32), 5)\n crop_param = np.concatenate([ratio, t0, t1], 1)\n coeff_3dmm_cat = np.concatenate([coeff_3dmm, crop_param], 1) \n coeff_3dmm_byte = coeff_3dmm_cat.tobytes()\n return coeff_3dmm_byte\n\n\ndef prepare_data(path, keypoint_path, coeff_3dmm_path, out, n_worker, sizes, chunksize, img_format):\n filenames = list()\n VIDEO_EXTENSIONS_LOWERCASE = {'mp4'}\n VIDEO_EXTENSIONS = VIDEO_EXTENSIONS_LOWERCASE.union({f.upper() for f in VIDEO_EXTENSIONS_LOWERCASE})\n extensions = VIDEO_EXTENSIONS\n for ext in extensions:\n filenames += glob(f'{path}/**/*.{ext}', recursive=True)\n train_video, test_video = [], []\n for item in filenames:\n if \"/train/\" in item:\n train_video.append(item)\n else:\n test_video.append(item)\n print(len(train_video), len(test_video))\n with open(os.path.join(out, 'train_list.txt'),'w') as f:\n for item in train_video:\n item = os.path.splitext(os.path.basename(item))[0]\n f.write(item + '\\n')\n\n with open(os.path.join(out, 'test_list.txt'),'w') as f:\n for item in test_video:\n item = os.path.splitext(os.path.basename(item))[0]\n f.write(item + '\\n') \n\n\n filenames = sorted(filenames)\n total = len(filenames)\n os.makedirs(out, exist_ok=True)\n for size in sizes:\n lmdb_path = os.path.join(out, str(size))\n with lmdb.open(lmdb_path, map_size=1024 ** 4, readahead=False) as env:\n with env.begin(write=True) as txn:\n txn.put(format_for_lmdb('length'), format_for_lmdb(total))\n resizer = Resizer(size, keypoint_path, coeff_3dmm_path, img_format)\n with multiprocessing.Pool(n_worker) as pool:\n for idx, result, filename in tqdm(\n pool.imap_unordered(resizer, enumerate(filenames), chunksize=chunksize),\n total=total):\n filename = os.path.basename(filename)\n video_name = os.path.splitext(filename)[0]\n txn.put(format_for_lmdb(video_name, 'length'), format_for_lmdb(len(result['img'])))\n\n for frame_idx, frame in enumerate(result['img']):\n txn.put(format_for_lmdb(video_name, frame_idx), frame)\n\n if result['kp']:\n txn.put(format_for_lmdb(video_name, 'keypoint'), result['kp'])\n if result['coeff_3dmm']:\n txn.put(format_for_lmdb(video_name, 'coeff_3dmm'), result['coeff_3dmm'])\n\n\nif __name__ == '__main__':\n parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)\n parser.add_argument('--path', type=str, help='a path to input directiory')\n parser.add_argument('--keypoint_path', type=str, help='a path to output directory', default=None)\n parser.add_argument('--coeff_3dmm_path', type=str, help='a path to output directory', default=None)\n parser.add_argument('--out', type=str, help='a path to output directory')\n parser.add_argument('--sizes', type=int, nargs='+', default=(256,))\n parser.add_argument('--n_worker', type=int, help='number of worker processes', default=8)\n parser.add_argument('--chunksize', type=int, help='approximate chunksize for each worker', default=10)\n parser.add_argument('--img_format', type=str, default='jpeg')\n args = parser.parse_args()\n prepare_data(**vars(args))"
},
{
"path": "third_part/PerceptualSimilarity/models/__init__.py",
"content": ""
},
{
"path": "third_part/PerceptualSimilarity/models/base_model.py",
"content": "import os\nimport torch\nfrom torch.autograd import Variable\nfrom pdb import set_trace as st\nfrom IPython import embed\n\nclass BaseModel():\n def __init__(self):\n pass;\n \n def name(self):\n return 'BaseModel'\n\n def initialize(self, use_gpu=True):\n self.use_gpu = use_gpu\n self.Tensor = torch.cuda.FloatTensor if self.use_gpu else torch.Tensor\n # self.save_dir = os.path.join(opt.checkpoints_dir, opt.name)\n\n def forward(self):\n pass\n\n def get_image_paths(self):\n pass\n\n def optimize_parameters(self):\n pass\n\n def get_current_visuals(self):\n return self.input\n\n def get_current_errors(self):\n return {}\n\n def save(self, label):\n pass\n\n # helper saving function that can be used by subclasses\n def save_network(self, network, path, network_label, epoch_label):\n save_filename = '%s_net_%s.pth' % (epoch_label, network_label)\n save_path = os.path.join(path, save_filename)\n torch.save(network.state_dict(), save_path)\n\n # helper loading function that can be used by subclasses\n def load_network(self, network, network_label, epoch_label):\n # embed()\n save_filename = '%s_net_%s.pth' % (epoch_label, network_label)\n save_path = os.path.join(self.save_dir, save_filename)\n print('Loading network from %s'%save_path)\n network.load_state_dict(torch.load(save_path))\n\n def update_learning_rate():\n pass\n\n def get_image_paths(self):\n return self.image_paths\n\n def save_done(self, flag=False):\n np.save(os.path.join(self.save_dir, 'done_flag'),flag)\n np.savetxt(os.path.join(self.save_dir, 'done_flag'),[flag,],fmt='%i')\n\n"
},
{
"path": "third_part/PerceptualSimilarity/models/dist_model.py",
"content": "\nfrom __future__ import absolute_import\n\nimport sys\nsys.path.append('..')\nsys.path.append('.')\nimport numpy as np\nimport torch\nfrom torch import nn\nimport os\nfrom collections import OrderedDict\nfrom torch.autograd import Variable\nimport itertools\nfrom .base_model import BaseModel\nfrom scipy.ndimage import zoom\nimport fractions\nimport functools\nimport skimage.transform\nfrom IPython import embed\n\nfrom . import networks_basic as networks\nfrom third_part.PerceptualSimilarity.util import util\n# from util import util\n\nclass DistModel(BaseModel):\n def name(self):\n return self.model_name\n\n def initialize(self, model='net-lin', net='alex', pnet_rand=False, pnet_tune=False, model_path=None, colorspace='Lab', use_gpu=True, printNet=False, spatial=False, spatial_shape=None, spatial_order=1, spatial_factor=None, is_train=False, lr=.0001, beta1=0.5, version='0.1'):\n '''\n INPUTS\n model - ['net-lin'] for linearly calibrated network\n ['net'] for off-the-shelf network\n ['L2'] for L2 distance in Lab colorspace\n ['SSIM'] for ssim in RGB colorspace\n net - ['squeeze','alex','vgg']\n model_path - if None, will look in weights/[NET_NAME].pth\n colorspace - ['Lab','RGB'] colorspace to use for L2 and SSIM\n use_gpu - bool - whether or not to use a GPU\n printNet - bool - whether or not to print network architecture out\n spatial - bool - whether to output an array containing varying distances across spatial dimensions\n spatial_shape - if given, output spatial shape. if None then spatial shape is determined automatically via spatial_factor (see below).\n spatial_factor - if given, specifies upsampling factor relative to the largest spatial extent of a convolutional layer. if None then resized to size of input images.\n spatial_order - spline order of filter for upsampling in spatial mode, by default 1 (bilinear).\n is_train - bool - [True] for training mode\n lr - float - initial learning rate\n beta1 - float - initial momentum term for adam\n version - 0.1 for latest, 0.0 was original\n '''\n BaseModel.initialize(self, use_gpu=use_gpu)\n\n self.model = model\n self.net = net\n self.use_gpu = use_gpu\n self.is_train = is_train\n self.spatial = spatial\n self.spatial_shape = spatial_shape\n self.spatial_order = spatial_order\n self.spatial_factor = spatial_factor\n\n self.model_name = '%s [%s]'%(model,net)\n if(self.model == 'net-lin'): # pretrained net + linear layer\n self.net = networks.PNetLin(use_gpu=use_gpu,pnet_rand=pnet_rand, pnet_tune=pnet_tune, pnet_type=net,use_dropout=True,spatial=spatial,version=version)\n kw = {}\n if not use_gpu:\n kw['map_location'] = 'cpu'\n if(model_path is None):\n import inspect\n # model_path = './PerceptualSimilarity/weights/v%s/%s.pth'%(version,net)\n model_path = os.path.abspath(os.path.join(inspect.getfile(self.initialize), '..', '..', 'weights/v%s/%s.pth'%(version,net)))\n\n if(not is_train):\n print('Loading model from: %s'%model_path)\n self.net.load_state_dict(torch.load(model_path, map_location=lambda storage, loc: storage))\n\n elif(self.model=='net'): # pretrained network\n assert not self.spatial, 'spatial argument not supported yet for uncalibrated networks'\n self.net = networks.PNet(use_gpu=use_gpu,pnet_type=net)\n self.is_fake_net = True\n elif(self.model in ['L2','l2']):\n self.net = networks.L2(use_gpu=use_gpu,colorspace=colorspace) # not really a network, only for testing\n self.model_name = 'L2'\n elif(self.model in ['DSSIM','dssim','SSIM','ssim']):\n self.net = networks.DSSIM(use_gpu=use_gpu,colorspace=colorspace)\n self.model_name = 'SSIM'\n else:\n raise ValueError(\"Model [%s] not recognized.\" % self.model)\n\n self.parameters = list(self.net.parameters())\n\n if self.is_train: # training mode\n # extra network on top to go from distances (d0,d1) => predicted human judgment (h*)\n self.rankLoss = networks.BCERankingLoss(use_gpu=use_gpu)\n self.parameters+=self.rankLoss.parameters\n self.lr = lr\n self.old_lr = lr\n self.optimizer_net = torch.optim.Adam(self.parameters, lr=lr, betas=(beta1, 0.999))\n else: # test mode\n self.net.eval()\n\n if(printNet):\n print('---------- Networks initialized -------------')\n networks.print_network(self.net)\n print('-----------------------------------------------')\n\n def forward_pair(self,in1,in2,retPerLayer=False):\n if(retPerLayer):\n return self.net.forward(in1,in2, retPerLayer=True)\n else:\n return self.net.forward(in1,in2)\n\n def forward(self, in0, in1, retNumpy=True):\n ''' Function computes the distance between image patches in0 and in1\n INPUTS\n in0, in1 - torch.Tensor object of shape Nx3xXxY - image patch scaled to [-1,1]\n retNumpy - [False] to return as torch.Tensor, [True] to return as numpy array\n OUTPUT\n computed distances between in0 and in1\n '''\n\n self.input_ref = in0\n self.input_p0 = in1\n\n if(self.use_gpu):\n self.input_ref = self.input_ref.cuda()\n self.input_p0 = self.input_p0.cuda()\n\n self.var_ref = Variable(self.input_ref,requires_grad=True)\n self.var_p0 = Variable(self.input_p0,requires_grad=True)\n\n self.d0 = self.forward_pair(self.var_ref, self.var_p0)\n self.loss_total = self.d0\n\n def convert_output(d0):\n if(retNumpy):\n ans = d0.cpu().data.numpy()\n if not self.spatial:\n ans = ans.flatten()\n else:\n assert(ans.shape[0] == 1 and len(ans.shape) == 4)\n return ans[0,...].transpose([1, 2, 0]) # Reshape to usual numpy image format: (height, width, channels)\n return ans\n else:\n return d0\n\n if self.spatial:\n L = [convert_output(x) for x in self.d0]\n spatial_shape = self.spatial_shape\n if spatial_shape is None:\n if(self.spatial_factor is None):\n spatial_shape = (in0.size()[2],in0.size()[3])\n else:\n spatial_shape = (max([x.shape[0] for x in L])*self.spatial_factor, max([x.shape[1] for x in L])*self.spatial_factor)\n \n L = [skimage.transform.resize(x, spatial_shape, order=self.spatial_order, mode='edge') for x in L]\n \n L = np.mean(np.concatenate(L, 2) * len(L), 2)\n return L\n else:\n return convert_output(self.d0)\n\n # ***** TRAINING FUNCTIONS *****\n def optimize_parameters(self):\n self.forward_train()\n self.optimizer_net.zero_grad()\n self.backward_train()\n self.optimizer_net.step()\n self.clamp_weights()\n\n def clamp_weights(self):\n for module in self.net.modules():\n if(hasattr(module, 'weight') and module.kernel_size==(1,1)):\n module.weight.data = torch.clamp(module.weight.data,min=0)\n\n def set_input(self, data):\n self.input_ref = data['ref']\n self.input_p0 = data['p0']\n self.input_p1 = data['p1']\n self.input_judge = data['judge']\n\n if(self.use_gpu):\n self.input_ref = self.input_ref.cuda()\n self.input_p0 = self.input_p0.cuda()\n self.input_p1 = self.input_p1.cuda()\n self.input_judge = self.input_judge.cuda()\n\n self.var_ref = Variable(self.input_ref,requires_grad=True)\n self.var_p0 = Variable(self.input_p0,requires_grad=True)\n self.var_p1 = Variable(self.input_p1,requires_grad=True)\n\n def forward_train(self): # run forward pass\n self.d0 = self.forward_pair(self.var_ref, self.var_p0)\n self.d1 = self.forward_pair(self.var_ref, self.var_p1)\n self.acc_r = self.compute_accuracy(self.d0,self.d1,self.input_judge)\n\n # var_judge\n self.var_judge = Variable(1.*self.input_judge).view(self.d0.size())\n\n self.loss_total = self.rankLoss.forward(self.d0, self.d1, self.var_judge*2.-1.)\n return self.loss_total\n\n def backward_train(self):\n torch.mean(self.loss_total).backward()\n\n def compute_accuracy(self,d0,d1,judge):\n ''' d0, d1 are Variables, judge is a Tensor '''\n d1_lt_d0 = (d1 %f' % (type,self.old_lr, lr))\n self.old_lr = lr\n\n\n\ndef score_2afc_dataset(data_loader,func):\n ''' Function computes Two Alternative Forced Choice (2AFC) score using\n distance function 'func' in dataset 'data_loader'\n INPUTS\n data_loader - CustomDatasetDataLoader object - contains a TwoAFCDataset inside\n func - callable distance function - calling d=func(in0,in1) should take 2\n pytorch tensors with shape Nx3xXxY, and return numpy array of length N\n OUTPUTS\n [0] - 2AFC score in [0,1], fraction of time func agrees with human evaluators\n [1] - dictionary with following elements\n d0s,d1s - N arrays containing distances between reference patch to perturbed patches \n gts - N array in [0,1], preferred patch selected by human evaluators\n (closer to \"0\" for left patch p0, \"1\" for right patch p1,\n \"0.6\" means 60pct people preferred right patch, 40pct preferred left)\n scores - N array in [0,1], corresponding to what percentage function agreed with humans\n CONSTS\n N - number of test triplets in data_loader\n '''\n\n d0s = []\n d1s = []\n gts = []\n\n # bar = pb.ProgressBar(max_value=data_loader.load_data().__len__())\n for (i,data) in enumerate(data_loader.load_data()):\n d0s+=func(data['ref'],data['p0']).tolist()\n d1s+=func(data['ref'],data['p1']).tolist()\n gts+=data['judge'].cpu().numpy().flatten().tolist()\n # bar.update(i)\n\n d0s = np.array(d0s)\n d1s = np.array(d1s)\n gts = np.array(gts)\n scores = (d0s 0:\n with self.doc.head:\n meta(http_equiv=\"reflesh\", content=str(reflesh))\n\n def get_image_dir(self):\n return self.img_dir\n\n def add_header(self, str):\n with self.doc:\n h3(str)\n\n def add_table(self, border=1):\n self.t = table(border=border, style=\"table-layout: fixed;\")\n self.doc.add(self.t)\n\n def add_images(self, ims, txts, links, width=400):\n self.add_table()\n with self.t:\n with tr():\n for im, txt, link in zip(ims, txts, links):\n with td(style=\"word-wrap: break-word;\", halign=\"center\", valign=\"top\"):\n with p():\n with a(href=os.path.join(link)):\n img(style=\"width:%dpx\" % width, src=os.path.join(im))\n br()\n p(txt)\n\n def save(self,file='index'):\n html_file = '%s/%s.html' % (self.web_dir,file)\n f = open(html_file, 'wt')\n f.write(self.doc.render())\n f.close()\n\n\nif __name__ == '__main__':\n html = HTML('web/', 'test_html')\n html.add_header('hello world')\n\n ims = []\n txts = []\n links = []\n for n in range(4):\n ims.append('image_%d.png' % n)\n txts.append('text_%d' % n)\n links.append('image_%d.png' % n)\n html.add_images(ims, txts, links)\n html.save()\n"
},
{
"path": "third_part/PerceptualSimilarity/util/util.py",
"content": "from __future__ import print_function\n\nimport numpy as np\nfrom PIL import Image\nimport inspect\nimport re\nimport numpy as np\nimport os\nimport collections\nimport matplotlib.pyplot as plt\nfrom scipy.ndimage.interpolation import zoom\nfrom skimage.measure import compare_ssim\nimport torch\nfrom IPython import embed\nimport cv2\nfrom datetime import datetime\n\ndef datetime_str():\n now = datetime.now()\n return '%04d-%02d-%02d-%02d-%02d-%02d'%(now.year,now.month,now.day,now.hour,now.minute,now.second)\n\ndef read_text_file(in_path):\n fid = open(in_path,'r')\n\n vals = []\n cur_line = fid.readline()\n while(cur_line!=''):\n vals.append(float(cur_line))\n cur_line = fid.readline()\n\n fid.close()\n return np.array(vals)\n\ndef bootstrap(in_vec,num_samples=100,bootfunc=np.mean):\n from astropy import stats\n return stats.bootstrap(np.array(in_vec),bootnum=num_samples,bootfunc=bootfunc)\n\ndef rand_flip(input1,input2):\n if(np.random.binomial(1,.5)==1):\n return (input1,input2)\n else:\n return (input2,input1)\n\ndef l2(p0, p1, range=255.):\n return .5*np.mean((p0 / range - p1 / range)**2)\n\ndef psnr(p0, p1, peak=255.):\n return 10*np.log10(peak**2/np.mean((1.*p0-1.*p1)**2))\n\ndef dssim(p0, p1, range=255.):\n # embed()\n return (1 - compare_ssim(p0, p1, data_range=range, multichannel=True)) / 2.\n\ndef rgb2lab(in_img,mean_cent=False):\n from skimage import color\n img_lab = color.rgb2lab(in_img)\n if(mean_cent):\n img_lab[:,:,0] = img_lab[:,:,0]-50\n return img_lab\n\ndef normalize_blob(in_feat,eps=1e-10):\n norm_factor = np.sqrt(np.sum(in_feat**2,axis=1,keepdims=True))\n return in_feat/(norm_factor+eps)\n\ndef cos_sim_blob(in0,in1):\n in0_norm = normalize_blob(in0)\n in1_norm = normalize_blob(in1)\n (N,C,X,Y) = in0_norm.shape\n\n return np.mean(np.mean(np.sum(in0_norm*in1_norm,axis=1),axis=1),axis=1)\n\ndef normalize_tensor(in_feat,eps=1e-10):\n # norm_factor = torch.sqrt(torch.sum(in_feat**2,dim=1)).view(in_feat.size()[0],1,in_feat.size()[2],in_feat.size()[3]).repeat(1,in_feat.size()[1],1,1)\n norm_factor = torch.sqrt(torch.sum(in_feat**2,dim=1)).view(in_feat.size()[0],1,in_feat.size()[2],in_feat.size()[3])\n return in_feat/(norm_factor.expand_as(in_feat)+eps)\n\ndef cos_sim(in0,in1):\n in0_norm = normalize_tensor(in0)\n in1_norm = normalize_tensor(in1)\n N = in0.size()[0]\n X = in0.size()[2]\n Y = in0.size()[3]\n\n return torch.mean(torch.mean(torch.sum(in0_norm*in1_norm,dim=1).view(N,1,X,Y),dim=2).view(N,1,1,Y),dim=3).view(N)\n\n# Converts a Tensor into a Numpy array\n# |imtype|: the desired type of the conve\n\ndef tensor2np(tensor_obj):\n # change dimension of a tensor object into a numpy array\n return tensor_obj[0].cpu().float().numpy().transpose((1,2,0))\n\ndef np2tensor(np_obj):\n # change dimenion of np array into tensor array\n return torch.Tensor(np_obj[:, :, :, np.newaxis].transpose((3, 2, 0, 1)))\n\ndef tensor2tensorlab(image_tensor,to_norm=True,mc_only=False):\n # image tensor to lab tensor\n from skimage import color\n\n img = tensor2im(image_tensor)\n # print('img_rgb',img.flatten())\n img_lab = color.rgb2lab(img)\n # print('img_lab',img_lab.flatten())\n if(mc_only):\n img_lab[:,:,0] = img_lab[:,:,0]-50\n if(to_norm and not mc_only):\n img_lab[:,:,0] = img_lab[:,:,0]-50\n img_lab = img_lab/100.\n\n return np2tensor(img_lab)\n\ndef tensorlab2tensor(lab_tensor,return_inbnd=False):\n from skimage import color\n import warnings\n warnings.filterwarnings(\"ignore\")\n\n lab = tensor2np(lab_tensor)*100.\n lab[:,:,0] = lab[:,:,0]+50\n # print('lab',lab)\n\n rgb_back = 255.*np.clip(color.lab2rgb(lab.astype('float')),0,1)\n # print('rgb',rgb_back)\n if(return_inbnd):\n # convert back to lab, see if we match\n lab_back = color.rgb2lab(rgb_back.astype('uint8'))\n # print('lab_back',lab_back)\n # print('lab==lab_back',np.isclose(lab_back,lab,atol=1.))\n # print('lab-lab_back',np.abs(lab-lab_back))\n mask = 1.*np.isclose(lab_back,lab,atol=2.)\n mask = np2tensor(np.prod(mask,axis=2)[:,:,np.newaxis])\n return (im2tensor(rgb_back),mask)\n else:\n return im2tensor(rgb_back)\n\ndef tensor2im(image_tensor, imtype=np.uint8, cent=1., factor=255./2.):\n# def tensor2im(image_tensor, imtype=np.uint8, cent=1., factor=1.):\n image_numpy = image_tensor[0].cpu().float().numpy()\n image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + cent) * factor\n return image_numpy.astype(imtype)\n\ndef im2tensor(image, imtype=np.uint8, cent=1., factor=255./2.):\n# def im2tensor(image, imtype=np.uint8, cent=1., factor=1.):\n return torch.Tensor((image / factor - cent)\n [:, :, :, np.newaxis].transpose((3, 2, 0, 1)))\n\ndef tensor2vec(vector_tensor):\n return vector_tensor.data.cpu().numpy()[:, :, 0, 0]\n\ndef diagnose_network(net, name='network'):\n mean = 0.0\n count = 0\n for param in net.parameters():\n if param.grad is not None:\n mean += torch.mean(torch.abs(param.grad.data))\n count += 1\n if count > 0:\n mean = mean / count\n print(name)\n print(mean)\n\ndef grab_patch(img_in, P, yy, xx):\n return img_in[yy:yy+P,xx:xx+P,:]\n\ndef load_image(path):\n if(path[-3:] == 'dng'):\n import rawpy\n with rawpy.imread(path) as raw:\n img = raw.postprocess()\n # img = plt.imread(path)\n elif(path[-3:]=='bmp' or path[-3:]=='jpg' or path[-3:]=='png'):\n import cv2\n return cv2.imread(path)[:,:,::-1]\n else:\n img = (255*plt.imread(path)[:,:,:3]).astype('uint8')\n\n return img\n\n\ndef resize_image(img, max_size=256):\n [Y, X] = img.shape[:2]\n\n # resize\n max_dim = max([Y, X])\n zoom_factor = 1. * max_size / max_dim\n img = zoom(img, [zoom_factor, zoom_factor, 1])\n\n return img\n\ndef resize_image_zoom(img, zoom_factor=1., order=3):\n if(zoom_factor==1):\n return img\n else:\n return zoom(img, [zoom_factor, zoom_factor, 1], order=order)\n\ndef save_image(image_numpy, image_path, ):\n image_pil = Image.fromarray(image_numpy)\n image_pil.save(image_path)\n\n\ndef prep_display_image(img, dtype='uint8'):\n if(dtype == 'uint8'):\n return np.clip(img, 0, 255).astype('uint8')\n else:\n return np.clip(img, 0, 1.)\n\n\ndef info(object, spacing=10, collapse=1):\n \"\"\"Print methods and doc strings.\n Takes module, class, list, dictionary, or string.\"\"\"\n methodList = [\n e for e in dir(object) if isinstance(\n getattr(\n object,\n e),\n collections.Callable)]\n processFunc = collapse and (lambda s: \" \".join(s.split())) or (lambda s: s)\n print(\"\\n\".join([\"%s %s\" %\n (method.ljust(spacing),\n processFunc(str(getattr(object, method).__doc__)))\n for method in methodList]))\n\n\ndef varname(p):\n for line in inspect.getframeinfo(inspect.currentframe().f_back)[3]:\n m = re.search(r'\\bvarname\\s*\\(\\s*([A-Za-z_][A-Za-z0-9_]*)\\s*\\)', line)\n if m:\n return m.group(1)\n\n\ndef print_numpy(x, val=True, shp=False):\n x = x.astype(np.float64)\n if shp:\n print('shape,', x.shape)\n if val:\n x = x.flatten()\n print(\n 'mean = %3.3f, min = %3.3f, max = %3.3f, median = %3.3f, std=%3.3f' %\n (np.mean(x), np.min(x), np.max(x), np.median(x), np.std(x)))\n\n\ndef mkdirs(paths):\n if isinstance(paths, list) and not isinstance(paths, str):\n for path in paths:\n mkdir(path)\n else:\n mkdir(paths)\n\n\ndef mkdir(path):\n if not os.path.exists(path):\n os.makedirs(path)\n\n\ndef rgb2lab(input):\n from skimage import color\n return color.rgb2lab(input / 255.)\n\n\ndef montage(\n imgs,\n PAD=5,\n RATIO=16 / 9.,\n EXTRA_PAD=(\n False,\n False),\n MM=-1,\n NN=-1,\n primeDir=0,\n verbose=False,\n returnGridPos=False,\n backClr=np.array(\n (0,\n 0,\n 0))):\n # INPUTS\n # imgs YxXxMxN or YxXxN\n # PAD scalar number of pixels in between\n # RATIO scalar target ratio of cols/rows\n # MM scalar # rows, if specified, overrides RATIO\n # NN scalar # columns, if specified, overrides RATIO\n # primeDir scalar 0 for top-to-bottom, 1 for left-to-right\n # OUTPUTS\n # mont_imgs MM*Y x NN*X x M big image with everything montaged\n # def montage(imgs, PAD=5, RATIO=16/9., MM=-1, NN=-1, primeDir=0,\n # verbose=False, forceFloat=False):\n if(imgs.ndim == 3):\n toExp = True\n imgs = imgs[:, :, np.newaxis, :]\n else:\n toExp = False\n\n Y = imgs.shape[0]\n X = imgs.shape[1]\n M = imgs.shape[2]\n N = imgs.shape[3]\n\n PADS = np.array((PAD))\n if(PADS.flatten().size == 1):\n PADY = PADS\n PADX = PADS\n else:\n PADY = PADS[0]\n PADX = PADS[1]\n\n if(MM == -1 and NN == -1):\n NN = np.ceil(np.sqrt(1.0 * N * RATIO))\n MM = np.ceil(1.0 * N / NN)\n NN = np.ceil(1.0 * N / MM)\n elif(MM == -1):\n MM = np.ceil(1.0 * N / NN)\n elif(NN == -1):\n NN = np.ceil(1.0 * N / MM)\n\n if(primeDir == 0): # write top-to-bottom\n [grid_mm, grid_nn] = np.meshgrid(\n np.arange(MM, dtype='uint'), np.arange(NN, dtype='uint'))\n elif(primeDir == 1): # write left-to-right\n [grid_nn, grid_mm] = np.meshgrid(\n np.arange(NN, dtype='uint'), np.arange(MM, dtype='uint'))\n\n grid_mm = np.uint(grid_mm.flatten()[0:N])\n grid_nn = np.uint(grid_nn.flatten()[0:N])\n\n EXTRA_PADY = EXTRA_PAD[0] * PADY\n EXTRA_PADX = EXTRA_PAD[0] * PADX\n\n # mont_imgs = np.zeros(((Y+PAD)*MM-PAD, (X+PAD)*NN-PAD, M), dtype=use_dtype)\n mont_imgs = np.zeros(\n (np.uint(\n (Y + PADY) * MM - PADY + EXTRA_PADY),\n np.uint(\n (X + PADX) * NN - PADX + EXTRA_PADX),\n M),\n dtype=imgs.dtype)\n mont_imgs = mont_imgs + \\\n backClr.flatten()[np.newaxis, np.newaxis, :].astype(mont_imgs.dtype)\n\n for ii in np.random.permutation(N):\n # print imgs[:,:,:,ii].shape\n # mont_imgs[grid_mm[ii]*(Y+PAD):(grid_mm[ii]*(Y+PAD)+Y), grid_nn[ii]*(X+PAD):(grid_nn[ii]*(X+PAD)+X),:]\n mont_imgs[np.uint(grid_mm[ii] *\n (Y +\n PADY)):np.uint((grid_mm[ii] *\n (Y +\n PADY) +\n Y)), np.uint(grid_nn[ii] *\n (X +\n PADX)):np.uint((grid_nn[ii] *\n (X +\n PADX) +\n X)), :] = imgs[:, :, :, ii]\n\n if(M == 1):\n imgs = imgs.reshape(imgs.shape[0], imgs.shape[1], imgs.shape[3])\n\n if(toExp):\n mont_imgs = mont_imgs[:, :, 0]\n\n if(returnGridPos):\n # return (mont_imgs,np.concatenate((grid_mm[:,:,np.newaxis]*(Y+PAD),\n # grid_nn[:,:,np.newaxis]*(X+PAD)),axis=2))\n return (mont_imgs, np.concatenate(\n (grid_mm[:, np.newaxis] * (Y + PADY), grid_nn[:, np.newaxis] * (X + PADX)), axis=1))\n # return (mont_imgs, (grid_mm,grid_nn))\n else:\n return mont_imgs\n\nclass zeroClipper(object):\n def __init__(self, frequency=1):\n self.frequency = frequency\n\n def __call__(self, module):\n embed()\n if hasattr(module, 'weight'):\n # module.weight.data = torch.max(module.weight.data, 0)\n module.weight.data = torch.max(module.weight.data, 0) + 100\n\ndef flatten_nested_list(nested_list):\n # only works for list of list\n accum = []\n for sublist in nested_list:\n for item in sublist:\n accum.append(item)\n return accum\n\ndef read_file(in_path,list_lines=False):\n agg_str = ''\n f = open(in_path,'r')\n cur_line = f.readline()\n while(cur_line!=''):\n agg_str+=cur_line\n cur_line = f.readline()\n f.close()\n if(list_lines==False):\n return agg_str.replace('\\n','')\n else:\n line_list = agg_str.split('\\n')\n ret_list = []\n for item in line_list:\n if(item!=''):\n ret_list.append(item)\n return ret_list\n\ndef read_csv_file_as_text(in_path):\n agg_str = []\n f = open(in_path,'r')\n cur_line = f.readline()\n while(cur_line!=''):\n agg_str.append(cur_line)\n cur_line = f.readline()\n f.close()\n return agg_str\n\ndef random_swap(obj0,obj1):\n if(np.random.rand() < .5):\n return (obj0,obj1,0)\n else:\n return (obj1,obj0,1)\n\ndef voc_ap(rec, prec, use_07_metric=False):\n \"\"\" ap = voc_ap(rec, prec, [use_07_metric])\n Compute VOC AP given precision and recall.\n If use_07_metric is true, uses the\n VOC 07 11 point method (default:False).\n \"\"\"\n if use_07_metric:\n # 11 point metric\n ap = 0.\n for t in np.arange(0., 1.1, 0.1):\n if np.sum(rec >= t) == 0:\n p = 0\n else:\n p = np.max(prec[rec >= t])\n ap = ap + p / 11.\n else:\n # correct AP calculation\n # first append sentinel values at the end\n mrec = np.concatenate(([0.], rec, [1.]))\n mpre = np.concatenate(([0.], prec, [0.]))\n\n # compute the precision envelope\n for i in range(mpre.size - 1, 0, -1):\n mpre[i - 1] = np.maximum(mpre[i - 1], mpre[i])\n\n # to calculate area under PR curve, look for points\n # where X axis (recall) changes value\n i = np.where(mrec[1:] != mrec[:-1])[0]\n\n # and sum (\\Delta recall) * prec\n ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1])\n return ap\n"
},
{
"path": "third_part/PerceptualSimilarity/util/visualizer.py",
"content": "import numpy as np\nimport os\nimport time\nfrom . import util\nfrom . import html\n# from pdb import set_trace as st\nimport matplotlib.pyplot as plt\nimport math\n# from IPython import embed\n\ndef zoom_to_res(img,res=256,order=0,axis=0):\n # img 3xXxX\n from scipy.ndimage import zoom\n zoom_factor = res/img.shape[1]\n if(axis==0):\n return zoom(img,[1,zoom_factor,zoom_factor],order=order)\n elif(axis==2):\n return zoom(img,[zoom_factor,zoom_factor,1],order=order)\n\nclass Visualizer():\n def __init__(self, opt):\n # self.opt = opt\n self.display_id = opt.display_id\n # self.use_html = opt.is_train and not opt.no_html\n self.win_size = opt.display_winsize\n self.name = opt.name\n self.display_cnt = 0 # display_current_results counter\n self.display_cnt_high = 0\n self.use_html = opt.use_html\n\n if self.display_id > 0:\n import visdom\n self.vis = visdom.Visdom(port = opt.display_port)\n\n self.web_dir = os.path.join(opt.checkpoints_dir, opt.name, 'web')\n util.mkdirs([self.web_dir,])\n if self.use_html:\n self.img_dir = os.path.join(self.web_dir, 'images')\n print('create web directory %s...' % self.web_dir)\n util.mkdirs([self.img_dir,])\n\n # |visuals|: dictionary of images to display or save\n def display_current_results(self, visuals, epoch, nrows=None, res=256):\n if self.display_id > 0: # show images in the browser\n title = self.name\n if(nrows is None):\n nrows = int(math.ceil(len(visuals.items()) / 2.0))\n images = []\n idx = 0\n for label, image_numpy in visuals.items():\n title += \" | \" if idx % nrows == 0 else \", \"\n title += label\n img = image_numpy.transpose([2, 0, 1])\n img = zoom_to_res(img,res=res,order=0)\n images.append(img)\n idx += 1\n if len(visuals.items()) % 2 != 0:\n white_image = np.ones_like(image_numpy.transpose([2, 0, 1]))*255\n white_image = zoom_to_res(white_image,res=res,order=0)\n images.append(white_image)\n self.vis.images(images, nrow=nrows, win=self.display_id + 1,\n opts=dict(title=title))\n\n if self.use_html: # save images to a html file\n for label, image_numpy in visuals.items():\n img_path = os.path.join(self.img_dir, 'epoch%.3d_cnt%.6d_%s.png' % (epoch, self.display_cnt, label))\n util.save_image(zoom_to_res(image_numpy, res=res, axis=2), img_path)\n\n self.display_cnt += 1\n self.display_cnt_high = np.maximum(self.display_cnt_high, self.display_cnt)\n\n # update website\n webpage = html.HTML(self.web_dir, 'Experiment name = %s' % self.name, reflesh=1)\n for n in range(epoch, 0, -1):\n webpage.add_header('epoch [%d]' % n)\n if(n==epoch):\n high = self.display_cnt\n else:\n high = self.display_cnt_high\n for c in range(high-1,-1,-1):\n ims = []\n txts = []\n links = []\n\n for label, image_numpy in visuals.items():\n img_path = 'epoch%.3d_cnt%.6d_%s.png' % (n, c, label)\n ims.append(os.path.join('images',img_path))\n txts.append(label)\n links.append(os.path.join('images',img_path))\n webpage.add_images(ims, txts, links, width=self.win_size)\n webpage.save()\n\n # save errors into a directory\n def plot_current_errors_save(self, epoch, counter_ratio, opt, errors,keys='+ALL',name='loss', to_plot=False):\n if not hasattr(self, 'plot_data'):\n self.plot_data = {'X':[],'Y':[], 'legend':list(errors.keys())}\n self.plot_data['X'].append(epoch + counter_ratio)\n self.plot_data['Y'].append([errors[k] for k in self.plot_data['legend']])\n\n # embed()\n if(keys=='+ALL'):\n plot_keys = self.plot_data['legend']\n else:\n plot_keys = keys\n\n if(to_plot):\n (f,ax) = plt.subplots(1,1)\n for (k,kname) in enumerate(plot_keys):\n kk = np.where(np.array(self.plot_data['legend'])==kname)[0][0]\n x = self.plot_data['X']\n y = np.array(self.plot_data['Y'])[:,kk]\n if(to_plot):\n ax.plot(x, y, 'o-', label=kname)\n np.save(os.path.join(self.web_dir,'%s_x')%kname,x)\n np.save(os.path.join(self.web_dir,'%s_y')%kname,y)\n\n if(to_plot):\n plt.legend(loc=0,fontsize='small')\n plt.xlabel('epoch')\n plt.ylabel('Value')\n f.savefig(os.path.join(self.web_dir,'%s.png'%name))\n f.clf()\n plt.close()\n\n # errors: dictionary of error labels and values\n def plot_current_errors(self, epoch, counter_ratio, opt, errors):\n if not hasattr(self, 'plot_data'):\n self.plot_data = {'X':[],'Y':[], 'legend':list(errors.keys())}\n self.plot_data['X'].append(epoch + counter_ratio)\n self.plot_data['Y'].append([errors[k] for k in self.plot_data['legend']])\n self.vis.line(\n X=np.stack([np.array(self.plot_data['X'])]*len(self.plot_data['legend']),1),\n Y=np.array(self.plot_data['Y']),\n opts={\n 'title': self.name + ' loss over time',\n 'legend': self.plot_data['legend'],\n 'xlabel': 'epoch',\n 'ylabel': 'loss'},\n win=self.display_id)\n\n # errors: same format as |errors| of plotCurrentErrors\n def print_current_errors(self, epoch, i, errors, t, t2=-1, t2o=-1, fid=None):\n message = '(ep: %d, it: %d, t: %.3f[s], ept: %.2f/%.2f[h]) ' % (epoch, i, t, t2o, t2)\n message += (', ').join(['%s: %.3f' % (k, v) for k, v in errors.items()])\n\n print(message)\n if(fid is not None):\n fid.write('%s\\n'%message)\n\n\n # save image to the disk\n def save_images_simple(self, webpage, images, names, in_txts, prefix='', res=256):\n image_dir = webpage.get_image_dir()\n ims = []\n txts = []\n links = []\n\n for name, image_numpy, txt in zip(names, images, in_txts):\n image_name = '%s_%s.png' % (prefix, name)\n save_path = os.path.join(image_dir, image_name)\n if(res is not None):\n util.save_image(zoom_to_res(image_numpy,res=res,axis=2), save_path)\n else:\n util.save_image(image_numpy, save_path)\n\n ims.append(os.path.join(webpage.img_subdir,image_name))\n # txts.append(name)\n txts.append(txt)\n links.append(os.path.join(webpage.img_subdir,image_name))\n # embed()\n webpage.add_images(ims, txts, links, width=self.win_size)\n\n # save image to the disk\n def save_images(self, webpage, images, names, image_path, title=''):\n image_dir = webpage.get_image_dir()\n # short_path = ntpath.basename(image_path)\n # name = os.path.splitext(short_path)[0]\n # name = short_path\n # webpage.add_header('%s, %s' % (name, title))\n ims = []\n txts = []\n links = []\n\n for label, image_numpy in zip(names, images):\n image_name = '%s.jpg' % (label,)\n save_path = os.path.join(image_dir, image_name)\n util.save_image(image_numpy, save_path)\n\n ims.append(image_name)\n txts.append(label)\n links.append(image_name)\n webpage.add_images(ims, txts, links, width=self.win_size)\n\n # save image to the disk\n # def save_images(self, webpage, visuals, image_path, short=False):\n # image_dir = webpage.get_image_dir()\n # if short:\n # short_path = ntpath.basename(image_path)\n # name = os.path.splitext(short_path)[0]\n # else:\n # name = image_path\n\n # webpage.add_header(name)\n # ims = []\n # txts = []\n # links = []\n\n # for label, image_numpy in visuals.items():\n # image_name = '%s_%s.png' % (name, label)\n # save_path = os.path.join(image_dir, image_name)\n # util.save_image(image_numpy, save_path)\n\n # ims.append(image_name)\n # txts.append(label)\n # links.append(image_name)\n # webpage.add_images(ims, txts, links, width=self.win_size)\n"
},
{
"path": "train.py",
"content": "import argparse\n\nimport data as Dataset\nfrom config import Config\nfrom util.logging import init_logging, make_logging_dir\nfrom util.trainer import get_model_optimizer_and_scheduler, set_random_seed, get_trainer\nfrom util.distributed import init_dist\nfrom util.distributed import master_only_print as print\n\n\ndef parse_args():\n parser = argparse.ArgumentParser(description='Training')\n parser.add_argument('--config', default='./config/face.yaml')\n parser.add_argument('--name', default=None)\n parser.add_argument('--checkpoints_dir', default='result',\n help='Dir for saving logs and models.')\n parser.add_argument('--seed', type=int, default=0, help='Random seed.')\n parser.add_argument('--which_iter', type=int, default=None)\n parser.add_argument('--no_resume', action='store_true')\n parser.add_argument('--local_rank', type=int, default=0)\n parser.add_argument('--single_gpu', action='store_true')\n parser.add_argument('--debug', action='store_true')\n\n args = parser.parse_args()\n return args\n\n\nif __name__ == '__main__':\n # get training options\n args = parse_args()\n set_random_seed(args.seed)\n opt = Config(args.config, args, is_train=True)\n\n if not args.single_gpu:\n opt.local_rank = args.local_rank\n init_dist(opt.local_rank) \n opt.device = opt.local_rank\n \n # create a visualizer\n date_uid, logdir = init_logging(opt)\n opt.logdir = logdir\n make_logging_dir(logdir, date_uid)\n # create a dataset\n val_dataset, train_dataset = Dataset.get_train_val_dataloader(opt.data)\n\n # create a model\n net_G, net_G_ema, opt_G, sch_G \\\n = get_model_optimizer_and_scheduler(opt)\n\n trainer = get_trainer(opt, net_G, net_G_ema, opt_G, sch_G, train_dataset)\n\n current_epoch, current_iteration = trainer.load_checkpoint(opt, args.which_iter) \n # training flag\n max_epoch = opt.max_epoch\n\n if args.debug:\n trainer.test_everything(train_dataset, val_dataset, current_epoch, current_iteration)\n exit()\n # Start training.\n for epoch in range(current_epoch, opt.max_epoch):\n print('Epoch {} ...'.format(epoch))\n if not args.single_gpu:\n train_dataset.sampler.set_epoch(current_epoch)\n trainer.start_of_epoch(current_epoch)\n for it, data in enumerate(train_dataset):\n data = trainer.start_of_iteration(data, current_iteration)\n trainer.optimize_parameters(data)\n current_iteration += 1\n trainer.end_of_iteration(data, current_epoch, current_iteration)\n \n if current_iteration >= opt.max_iter:\n print('Done with training!!!')\n break\n current_epoch += 1\n trainer.end_of_epoch(data, val_dataset, current_epoch, current_iteration)\n"
},
{
"path": "trainers/__init__.py",
"content": ""
},
{
"path": "trainers/base.py",
"content": "import os\nimport time\nimport glob\nfrom tqdm import tqdm\n\nimport torch\nimport torchvision\nfrom torch import nn\n\nfrom util.distributed import is_master, master_only\nfrom util.distributed import master_only_print as print\nfrom util.meters import Meter, add_hparams\nfrom util.misc import to_cuda, to_device, requires_grad\nfrom util.lpips import LPIPS\n\n\n\nclass BaseTrainer(object):\n r\"\"\"Base trainer. We expect that all trainers inherit this class.\n\n Args:\n opt (obj): Global configuration.\n net_G (obj): Generator network.\n net_D (obj): Discriminator network.\n opt_G (obj): Optimizer for the generator network.\n opt_D (obj): Optimizer for the discriminator network.\n sch_G (obj): Scheduler for the generator optimizer.\n sch_D (obj): Scheduler for the discriminator optimizer.\n train_data_loader (obj): Train data loader.\n val_data_loader (obj): Validation data loader.\n \"\"\"\n\n def __init__(self,\n opt,\n net_G,\n net_G_ema,\n opt_G,\n sch_G,\n train_data_loader,\n val_data_loader=None):\n super(BaseTrainer, self).__init__()\n print('Setup trainer.')\n\n # Initialize models and data loaders.\n self.opt = opt\n self.net_G = net_G\n if opt.distributed:\n self.net_G_module = self.net_G.module\n else:\n self.net_G_module = self.net_G\n\n self.is_inference = train_data_loader is None\n self.net_G_ema = net_G_ema\n self.opt_G = opt_G\n self.sch_G = sch_G\n self.train_data_loader = train_data_loader\n\n self.criteria = nn.ModuleDict()\n self.weights = dict()\n self.losses = dict(gen_update=dict(), dis_update=dict())\n self.gen_losses = self.losses['gen_update']\n self._init_loss(opt)\n for loss_name, loss_weight in self.weights.items():\n print(\"Loss {:<20} Weight {}\".format(loss_name, loss_weight))\n if loss_name in self.criteria.keys() and \\\n self.criteria[loss_name] is not None:\n self.criteria[loss_name].to('cuda')\n\n if self.is_inference:\n # The initialization steps below can be skipped during inference.\n return\n\n # Initialize logging attributes.\n self.current_iteration = 0\n self.current_epoch = 0\n self.start_iteration_time = None\n self.start_epoch_time = None\n self.elapsed_iteration_time = 0\n self.time_iteration = -1\n self.time_epoch = -1\n if getattr(self.opt, 'speed_benchmark', False):\n self.accu_gen_forw_iter_time = 0\n self.accu_gen_loss_iter_time = 0\n self.accu_gen_back_iter_time = 0\n self.accu_gen_step_iter_time = 0\n self.accu_gen_avg_iter_time = 0\n\n # Initialize tensorboard and hparams.\n self._init_tensorboard()\n self._init_hparams()\n self.lpips = LPIPS()\n self.best_lpips = None\n\n def _init_tensorboard(self):\n r\"\"\"Initialize the tensorboard. Different algorithms might require\n different performance metrics. Hence, custom tensorboard\n initialization might be necessary.\n \"\"\"\n # Logging frequency: self.opt.logging_iter\n self.meters = {}\n names = ['optim/gen_lr', 'time/iteration', 'time/epoch', \n 'metric/best_lpips', 'metric/lpips']\n for name in names:\n self.meters[name] = Meter(name)\n\n # Logging frequency: self.opt.image_display_iter\n self.image_meter = Meter('images')\n\n # Logging frequency: self.opt.snapshot_save_iter\n # self.meters['metric/lpips'] = Meter('metric/lpips')\n\n\n def _init_hparams(self):\n r\"\"\"Initialize a dictionary of hyperparameters that we want to monitor\n in the HParams dashboard in tensorBoard.\n \"\"\"\n self.hparam_dict = {}\n\n def _write_tensorboard(self):\n r\"\"\"Write values to tensorboard. By default, we will log the time used\n per iteration, time used per epoch, generator learning rate, and\n discriminator learning rate. We will log all the losses as well as\n custom meters.\n \"\"\"\n # Logs that are shared by all models.\n self._write_to_meters({'time/iteration': self.time_iteration,\n 'time/epoch': self.time_epoch,\n 'optim/gen_lr': self.sch_G.get_last_lr()[0]},\n self.meters)\n # Logs for loss values. Different models have different losses.\n self._write_loss_meters()\n # Other custom logs.\n self._write_custom_meters()\n # Write all logs to tensorboard.\n self._flush_meters(self.meters)\n\n def _write_loss_meters(self):\n r\"\"\"Write all loss values to tensorboard.\"\"\"\n for loss_name, loss in self.gen_losses.items():\n full_loss_name = 'gen_update' + '/' + loss_name\n if full_loss_name not in self.meters.keys():\n # Create a new meter if it doesn't exist.\n self.meters[full_loss_name] = Meter(full_loss_name)\n self.meters[full_loss_name].write(loss.item())\n\n def test_everything(self, train_dataset, val_dataset, current_epoch, current_iteration):\n r\"\"\"Test the functions defined in the models. by default, we will test the \n training function, the inference function, the visualization function.\n \"\"\" \n self._set_custom_debug_parameter()\n self.start_of_epoch(current_epoch)\n print('Start testing your functions')\n for it in tqdm(range(30)):\n data = iter(train_dataset).next()\n data = self.start_of_iteration(data, current_iteration)\n self.optimize_parameters(data)\n current_iteration += 1\n self.end_of_iteration(data, current_epoch, current_iteration)\n \n self.save_image(self._get_save_path('image', 'jpg'), data)\n self._write_tensorboard()\n self._print_current_errors()\n self.write_metrics(data)\n self.end_of_epoch(data, val_dataset, current_epoch, current_iteration)\n print('End debugging')\n \n\n def _set_custom_debug_parameter(self):\n r\"\"\"Set custom debug parame.\n \"\"\"\n self.opt.logging_iter = 10\n self.opt.image_save_iter = 10\n \n\n def _write_custom_meters(self):\n r\"\"\"Dummy member function to be overloaded by the child class.\n In the child class, you can write down whatever you want to track.\n \"\"\"\n pass\n\n @staticmethod\n def _write_to_meters(data, meters):\n r\"\"\"Write values to meters.\"\"\"\n for key, value in data.items():\n meters[key].write(value)\n\n def _flush_meters(self, meters):\n r\"\"\"Flush all meters using the current iteration.\"\"\"\n for meter in meters.values():\n meter.flush(self.current_iteration)\n\n def _pre_save_checkpoint(self):\n r\"\"\"Implement the things you want to do before saving a checkpoint.\n For example, you can compute the K-mean features (pix2pixHD) before\n saving the model weights to a checkpoint.\n \"\"\"\n pass\n\n def save_checkpoint(self, current_epoch, current_iteration):\n r\"\"\"Save network weights, optimizer parameters, scheduler parameters\n to a checkpoint.\n \"\"\"\n self._pre_save_checkpoint()\n _save_checkpoint(self.opt,\n self.net_G, self.net_G_ema, \n self.opt_G, self.sch_G,\n current_epoch, current_iteration)\n\n def load_checkpoint(self, opt, which_iter=None):\n if which_iter is not None:\n model_path = os.path.join(\n opt.logdir, '*_iteration_{:09}_checkpoint.pt'.format(which_iter))\n latest_checkpoint_path = glob.glob(model_path)\n assert len(latest_checkpoint_path) <= 1, \"please check the saved model {}\".format(\n model_path)\n if len(latest_checkpoint_path) == 0:\n current_epoch = 0\n current_iteration = 0\n print('No checkpoint found at iteration {}.'.format(which_iter))\n return current_epoch, current_iteration\n checkpoint_path = latest_checkpoint_path[0]\n\n elif os.path.exists(os.path.join(opt.logdir, 'latest_checkpoint.txt')):\n with open(os.path.join(opt.logdir, 'latest_checkpoint.txt'), 'r') as f:\n line = f.readlines()[0].replace('\\n', '')\n checkpoint_path = os.path.join(opt.logdir, line.split(' ')[-1])\n else:\n current_epoch = 0\n current_iteration = 0\n print('No checkpoint found.')\n return current_epoch, current_iteration\n resume = opt.phase == 'train' and opt.resume\n current_epoch, current_iteration = self._load_checkpoint(\n checkpoint_path, resume)\n return current_epoch, current_iteration\n\n def _load_checkpoint(self, checkpoint_path, resume=True):\n checkpoint = torch.load(checkpoint_path, map_location=lambda storage, loc: storage)\n self.net_G.load_state_dict(checkpoint['net_G'], strict=False)\n self.net_G_ema.load_state_dict(checkpoint['net_G_ema'], strict=False)\n print('load [net_G] and [net_G_ema] from {}'.format(checkpoint_path))\n if self.opt.phase == 'train' and resume:\n # the checkpoint we provided does not contains \n # the parameters of the optimizer and schdule \n # because we train the model use another code \n # which does not save these parameters\n self.opt_G.load_state_dict(checkpoint['opt_G'])\n self.sch_G.load_state_dict(checkpoint['sch_G'])\n print('load optimizers and schdules from {}'.format(checkpoint_path))\n\n if resume or self.opt.phase == 'test':\n current_epoch = checkpoint['current_epoch']\n current_iteration = checkpoint['current_iteration']\n else:\n current_epoch = 0\n current_iteration = 0\n print('Done with loading the checkpoint.')\n return current_epoch, current_iteration \n\n def start_of_epoch(self, current_epoch):\n r\"\"\"Things to do before an epoch.\n\n Args:\n current_epoch (int): Current number of epoch.\n \"\"\"\n self._start_of_epoch(current_epoch)\n self.current_epoch = current_epoch\n self.start_epoch_time = time.time()\n\n def start_of_iteration(self, data, current_iteration):\n r\"\"\"Things to do before an iteration.\n\n Args:\n data (dict): Data used for the current iteration.\n current_iteration (int): Current number of iteration.\n \"\"\"\n data = self._start_of_iteration(data, current_iteration)\n data = to_cuda(data)\n self.current_iteration = current_iteration\n if not self.is_inference:\n self.net_G.train()\n self.start_iteration_time = time.time()\n return data\n\n def end_of_iteration(self, data, current_epoch, current_iteration):\n r\"\"\"Things to do after an iteration.\n\n Args:\n data (dict): Data used for the current iteration.\n current_epoch (int): Current number of epoch.\n current_iteration (int): Current number of iteration.\n \"\"\"\n self.current_iteration = current_iteration\n self.current_epoch = current_epoch\n # Update the learning rate policy for the generator if operating in the\n # iteration mode.\n if self.opt.gen_optimizer.lr_policy.iteration_mode:\n self.sch_G.step()\n\n # Accumulate time\n # torch.cuda.synchronize()\n self.elapsed_iteration_time += time.time() - self.start_iteration_time\n # Logging.\n if current_iteration % self.opt.logging_iter == 0:\n ave_t = self.elapsed_iteration_time / self.opt.logging_iter\n self.time_iteration = ave_t\n print('Iteration: {}, average iter time: '\n '{:6f}.'.format(current_iteration, ave_t))\n self.elapsed_iteration_time = 0\n\n if getattr(self.opt, 'speed_benchmark', False):\n # Below code block only needed when analyzing computation\n # bottleneck.\n print('\\tGenerator FWD time {:6f}'.format(\n self.accu_gen_forw_iter_time / self.opt.logging_iter))\n print('\\tGenerator LOS time {:6f}'.format(\n self.accu_gen_loss_iter_time / self.opt.logging_iter))\n print('\\tGenerator BCK time {:6f}'.format(\n self.accu_gen_back_iter_time / self.opt.logging_iter))\n print('\\tGenerator STP time {:6f}'.format(\n self.accu_gen_step_iter_time / self.opt.logging_iter))\n print('\\tGenerator AVG time {:6f}'.format(\n self.accu_gen_avg_iter_time / self.opt.logging_iter))\n print('{:6f}'.format(ave_t))\n\n self.accu_gen_forw_iter_time = 0\n self.accu_gen_loss_iter_time = 0\n self.accu_gen_back_iter_time = 0\n self.accu_gen_step_iter_time = 0\n self.accu_gen_avg_iter_time = 0\n\n self._end_of_iteration(data, current_epoch, current_iteration)\n # Save everything to the checkpoint.\n if current_iteration >= self.opt.snapshot_save_start_iter and \\\n current_iteration % self.opt.snapshot_save_iter == 0:\n self.save_image(self._get_save_path('image', 'jpg'), data)\n self.save_checkpoint(current_epoch, current_iteration)\n self.write_metrics(data)\n # Compute image to be saved.\n elif current_iteration % self.opt.image_save_iter == 0:\n self.save_image(self._get_save_path('image', 'jpg'), data)\n\n if current_iteration % self.opt.logging_iter == 0:\n self._write_tensorboard()\n self._print_current_errors()\n\n\n def _print_current_errors(self):\n epoch, iteration = self.current_epoch, self.current_iteration\n message = '(epoch: %d, iters: %d) ' % (epoch, iteration)\n for loss_name, losses in self.gen_losses.items():\n full_loss_name = 'gen_update' + '/' + loss_name\n message += '%s: %.3f ' % (full_loss_name, losses)\n\n print(message)\n log_name = os.path.join(self.opt.logdir, 'loss_log.txt')\n with open(log_name, \"a\") as log_file:\n log_file.write('%s\\n' % message)\n\n def end_of_epoch(self, data, val_dataset, current_epoch, current_iteration):\n r\"\"\"Things to do after an epoch.\n\n Args:\n data (dict): Data used for the current iteration.\n\n current_epoch (int): Current number of epoch.\n current_iteration (int): Current number of iteration.\n \"\"\"\n # Update the learning rate policy for the generator if operating in the\n # epoch mode.\n self.current_iteration = current_iteration\n self.current_epoch = current_epoch\n if not self.opt.gen_optimizer.lr_policy.iteration_mode:\n self.sch_G.step()\n\n elapsed_epoch_time = time.time() - self.start_epoch_time\n # Logging.\n print('Epoch: {}, total time: {:6f}.'.format(current_epoch,\n elapsed_epoch_time))\n self.time_epoch = elapsed_epoch_time\n self._end_of_epoch(data, current_epoch, current_iteration)\n # Save everything to the checkpoint.\n if current_epoch >= self.opt.snapshot_save_start_epoch and \\\n current_epoch % self.opt.snapshot_save_epoch == 0:\n self.save_image(self._get_save_path('image', 'jpg'), data)\n self.save_checkpoint(current_epoch, current_iteration)\n self.write_metrics(data)\n if self.current_epoch % self.opt.eval_epoch == 0 and self.current_epoch >= self.opt.start_eval_epoch:\n self.eval(val_dataset)\n \n\n # def eval(self, val_dataset):\n # output_dir = os.path.join(\n # self.opt.logdir, 'evaluation',\n # 'epoch_{:05}_iteration_{:09}'.format(self.current_epoch, self.current_iteration)\n # ) \n # os.makedirs(output_dir, exist_ok=True)\n # lpips = self.test(val_dataset, output_dir, self.current_iteration)\n # self.write_data_tensorboard({'test_lpips': lpips.mean()},\n # self.current_epoch, self.current_iteration)\n \n\n def write_data_tensorboard(self, data, epoch, iteration):\n for name, value in data.items():\n full_name = 'eval/' + name\n if full_name not in self.meters.keys():\n # Create a new meter if it doesn't exist.\n self.meters[full_name] = Meter(full_name)\n self.meters[full_name].write(value)\n self.meters[full_name].flush(iteration)\n\n # def pre_process(self, data):\n # r\"\"\"Custom data pre-processing function. Utilize this function if you\n # need to preprocess your data before sending it to the generator and\n # discriminator.\n\n # Args:\n # data (dict): Data used for the current iteration.\n # \"\"\"\n\n\n def save_image(self, path, data):\n r\"\"\"Compute visualization images and save them to the disk.\n\n Args:\n path (str): Location of the file.\n data (dict): Data used for the current iteration.\n \"\"\"\n self.net_G.eval()\n vis_images = self._get_visualizations(data)\n if is_master() and vis_images is not None:\n vis_images = (vis_images + 1) / 2\n print('Save output images to {}'.format(path))\n vis_images.clamp_(0, 1)\n os.makedirs(os.path.dirname(path), exist_ok=True)\n image_grid = torchvision.utils.make_grid(\n vis_images, nrow=1, padding=0, normalize=False)\n if self.opt.trainer.image_to_tensorboard:\n self.image_meter.write_image(image_grid, self.current_iteration)\n torchvision.utils.save_image(image_grid, path, nrow=1)\n\n def write_metrics(self, data):\n r\"\"\"Write metrics to the tensorboard.\"\"\"\n cur_metrics = self._compute_metrics(data, self.current_iteration)\n if cur_metrics is not None:\n if self.best_lpips is not None:\n self.best_lpips = min(self.best_lpips, cur_metrics['lpips'])\n else:\n self.best_lpips = cur_metrics['lpips']\n metric_dict = {\n 'metric/lpips': cur_metrics['lpips'], 'metric/best_lpips': self.best_lpips\n }\n self._write_to_meters(metric_dict, self.meters)\n self._flush_meters(self.meters)\n if self.opt.trainer.hparam_to_tensorboard:\n add_hparams(self.hparam_dict, metric_dict)\n\n def _get_save_path(self, subdir, ext):\n r\"\"\"Get the image save path.\n\n Args:\n subdir (str): Sub-directory under the main directory for saving\n the outputs.\n ext (str): Filename extension for the image (e.g., jpg, png, ...).\n Return:\n (str): image filename to be used to save the visualization results.\n \"\"\"\n subdir_path = os.path.join(self.opt.logdir, subdir)\n if not os.path.exists(subdir_path):\n os.makedirs(subdir_path, exist_ok=True)\n return os.path.join(\n subdir_path, 'epoch_{:05}_iteration_{:09}.{}'.format(\n self.current_epoch, self.current_iteration, ext))\n\n\n def _compute_metrics(self, data, current_iteration):\n r\"\"\"Return the evaluation result.\n \"\"\"\n return None\n\n\n def _start_of_epoch(self, current_epoch):\n r\"\"\"Operations to do before starting an epoch.\n\n Args:\n current_epoch (int): Current number of epoch.\n \"\"\"\n pass\n\n def _start_of_iteration(self, data, current_iteration):\n r\"\"\"Operations to do before starting an iteration.\n\n Args:\n data (dict): Data used for the current iteration.\n current_iteration (int): Current epoch number.\n Returns:\n (dict): Data used for the current iteration. They might be\n processed by the custom _start_of_iteration function.\n \"\"\"\n return data\n\n def _end_of_iteration(self, data, current_epoch, current_iteration):\n r\"\"\"Operations to do after an iteration.\n\n Args:\n data (dict): Data used for the current iteration.\n current_epoch (int): Current number of epoch.\n current_iteration (int): Current epoch number.\n \"\"\"\n pass\n \n\n def _end_of_epoch(self, data, current_epoch, current_iteration):\n r\"\"\"Operations to do after an epoch.\n\n Args:\n data (dict): Data used for the current iteration.\n current_epoch (int): Current number of epoch.\n current_iteration (int): Current epoch number.\n \"\"\"\n pass\n\n def _get_visualizations(self, data):\n r\"\"\"Compute visualization outputs.\n\n Args:\n data (dict): Data used for the current iteration.\n \"\"\"\n return None\n\n def _init_loss(self, opt):\n r\"\"\"Every trainer should implement its own init loss function.\"\"\"\n raise NotImplementedError\n\n def gen_forward(self, data):\n r\"\"\"Every trainer should implement its own generator forward.\"\"\"\n raise NotImplementedError\n\n\n def test(self, data_loader, output_dir, current_iteration):\n r\"\"\"Compute results images for a batch of input data and save the\n results in the specified folder.\n\n Args:\n data_loader (torch.utils.data.DataLoader): PyTorch dataloader.\n output_dir (str): Target location for saving the output image.\n \"\"\"\n raise NotImplementedError\n\n # def _get_total_loss(self, gen_forward):\n # r\"\"\"Return the total loss to be backpropagated.\n # Args:\n # gen_forward (bool): If ``True``, backpropagates the generator loss,\n # otherwise the discriminator loss.\n # \"\"\"\n # losses = self.gen_losses if gen_forward else self.dis_losses\n # total_loss = torch.tensor(0., device=torch.device('cuda'))\n # # Iterates over all possible losses.\n # for loss_name in self.weights:\n # # If it is for the current model (gen/dis).\n # if loss_name in losses:\n # # Multiply it with the corresponding weight\n # # and add it to the total loss.\n # total_loss += losses[loss_name] * self.weights[loss_name]\n # losses['total'] = total_loss # logging purpose\n # return total_loss\n\n # def _detach_losses(self):\n # r\"\"\"Detach all logging variables to prevent potential memory leak.\"\"\"\n # for loss_name in self.gen_losses:\n # self.gen_losses[loss_name] = self.gen_losses[loss_name].detach()\n # for loss_name in self.dis_losses:\n # self.dis_losses[loss_name] = self.dis_losses[loss_name].detach()\n\n # def _time_before_forward(self):\n # r\"\"\"\n # Record time before applying forward.\n # \"\"\"\n # if getattr(self.opt, 'speed_benchmark', False):\n # torch.cuda.synchronize()\n # self.forw_time = time.time()\n\n # def _time_before_loss(self):\n # r\"\"\"\n # Record time before computing loss.\n # \"\"\"\n # if getattr(self.opt, 'speed_benchmark', False):\n # torch.cuda.synchronize()\n # self.loss_time = time.time()\n\n # def _time_before_backward(self):\n # r\"\"\"\n # Record time before applying backward.\n # \"\"\"\n # if getattr(self.opt, 'speed_benchmark', False):\n # torch.cuda.synchronize()\n # self.back_time = time.time()\n\n # def _time_before_step(self):\n # r\"\"\"\n # Record time before updating the weights\n # \"\"\"\n # if getattr(self.opt, 'speed_benchmark', False):\n # torch.cuda.synchronize()\n # self.step_time = time.time()\n\n # def _time_before_model_avg(self):\n # r\"\"\"\n # Record time before applying model average.\n # \"\"\"\n # if getattr(self.opt, 'speed_benchmark', False):\n # torch.cuda.synchronize()\n # self.avg_time = time.time()\n\n # def _time_before_leave_gen(self):\n # r\"\"\"\n # Record forward, backward, loss, and model average time for the\n # generator update.\n # \"\"\"\n # if getattr(self.opt, 'speed_benchmark', False):\n # torch.cuda.synchronize()\n # end_time = time.time()\n # self.accu_gen_forw_iter_time += self.loss_time - self.forw_time\n # self.accu_gen_loss_iter_time += self.back_time - self.loss_time\n # self.accu_gen_back_iter_time += self.step_time - self.back_time\n # self.accu_gen_step_iter_time += self.avg_time - self.step_time\n # self.accu_gen_avg_iter_time += end_time - self.avg_time\n\n # def _time_before_leave_dis(self):\n # r\"\"\"\n # Record forward, backward, loss time for the discriminator update.\n # \"\"\"\n # if getattr(self.opt, 'speed_benchmark', False):\n # torch.cuda.synchronize()\n # end_time = time.time()\n # self.accu_dis_forw_iter_time += self.loss_time - self.forw_time\n # self.accu_dis_loss_iter_time += self.back_time - self.loss_time\n # self.accu_dis_back_iter_time += self.step_time - self.back_time\n # self.accu_dis_step_iter_time += end_time - self.step_time\n\n\n@master_only\ndef _save_checkpoint(opt,\n net_G, net_G_ema, opt_G, sch_G,\n current_epoch, current_iteration):\n r\"\"\"Save network weights, optimizer parameters, scheduler parameters\n in the checkpoint.\n\n Args:\n opt (obj): Global configuration.\n opt_G (obj): Optimizer for the generator network.\n sch_G (obj): Scheduler for the generator optimizer.\n current_epoch (int): Current epoch.\n current_iteration (int): Current iteration.\n \"\"\"\n latest_checkpoint_path = 'epoch_{:05}_iteration_{:09}_checkpoint.pt'.format(\n current_epoch, current_iteration)\n save_path = os.path.join(opt.logdir, latest_checkpoint_path)\n torch.save(\n {\n 'net_G': net_G.state_dict(),\n 'net_G_ema': net_G_ema.state_dict(),\n 'opt_G': opt_G.state_dict(),\n 'sch_G': sch_G.state_dict(),\n 'current_epoch': current_epoch,\n 'current_iteration': current_iteration,\n },\n save_path,\n )\n fn = os.path.join(opt.logdir, 'latest_checkpoint.txt')\n with open(fn, 'wt') as f:\n f.write('latest_checkpoint: %s' % latest_checkpoint_path)\n print('Save checkpoint to {}'.format(save_path))\n return save_path\n\n\n"
},
{
"path": "trainers/face_trainer.py",
"content": "import math\n\nimport torch\n\nfrom trainers.base import BaseTrainer\nfrom util.trainer import accumulate, get_optimizer\nfrom loss.perceptual import PerceptualLoss\n\nclass FaceTrainer(BaseTrainer):\n r\"\"\"Initialize lambda model trainer.\n\n Args:\n cfg (obj): Global configuration.\n net_G (obj): Generator network.\n opt_G (obj): Optimizer for the generator network.\n sch_G (obj): Scheduler for the generator optimizer.\n train_data_loader (obj): Train data loader.\n val_data_loader (obj): Validation data loader.\n \"\"\"\n\n def __init__(self, opt, net_G, opt_G, sch_G,\n train_data_loader, val_data_loader=None):\n super(FaceTrainer, self).__init__(opt, net_G, opt_G, sch_G, train_data_loader, val_data_loader)\n self.accum = 0.5 ** (32 / (10 * 1000))\n self.log_size = int(math.log(opt.data.resolution, 2))\n\n def _init_loss(self, opt):\n self._assign_criteria(\n 'perceptual_warp',\n PerceptualLoss(\n network=opt.trainer.vgg_param_warp.network,\n layers=opt.trainer.vgg_param_warp.layers,\n num_scales=getattr(opt.trainer.vgg_param_warp, 'num_scales', 1),\n use_style_loss=getattr(opt.trainer.vgg_param_warp, 'use_style_loss', False),\n weight_style_to_perceptual=getattr(opt.trainer.vgg_param_warp, 'style_to_perceptual', 0)\n ).to('cuda'),\n opt.trainer.loss_weight.weight_perceptual_warp)\n\n self._assign_criteria(\n 'perceptual_final',\n PerceptualLoss(\n network=opt.trainer.vgg_param_final.network,\n layers=opt.trainer.vgg_param_final.layers,\n num_scales=getattr(opt.trainer.vgg_param_final, 'num_scales', 1),\n use_style_loss=getattr(opt.trainer.vgg_param_final, 'use_style_loss', False),\n weight_style_to_perceptual=getattr(opt.trainer.vgg_param_final, 'style_to_perceptual', 0)\n ).to('cuda'),\n opt.trainer.loss_weight.weight_perceptual_final)\n\n def _assign_criteria(self, name, criterion, weight):\n self.criteria[name] = criterion\n self.weights[name] = weight\n\n def optimize_parameters(self, data):\n self.gen_losses = {}\n source_image, target_image = data['source_image'], data['target_image']\n source_semantic, target_semantic = data['source_semantics'], data['target_semantics']\n\n input_image = torch.cat((source_image, target_image), 0)\n input_semantic = torch.cat((target_semantic, source_semantic), 0)\n gt_image = torch.cat((target_image, source_image), 0) \n\n output_dict = self.net_G(input_image, input_semantic, self.training_stage)\n\n if self.training_stage == 'gen':\n fake_img = output_dict['fake_image']\n warp_img = output_dict['warp_image']\n self.gen_losses[\"perceptual_final\"] = self.criteria['perceptual_final'](fake_img, gt_image)\n self.gen_losses[\"perceptual_warp\"] = self.criteria['perceptual_warp'](warp_img, gt_image)\n else:\n warp_img = output_dict['warp_image']\n self.gen_losses[\"perceptual_warp\"] = self.criteria['perceptual_warp'](warp_img, gt_image)\n\n total_loss = 0\n for key in self.gen_losses:\n self.gen_losses[key] = self.gen_losses[key] * self.weights[key]\n total_loss += self.gen_losses[key]\n\n self.gen_losses['total_loss'] = total_loss\n\n self.net_G.zero_grad()\n total_loss.backward()\n self.opt_G.step()\n\n accumulate(self.net_G_ema, self.net_G_module, self.accum)\n\n def _start_of_iteration(self, data, current_iteration):\n self.training_stage = 'gen' if current_iteration >= self.opt.trainer.pretrain_warp_iteration else 'warp'\n if current_iteration == self.opt.trainer.pretrain_warp_iteration:\n self.reset_trainer()\n return data\n\n def reset_trainer(self):\n self.opt_G = get_optimizer(self.opt.gen_optimizer, self.net_G.module)\n\n def _get_visualizations(self, data):\n source_image, target_image = data['source_image'], data['target_image']\n source_semantic, target_semantic = data['source_semantics'], data['target_semantics']\n\n input_image = torch.cat((source_image, target_image), 0)\n input_semantic = torch.cat((target_semantic, source_semantic), 0) \n with torch.no_grad():\n self.net_G_ema.eval()\n output_dict = self.net_G_ema(\n input_image, input_semantic, self.training_stage\n )\n if self.training_stage == 'gen':\n fake_img = torch.cat([output_dict['warp_image'], output_dict['fake_image']], 3)\n else:\n fake_img = output_dict['warp_image']\n\n fake_source, fake_target = torch.chunk(fake_img, 2, dim=0)\n sample_source = torch.cat([source_image, fake_source, target_image], 3)\n sample_target = torch.cat([target_image, fake_target, source_image], 3) \n sample = torch.cat([sample_source, sample_target], 2)\n sample = torch.cat(torch.chunk(sample, sample.size(0), 0)[:3], 2)\n return sample\n\n def test(self, data_loader, output_dir, current_iteration=-1):\n pass\n\n def _compute_metrics(self, data, current_iteration):\n if self.training_stage == 'gen':\n source_image, target_image = data['source_image'], data['target_image']\n source_semantic, target_semantic = data['source_semantics'], data['target_semantics']\n\n input_image = torch.cat((source_image, target_image), 0)\n input_semantic = torch.cat((target_semantic, source_semantic), 0) \n gt_image = torch.cat((target_image, source_image), 0) \n metrics = {}\n with torch.no_grad():\n self.net_G_ema.eval()\n output_dict = self.net_G_ema(\n input_image, input_semantic, self.training_stage\n )\n fake_image = output_dict['fake_image']\n metrics['lpips'] = self.lpips(fake_image, gt_image).mean()\n return metrics"
},
{
"path": "util/cudnn.py",
"content": "import torch.backends.cudnn as cudnn\n\nfrom util.distributed import master_only_print as print\n\n\ndef init_cudnn(deterministic, benchmark):\n r\"\"\"Initialize the cudnn module. The two things to consider is whether to\n use cudnn benchmark and whether to use cudnn deterministic. If cudnn\n benchmark is set, then the cudnn deterministic is automatically false.\n\n Args:\n deterministic (bool): Whether to use cudnn deterministic.\n benchmark (bool): Whether to use cudnn benchmark.\n \"\"\"\n cudnn.deterministic = deterministic\n cudnn.benchmark = benchmark\n print('cudnn benchmark: {}'.format(benchmark))\n print('cudnn deterministic: {}'.format(deterministic))\n"
},
{
"path": "util/distributed.py",
"content": "import functools\n\nimport torch\nimport torch.distributed as dist\n\ndef init_dist(local_rank, backend='nccl', **kwargs):\n r\"\"\"Initialize distributed training\"\"\"\n if dist.is_available():\n if dist.is_initialized():\n return torch.cuda.current_device()\n torch.cuda.set_device(local_rank)\n dist.init_process_group(backend=backend, init_method='env://', **kwargs)\n \n\ndef get_rank():\n r\"\"\"Get rank of the thread.\"\"\"\n rank = 0\n if dist.is_available():\n if dist.is_initialized():\n rank = dist.get_rank()\n return rank\n\n\ndef get_world_size():\n r\"\"\"Get world size. How many GPUs are available in this job.\"\"\"\n world_size = 1\n if dist.is_available():\n if dist.is_initialized():\n world_size = dist.get_world_size()\n return world_size\n\n\ndef master_only(func):\n r\"\"\"Apply this function only to the master GPU.\"\"\"\n @functools.wraps(func)\n def wrapper(*args, **kwargs):\n r\"\"\"Simple function wrapper for the master function\"\"\"\n if get_rank() == 0:\n return func(*args, **kwargs)\n else:\n return None\n return wrapper\n\n\ndef is_master():\n r\"\"\"check if current process is the master\"\"\"\n return get_rank() == 0\n\n\n@master_only\ndef master_only_print(*args):\n r\"\"\"master-only print\"\"\"\n print(*args)\n\n\ndef dist_reduce_tensor(tensor):\n r\"\"\" Reduce to rank 0 \"\"\"\n world_size = get_world_size()\n if world_size < 2:\n return tensor\n with torch.no_grad():\n dist.reduce(tensor, dst=0)\n if get_rank() == 0:\n tensor /= world_size\n return tensor\n\n\ndef dist_all_reduce_tensor(tensor):\n r\"\"\" Reduce to all ranks \"\"\"\n world_size = get_world_size()\n if world_size < 2:\n return tensor\n with torch.no_grad():\n dist.all_reduce(tensor)\n tensor.div_(world_size)\n return tensor\n\n\ndef dist_all_gather_tensor(tensor):\n r\"\"\" gather to all ranks \"\"\"\n world_size = get_world_size()\n if world_size < 2:\n return [tensor]\n tensor_list = [\n torch.ones_like(tensor) for _ in range(dist.get_world_size())]\n with torch.no_grad():\n dist.all_gather(tensor_list, tensor)\n return tensor_list\n"
},
{
"path": "util/flow_util.py",
"content": "import torch\n\ndef convert_flow_to_deformation(flow):\n r\"\"\"convert flow fields to deformations.\n\n Args:\n flow (tensor): Flow field obtained by the model\n Returns:\n deformation (tensor): The deformation used for warpping\n \"\"\"\n b,c,h,w = flow.shape\n flow_norm = 2 * torch.cat([flow[:,:1,...]/(w-1),flow[:,1:,...]/(h-1)], 1)\n grid = make_coordinate_grid(flow)\n deformation = grid + flow_norm.permute(0,2,3,1)\n return deformation\n\ndef make_coordinate_grid(flow):\n r\"\"\"obtain coordinate grid with the same size as the flow filed.\n\n Args:\n flow (tensor): Flow field obtained by the model\n Returns:\n grid (tensor): The grid with the same size as the input flow\n \"\"\" \n b,c,h,w = flow.shape\n\n x = torch.arange(w).to(flow)\n y = torch.arange(h).to(flow)\n\n x = (2 * (x / (w - 1)) - 1)\n y = (2 * (y / (h - 1)) - 1)\n\n yy = y.view(-1, 1).repeat(1, w)\n xx = x.view(1, -1).repeat(h, 1)\n\n meshed = torch.cat([xx.unsqueeze_(2), yy.unsqueeze_(2)], 2)\n meshed = meshed.expand(b, -1, -1, -1)\n return meshed \n\n \ndef warp_image(source_image, deformation):\n r\"\"\"warp the input image according to the deformation\n\n Args:\n source_image (tensor): source images to be warpped\n deformation (tensor): deformations used to warp the images; value in range (-1, 1)\n Returns:\n output (tensor): the warpped images\n \"\"\" \n _, h_old, w_old, _ = deformation.shape\n _, _, h, w = source_image.shape\n if h_old != h or w_old != w:\n deformation = deformation.permute(0, 3, 1, 2)\n deformation = torch.nn.functional.interpolate(deformation, size=(h, w), mode='bilinear')\n deformation = deformation.permute(0, 2, 3, 1)\n return torch.nn.functional.grid_sample(source_image, deformation) "
},
{
"path": "util/init_weight.py",
"content": "from torch.nn import init\n\n\ndef weights_init(init_type='normal', gain=0.02, bias=None):\n r\"\"\"Initialize weights in the network.\n\n Args:\n init_type (str): The name of the initialization scheme.\n gain (float): The parameter that is required for the initialization\n scheme.\n bias (object): If not ``None``, specifies the initialization parameter\n for bias.\n\n Returns:\n (obj): init function to be applied.\n \"\"\"\n\n def init_func(m):\n r\"\"\"Init function\n\n Args:\n m: module to be weight initialized.\n \"\"\"\n class_name = m.__class__.__name__\n if hasattr(m, 'weight') and (\n class_name.find('Conv') != -1 or\n class_name.find('Linear') != -1 or\n class_name.find('Embedding') != -1):\n if init_type == 'normal':\n init.normal_(m.weight.data, 0.0, gain)\n elif init_type == 'xavier':\n init.xavier_normal_(m.weight.data, gain=gain)\n elif init_type == 'xavier_uniform':\n init.xavier_uniform_(m.weight.data, gain=1.0)\n elif init_type == 'kaiming':\n init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')\n elif init_type == 'orthogonal':\n init.orthogonal_(m.weight.data, gain=gain)\n elif init_type == 'none':\n m.reset_parameters()\n else:\n raise NotImplementedError(\n 'initialization method [%s] is '\n 'not implemented' % init_type)\n if hasattr(m, 'bias') and m.bias is not None:\n if bias is not None:\n bias_type = getattr(bias, 'type', 'normal')\n if bias_type == 'normal':\n bias_gain = getattr(bias, 'gain', 0.5)\n init.normal_(m.bias.data, 0.0, bias_gain)\n else:\n raise NotImplementedError(\n 'initialization method [%s] is '\n 'not implemented' % bias_type)\n else:\n init.constant_(m.bias.data, 0.0)\n return init_func\n"
},
{
"path": "util/io.py",
"content": "import os\n\nimport requests\nimport torch.distributed as dist\nimport torchvision.utils\n\nfrom util.distributed import is_master\n\n\ndef save_pilimage_in_jpeg(fullname, output_img):\n r\"\"\"Save PIL Image to JPEG.\n\n Args:\n fullname (str): Full save path.\n output_img (PIL Image): Image to be saved.\n \"\"\"\n dirname = os.path.dirname(fullname)\n os.makedirs(dirname, exist_ok=True)\n output_img.save(fullname, 'JPEG', quality=99)\n\n\ndef save_intermediate_training_results(\n visualization_images, logdir, current_epoch, current_iteration):\n r\"\"\"Save intermediate training results for debugging purpose.\n\n Args:\n visualization_images (tensor): Image where pixel values are in [-1, 1].\n logdir (str): Where to save the image.\n current_epoch (int): Current training epoch.\n current_iteration (int): Current training iteration.\n \"\"\"\n visualization_images = (visualization_images + 1) / 2\n output_filename = os.path.join(\n logdir, 'images',\n 'epoch_{:05}iteration{:09}.jpg'.format(\n current_epoch, current_iteration))\n print('Save output images to {}'.format(output_filename))\n os.makedirs(os.path.dirname(output_filename), exist_ok=True)\n image_grid = torchvision.utils.make_grid(\n visualization_images.data, nrow=1, padding=0, normalize=False)\n torchvision.utils.save_image(image_grid, output_filename, nrow=1)\n\n\ndef download_file_from_google_drive(file_id, destination):\n r\"\"\"Download a file from the google drive by using the file ID.\n\n Args:\n file_id: Google drive file ID\n destination: Path to save the file.\n\n Returns:\n\n \"\"\"\n URL = \"https://docs.google.com/uc?export=download\"\n session = requests.Session()\n response = session.get(URL, params={'id': file_id}, stream=True)\n token = get_confirm_token(response)\n if token:\n params = {'id': file_id, 'confirm': token}\n response = session.get(URL, params=params, stream=True)\n save_response_content(response, destination)\n\n\ndef get_confirm_token(response):\n r\"\"\"Get confirm token\n\n Args:\n response: Check if the file exists.\n\n Returns:\n\n \"\"\"\n for key, value in response.cookies.items():\n if key.startswith('download_warning'):\n return value\n return None\n\n\ndef save_response_content(response, destination):\n r\"\"\"Save response content\n\n Args:\n response:\n destination: Path to save the file.\n\n Returns:\n\n \"\"\"\n chunk_size = 32768\n with open(destination, \"wb\") as f:\n for chunk in response.iter_content(chunk_size):\n if chunk:\n f.write(chunk)\n\n\ndef get_checkpoint(checkpoint_path, url=''):\n r\"\"\"Get the checkpoint path. If it does not exist yet, download it from\n the url.\n\n Args:\n checkpoint_path (str): Checkpoint path.\n url (str): URL to download checkpoint.\n Returns:\n (str): Full checkpoint path.\n \"\"\"\n if 'TORCH_HOME' not in os.environ:\n os.environ['TORCH_HOME'] = os.getcwd()\n save_dir = os.path.join(os.environ['TORCH_HOME'], 'checkpoints')\n os.makedirs(save_dir, exist_ok=True)\n full_checkpoint_path = os.path.join(save_dir, checkpoint_path)\n if not os.path.exists(full_checkpoint_path):\n os.makedirs(os.path.dirname(full_checkpoint_path), exist_ok=True)\n if is_master():\n print('Download {}'.format(url))\n download_file_from_google_drive(url, full_checkpoint_path)\n if dist.is_available() and dist.is_initialized():\n dist.barrier()\n return full_checkpoint_path\n"
},
{
"path": "util/logging.py",
"content": "import os\nimport datetime\n\nfrom util.meters import set_summary_writer\nfrom util.distributed import master_only_print as print\nfrom util.distributed import master_only\n\ndef get_date_uid():\n \"\"\"Generate a unique id based on date.\n Returns:\n str: Return uid string, e.g. '20171122171307111552'.\n \"\"\"\n return str(datetime.datetime.now().strftime(\"%Y_%m%d_%H%M_%S\"))\n\n\ndef init_logging(opt):\n date_uid = get_date_uid()\n if opt.name is not None:\n logdir = os.path.join(opt.checkpoints_dir, opt.name)\n else:\n logdir = os.path.join(opt.checkpoints_dir, date_uid)\n opt.logdir = logdir\n return date_uid, logdir\n \n@master_only\ndef make_logging_dir(logdir, date_uid):\n r\"\"\"Create the logging directory\n\n Args:\n logdir (str): Log directory name\n \"\"\"\n\n \n print('Make folder {}'.format(logdir))\n os.makedirs(logdir, exist_ok=True)\n tensorboard_dir = os.path.join(logdir, 'tensorboard')\n image_dir = os.path.join(logdir, 'image')\n eval_dir = os.path.join(logdir, 'evaluation')\n os.makedirs(tensorboard_dir, exist_ok=True)\n os.makedirs(image_dir, exist_ok=True)\n os.makedirs(eval_dir, exist_ok=True)\n\n set_summary_writer(tensorboard_dir)\n loss_log_name = os.path.join(logdir, 'loss_log.txt')\n with open(loss_log_name, \"a\") as log_file:\n log_file.write('================ Training Loss (%s) ================\\n' % date_uid)\n"
},
{
"path": "util/lpips.py",
"content": "import os \nimport glob\nimport numpy as np\nfrom imageio import imread\n\nimport torch\n\nfrom third_part.PerceptualSimilarity.models import dist_model as dm\n\ndef get_image_list(flist):\n if isinstance(flist, list):\n return flist\n\n # flist: image file path, image directory path, text file flist path\n if isinstance(flist, str):\n if os.path.isdir(flist):\n flist = list(glob.glob(flist + '/*.jpg')) + list(glob.glob(flist + '/*.png'))\n flist.sort()\n return flist\n\n if os.path.isfile(flist):\n try:\n return np.genfromtxt(flist, dtype=np.str)\n except:\n return [flist]\n print('can not read files from %s return empty list'%flist)\n return []\n\ndef preprocess_path_for_deform_task(gt_path, distorted_path):\n distorted_image_list = sorted(get_image_list(distorted_path))\n gt_list=[]\n distorated_list=[]\n\n for distorted_image in distorted_image_list:\n image = os.path.basename(distorted_image)\n image = image.split('_2_')[-1]\n image = image.split('_vis')[0] +'.jpg'\n gt_image = os.path.join(gt_path, image)\n if not os.path.isfile(gt_image):\n gt_image = gt_image.replace('.jpg', '.png')\n gt_list.append(gt_image)\n distorated_list.append(distorted_image)\n return gt_list, distorated_list\n\nclass LPIPS():\n def __init__(self, use_gpu=True):\n self.model = dm.DistModel()\n self.model.initialize(model='net-lin', net='alex', use_gpu=use_gpu)\n self.use_gpu=use_gpu\n\n def __call__(self, image_1, image_2):\n \"\"\"\n image_1: images with size (n, 3, w, h) with value [-1, 1]\n image_2: images with size (n, 3, w, h) with value [-1, 1]\n \"\"\"\n result = self.model.forward(image_1, image_2)\n return result\n\n def calculate_from_disk(self, gt_path, distorted_path, batch_size=64, verbose=False, for_deformation=True):\n # if sort:\n if for_deformation:\n files_1, files_2 = preprocess_path_for_deform_task(gt_path, distorted_path)\n else:\n files_1 = sorted(get_image_list(gt_path))\n files_2 = sorted(get_image_list(distorted_path))\n\n new_files_1, new_files_2 = [], []\n for item1,item2 in zip(files_1, files_2):\n if os.path.isfile(item1) and os.path.isfile(item2):\n new_files_1.append(item1)\n new_files_2.append(item2)\n else:\n print(item2)\n imgs_1 = np.array([imread(str(fn)).astype(np.float32)/127.5-1 for fn in new_files_1])\n imgs_2 = np.array([imread(str(fn)).astype(np.float32)/127.5-1 for fn in new_files_2])\n\n # Bring images to shape (B, 3, H, W)\n imgs_1 = imgs_1.transpose((0, 3, 1, 2))\n imgs_2 = imgs_2.transpose((0, 3, 1, 2))\n\n result=[]\n\n\n d0 = imgs_1.shape[0]\n if batch_size > d0:\n print(('Warning: batch size is bigger than the data size. '\n 'Setting batch size to data size'))\n batch_size = d0\n\n n_batches = d0 // batch_size\n n_used_imgs = n_batches * batch_size\n\n # imgs_1_arr = np.empty((n_used_imgs, self.dims))\n # imgs_2_arr = np.empty((n_used_imgs, self.dims))\n for i in range(n_batches):\n if verbose:\n print('\\rPropagating batch %d/%d' % (i + 1, n_batches))\n # end='', flush=True)\n start = i * batch_size\n end = start + batch_size\n\n img_1_batch = torch.from_numpy(imgs_1[start:end]).type(torch.FloatTensor)\n img_2_batch = torch.from_numpy(imgs_2[start:end]).type(torch.FloatTensor)\n\n if self.use_gpu:\n img_1_batch = img_1_batch.cuda()\n img_2_batch = img_2_batch.cuda()\n\n\n result.append(self.model.forward(img_1_batch, img_2_batch))\n\n\n distance = np.average(result)\n print('lpips: %.3f'%distance)\n return distance"
},
{
"path": "util/meters.py",
"content": "import math\n\nimport torch\nfrom torch.utils.tensorboard import SummaryWriter\nfrom torch.utils.tensorboard.summary import hparams\n\n\nfrom util.distributed import master_only\nfrom util.distributed import master_only_print as print\n\nLOG_WRITER = None\nLOG_DIR = None\n\n\n@torch.no_grad()\ndef sn_reshape_weight_to_matrix(weight):\n r\"\"\"Reshape weight to obtain the matrix form.\n\n Args:\n weight (Parameters): pytorch layer parameter tensor.\n \"\"\"\n weight_mat = weight\n height = weight_mat.size(0)\n return weight_mat.reshape(height, -1)\n\n\n@torch.no_grad()\ndef get_weight_stats(mod, cfg, loss_id):\n r\"\"\"Get weight state\n\n Args:\n mod: Pytorch module\n cfg: Configuration object\n loss_id: Needed when using AMP.\n \"\"\"\n loss_scale = 1.0\n if cfg.trainer.amp == 'O1' or cfg.trainer.amp == 'O2':\n # AMP rescales the gradient so we have to undo it.\n loss_scale = amp._amp_state.loss_scalers[loss_id].loss_scale()\n if mod.weight_orig.grad is not None:\n grad_norm = mod.weight_orig.grad.data.norm().item() / float(loss_scale)\n else:\n grad_norm = 0.\n weight_norm = mod.weight_orig.data.norm().item()\n weight_mat = sn_reshape_weight_to_matrix(mod.weight_orig)\n sigma = torch.sum(mod.weight_u * torch.mv(weight_mat, mod.weight_v))\n return grad_norm, weight_norm, sigma\n\n\n@master_only\ndef set_summary_writer(log_dir):\n r\"\"\"Set summary writer\n\n Args:\n log_dir (str): Log directory.\n \"\"\"\n global LOG_DIR, LOG_WRITER\n LOG_DIR = log_dir\n LOG_WRITER = SummaryWriter(log_dir=log_dir)\n\n\n@master_only\ndef write_summary(name, summary, step, hist=False):\n \"\"\"Utility function for write summary to log_writer.\n \"\"\"\n global LOG_WRITER\n lw = LOG_WRITER\n if lw is None:\n raise Exception(\"Log writer not set.\")\n if hist:\n lw.add_histogram(name, summary, step)\n else:\n lw.add_scalar(name, summary, step)\n\n\n@master_only\ndef add_hparams(hparam_dict=None, metric_dict=None):\n r\"\"\"Add a set of hyperparameters to be compared in tensorboard.\n\n Args:\n hparam_dict (dictionary): Each key-value pair in the dictionary is the\n name of the hyper parameter and it's corresponding value.\n The type of the value can be one of `bool`, `string`, `float`,\n `int`, or `None`.\n metric_dict (dictionary): Each key-value pair in the dictionary is the\n name of the metric and it's corresponding value. Note that the key\n used here should be unique in the tensorboard record. Otherwise the\n value you added by `add_scalar` will be displayed in hparam plugin.\n In most cases, this is unwanted.\n \"\"\"\n if type(hparam_dict) is not dict or type(metric_dict) is not dict:\n raise TypeError('hparam_dict and metric_dict should be dictionary.')\n global LOG_WRITER\n lw = LOG_WRITER\n\n exp, ssi, sei = hparams(hparam_dict, metric_dict)\n\n lw.file_writer.add_summary(exp)\n lw.file_writer.add_summary(ssi)\n lw.file_writer.add_summary(sei)\n\n\nclass Meter(object):\n \"\"\"Meter is to keep track of statistics along steps.\n Meters write values for purpose like printing average values.\n Meters can be flushed to log files (i.e. TensorBoard for now)\n regularly.\n\n Args:\n name (str): the name of meter\n \"\"\"\n\n @master_only\n def __init__(self, name):\n self.name = name\n self.values = []\n\n @master_only\n def reset(self):\n r\"\"\"Reset the meter values\"\"\"\n self.values = []\n\n @master_only\n def write(self, value):\n r\"\"\"Record the value\"\"\"\n self.values.append(value)\n\n @master_only\n def flush(self, step):\n r\"\"\"Write the value in the tensorboard.\n\n Args:\n step (int): Epoch or iteration number.\n \"\"\"\n if not all(math.isfinite(x) for x in self.values):\n print(\"meter {} contained a nan or inf.\".format(self.name))\n filtered_values = list(filter(lambda x: math.isfinite(x), self.values))\n if float(len(filtered_values)) != 0:\n value = float(sum(filtered_values)) / float(len(filtered_values))\n write_summary(self.name, value, step)\n self.reset()\n\n @master_only\n def write_image(self, img_grid, step):\n r\"\"\"Write the value in the tensorboard.\n\n Args:\n img_grid:\n step (int): Epoch or iteration number.\n \"\"\"\n global LOG_WRITER\n lw = LOG_WRITER\n if lw is None:\n raise Exception(\"Log writer not set.\")\n lw.add_image(\"Visualizations\", img_grid, step)\n"
},
{
"path": "util/misc.py",
"content": "\"\"\"Miscellaneous utils.\"\"\"\nfrom collections import OrderedDict\n\nimport numpy as np\nimport torch\nimport torch.nn.functional as F\nfrom scipy.stats import truncnorm\nfrom torch._six import container_abcs, string_classes\n\n\ndef split_labels(labels, label_lengths):\n r\"\"\"Split concatenated labels into their parts.\n\n Args:\n labels (torch.Tensor): Labels obtained through concatenation.\n label_lengths (OrderedDict): Containing order of labels & their lengths.\n\n Returns:\n\n \"\"\"\n assert isinstance(label_lengths, OrderedDict)\n start = 0\n outputs = {}\n for data_type, length in label_lengths.items():\n end = start + length\n if labels.dim() == 5:\n outputs[data_type] = labels[:, :, start:end]\n elif labels.dim() == 4:\n outputs[data_type] = labels[:, start:end]\n elif labels.dim() == 3:\n outputs[data_type] = labels[start:end]\n start = end\n return outputs\n\n\ndef requires_grad(model, require=True):\n r\"\"\" Set a model to require gradient or not.\n\n Args:\n model (nn.Module): Neural network model.\n require (bool): Whether the network requires gradient or not.\n\n Returns:\n\n \"\"\"\n for p in model.parameters():\n p.requires_grad = require\n\n\ndef to_device(data, device):\n r\"\"\"Move all tensors inside data to device.\n\n Args:\n data (dict, list, or tensor): Input data.\n device (str): 'cpu' or 'cuda'.\n \"\"\"\n assert device in ['cpu', 'cuda']\n if isinstance(data, torch.Tensor):\n data = data.to(torch.device(device))\n return data\n elif isinstance(data, container_abcs.Mapping):\n return {key: to_device(data[key], device) for key in data}\n elif isinstance(data, container_abcs.Sequence) and \\\n not isinstance(data, string_classes):\n return [to_device(d, device) for d in data]\n else:\n return data\n\n\ndef to_cuda(data):\n r\"\"\"Move all tensors inside data to gpu.\n\n Args:\n data (dict, list, or tensor): Input data.\n \"\"\"\n return to_device(data, 'cuda')\n\n\ndef to_cpu(data):\n r\"\"\"Move all tensors inside data to cpu.\n\n Args:\n data (dict, list, or tensor): Input data.\n \"\"\"\n return to_device(data, 'cpu')\n\n\ndef to_half(data):\n r\"\"\"Move all floats to half.\n\n Args:\n data (dict, list or tensor): Input data.\n \"\"\"\n if isinstance(data, torch.Tensor) and torch.is_floating_point(data):\n data = data.half()\n return data\n elif isinstance(data, container_abcs.Mapping):\n return {key: to_half(data[key]) for key in data}\n elif isinstance(data, container_abcs.Sequence) and \\\n not isinstance(data, string_classes):\n return [to_half(d) for d in data]\n else:\n return data\n\n\ndef to_float(data):\n r\"\"\"Move all halfs to float.\n\n Args:\n data (dict, list or tensor): Input data.\n \"\"\"\n if isinstance(data, torch.Tensor) and torch.is_floating_point(data):\n data = data.float()\n return data\n elif isinstance(data, container_abcs.Mapping):\n return {key: to_float(data[key]) for key in data}\n elif isinstance(data, container_abcs.Sequence) and \\\n not isinstance(data, string_classes):\n return [to_float(d) for d in data]\n else:\n return data\n\n\ndef get_and_setattr(cfg, name, default):\n r\"\"\"Get attribute with default choice. If attribute does not exist, set it\n using the default value.\n\n Args:\n cfg (obj) : Config options.\n name (str) : Attribute name.\n default (obj) : Default attribute.\n\n Returns:\n (obj) : Desired attribute.\n \"\"\"\n if not hasattr(cfg, name) or name not in cfg.__dict__:\n setattr(cfg, name, default)\n return getattr(cfg, name)\n\n\ndef get_nested_attr(cfg, attr_name, default):\n r\"\"\"Iteratively try to get the attribute from cfg. If not found, return\n default.\n\n Args:\n cfg (obj): Config file.\n attr_name (str): Attribute name (e.g. XXX.YYY.ZZZ).\n default (obj): Default return value for the attribute.\n\n Returns:\n (obj): Attribute value.\n \"\"\"\n names = attr_name.split('.')\n atr = cfg\n for name in names:\n if not hasattr(atr, name):\n return default\n atr = getattr(atr, name)\n return atr\n\n\ndef gradient_norm(model):\n r\"\"\"Return the gradient norm of model.\n\n Args:\n model (PyTorch module): Your network.\n\n \"\"\"\n total_norm = 0\n for p in model.parameters():\n if p.grad is not None:\n param_norm = p.grad.norm(2)\n total_norm += param_norm.item() ** 2\n return total_norm ** (1. / 2)\n\n\ndef random_shift(x, offset=0.05, mode='bilinear', padding_mode='reflection'):\n r\"\"\"Randomly shift the input tensor.\n\n Args:\n x (4D tensor): The input batch of images.\n offset (int): The maximum offset ratio that is between [0, 1].\n The maximum shift is offset * image_size for each direction.\n mode (str): The resample mode for 'F.grid_sample'.\n padding_mode (str): The padding mode for 'F.grid_sample'.\n\n Returns:\n x (4D tensor) : The randomly shifted image.\n \"\"\"\n assert x.dim() == 4, \"Input must be a 4D tensor.\"\n batch_size = x.size(0)\n theta = torch.eye(2, 3, device=x.device).unsqueeze(0).repeat(\n batch_size, 1, 1)\n theta[:, :, 2] = 2 * offset * torch.rand(batch_size, 2) - offset\n grid = F.affine_grid(theta, x.size())\n x = F.grid_sample(x, grid, mode=mode, padding_mode=padding_mode)\n return x\n\n\ndef truncated_gaussian(threshold, size, seed=None, device=None):\n r\"\"\"Apply the truncated gaussian trick to trade diversity for quality\n\n Args:\n threshold (float): Truncation threshold.\n size (list of integer): Tensor size.\n seed (int): Random seed.\n device:\n \"\"\"\n state = None if seed is None else np.random.RandomState(seed)\n values = truncnorm.rvs(-threshold, threshold,\n size=size, random_state=state)\n return torch.tensor(values, device=device).float()\n\n\ndef apply_imagenet_normalization(input):\n r\"\"\"Normalize using ImageNet mean and std.\n\n Args:\n input (4D tensor NxCxHxW): The input images, assuming to be [-1, 1].\n\n Returns:\n Normalized inputs using the ImageNet normalization.\n \"\"\"\n # normalize the input back to [0, 1]\n normalized_input = (input + 1) / 2\n # normalize the input using the ImageNet mean and std\n mean = normalized_input.new_tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)\n std = normalized_input.new_tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)\n output = (normalized_input - mean) / std\n return output\n"
},
{
"path": "util/trainer.py",
"content": "import random\nimport importlib\nimport numpy as np\n\nimport torch\nimport torch.nn as nn\nfrom torch.optim import Adam, lr_scheduler\n\nfrom util.distributed import master_only_print as print\nfrom util.init_weight import weights_init\n\ndef accumulate(model1, model2, decay=0.999):\n par1 = dict(model1.named_parameters())\n par2 = dict(model2.named_parameters())\n\n for k in par1.keys():\n par1[k].data.mul_(decay).add_(par2[k].data, alpha=1 - decay)\n\ndef set_random_seed(seed):\n r\"\"\"Set random seeds for everything.\n\n Args:\n seed (int): Random seed.\n by_rank (bool):\n \"\"\"\n random.seed(seed)\n np.random.seed(seed)\n torch.manual_seed(seed)\n torch.cuda.manual_seed(seed)\n torch.cuda.manual_seed_all(seed)\n\n\n\ndef get_trainer(opt, net_G, net_G_ema, opt_G, sch_G, train_dataset):\n module, trainer_name = opt.trainer.type.split('::')\n\n trainer_lib = importlib.import_module(module)\n trainer_class = getattr(trainer_lib, trainer_name)\n trainer = trainer_class(opt, net_G, net_G_ema, opt_G, sch_G, train_dataset)\n return trainer\n\ndef get_model_optimizer_and_scheduler(opt):\n gen_module, gen_network_name = opt.gen.type.split('::')\n lib = importlib.import_module(gen_module)\n network = getattr(lib, gen_network_name)\n net_G = network(**opt.gen.param).to(opt.device)\n init_bias = getattr(opt.trainer.init, 'bias', None)\n net_G.apply(weights_init(\n opt.trainer.init.type, opt.trainer.init.gain, init_bias))\n\n net_G_ema = network(**opt.gen.param).to(opt.device)\n net_G_ema.eval()\n accumulate(net_G_ema, net_G, 0)\n print('net [{}] parameter count: {:,}'.format(\n 'net_G', _calculate_model_size(net_G)))\n print('Initialize net_G weights using '\n 'type: {} gain: {}'.format(opt.trainer.init.type,\n opt.trainer.init.gain))\n\n\n opt_G = get_optimizer(opt.gen_optimizer, net_G)\n\n if opt.distributed:\n net_G = nn.parallel.DistributedDataParallel(\n net_G,\n device_ids=[opt.local_rank],\n output_device=opt.local_rank,\n broadcast_buffers=False,\n find_unused_parameters=True,\n )\n\n # Scheduler\n sch_G = get_scheduler(opt.gen_optimizer, opt_G)\n return net_G, net_G_ema, opt_G, sch_G\n\n\ndef _calculate_model_size(model):\n r\"\"\"Calculate number of parameters in a PyTorch network.\n\n Args:\n model (obj): PyTorch network.\n\n Returns:\n (int): Number of parameters.\n \"\"\"\n return sum(p.numel() for p in model.parameters() if p.requires_grad)\n\n\ndef get_scheduler(opt_opt, opt):\n \"\"\"Return the scheduler object.\n\n Args:\n opt_opt (obj): Config for the specific optimization module (gen/dis).\n opt (obj): PyTorch optimizer object.\n\n Returns:\n (obj): Scheduler\n \"\"\"\n if opt_opt.lr_policy.type == 'step':\n scheduler = lr_scheduler.StepLR(\n opt,\n step_size=opt_opt.lr_policy.step_size,\n gamma=opt_opt.lr_policy.gamma)\n elif opt_opt.lr_policy.type == 'constant':\n scheduler = lr_scheduler.LambdaLR(opt, lambda x: 1)\n else:\n return NotImplementedError('Learning rate policy {} not implemented.'.\n format(opt_opt.lr_policy.type))\n return scheduler\n\n\ndef get_optimizer(opt_opt, net):\n return get_optimizer_for_params(opt_opt, net.parameters())\n\n\ndef get_optimizer_for_params(opt_opt, params):\n r\"\"\"Return the scheduler object.\n\n Args:\n opt_opt (obj): Config for the specific optimization module (gen/dis).\n params (obj): Parameters to be trained by the parameters.\n\n Returns:\n (obj): Optimizer\n \"\"\"\n # We will use fuse optimizers by default.\n if opt_opt.type == 'adam':\n opt = Adam(params,\n lr=opt_opt.lr,\n betas=(opt_opt.adam_beta1, opt_opt.adam_beta2))\n else:\n raise NotImplementedError(\n 'Optimizer {} is not yet implemented.'.format(opt_opt.type))\n return opt\n\n\n"
}
]