update
This commit is contained in:
chenxuanhong
@@ -0,0 +1,110 @@
|
||||
check_points/
|
||||
pretrain_models*
|
||||
test_dir_enhance_results/
|
||||
test_dir_align_results/
|
||||
test_unalign_results/
|
||||
tmp*
|
||||
local*
|
||||
*.pth
|
||||
|
||||
# Byte-compiled / optimized / DLL files
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*$py.class
|
||||
|
||||
# C extensions
|
||||
*.so
|
||||
|
||||
# Distribution / packaging
|
||||
.Python
|
||||
env/
|
||||
build/
|
||||
develop-eggs/
|
||||
dist/
|
||||
downloads/
|
||||
eggs/
|
||||
.eggs/
|
||||
lib/
|
||||
lib64/
|
||||
parts/
|
||||
sdist/
|
||||
var/
|
||||
wheels/
|
||||
*.egg-info/
|
||||
.installed.cfg
|
||||
*.egg
|
||||
|
||||
# PyInstaller
|
||||
# Usually these files are written by a python script from a template
|
||||
# before PyInstaller builds the exe, so as to inject date/other infos into it.
|
||||
*.manifest
|
||||
*.spec
|
||||
|
||||
# Installer logs
|
||||
pip-log.txt
|
||||
pip-delete-this-directory.txt
|
||||
|
||||
# Unit test / coverage reports
|
||||
htmlcov/
|
||||
.tox/
|
||||
.coverage
|
||||
.coverage.*
|
||||
.cache
|
||||
nosetests.xml
|
||||
coverage.xml
|
||||
*.cover
|
||||
.hypothesis/
|
||||
|
||||
# Translations
|
||||
*.mo
|
||||
*.pot
|
||||
|
||||
# Django stuff:
|
||||
*.log
|
||||
local_settings.py
|
||||
|
||||
# Flask stuff:
|
||||
instance/
|
||||
.webassets-cache
|
||||
|
||||
# Scrapy stuff:
|
||||
.scrapy
|
||||
|
||||
# Sphinx documentation
|
||||
docs/_build/
|
||||
|
||||
# PyBuilder
|
||||
target/
|
||||
|
||||
# Jupyter Notebook
|
||||
.ipynb_checkpoints
|
||||
|
||||
# pyenv
|
||||
.python-version
|
||||
|
||||
# celery beat schedule file
|
||||
celerybeat-schedule
|
||||
|
||||
# SageMath parsed files
|
||||
*.sage.py
|
||||
|
||||
# dotenv
|
||||
.env
|
||||
|
||||
# virtualenv
|
||||
.venv
|
||||
venv/
|
||||
ENV/
|
||||
|
||||
# Spyder project settings
|
||||
.spyderproject
|
||||
.spyproject
|
||||
|
||||
# Rope project settings
|
||||
.ropeproject
|
||||
|
||||
# mkdocs documentation
|
||||
/site
|
||||
|
||||
# mypy
|
||||
.mypy_cache/
|
||||
@@ -0,0 +1,445 @@
|
||||
PSFR-GAN (c) by Chaofeng Chen
|
||||
|
||||
PSFR-GAN is licensed under a
|
||||
Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International License.
|
||||
|
||||
You should have received a copy of the license along with this
|
||||
work. If not, see <http://creativecommons.org/licenses/by-nc-sa/4.0/>.
|
||||
|
||||
Attribution-NonCommercial-ShareAlike 4.0 International
|
||||
|
||||
=======================================================================
|
||||
|
||||
Creative 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.
|
||||
|
||||
Using Creative Commons Public Licenses
|
||||
|
||||
Creative 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.
|
||||
|
||||
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:
|
||||
wiki.creativecommons.org/Considerations_for_licensors
|
||||
|
||||
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:
|
||||
wiki.creativecommons.org/Considerations_for_licensees
|
||||
|
||||
=======================================================================
|
||||
|
||||
Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International
|
||||
Public License
|
||||
|
||||
By exercising the Licensed Rights (defined below), You accept and agree
|
||||
to be bound by the terms and conditions of this Creative Commons
|
||||
Attribution-NonCommercial-ShareAlike 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.
|
||||
|
||||
|
||||
Section 1 -- Definitions.
|
||||
|
||||
a. 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.
|
||||
|
||||
b. 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.
|
||||
|
||||
c. BY-NC-SA Compatible License means a license listed at
|
||||
creativecommons.org/compatiblelicenses, approved by Creative
|
||||
Commons as essentially the equivalent of this Public License.
|
||||
|
||||
d. 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.
|
||||
|
||||
e. 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.
|
||||
|
||||
f. 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.
|
||||
|
||||
g. License Elements means the license attributes listed in the name
|
||||
of a Creative Commons Public License. The License Elements of this
|
||||
Public License are Attribution, NonCommercial, and ShareAlike.
|
||||
|
||||
h. Licensed Material means the artistic or literary work, database,
|
||||
or other material to which the Licensor applied this Public
|
||||
License.
|
||||
|
||||
i. 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.
|
||||
|
||||
j. Licensor means the individual(s) or entity(ies) granting rights
|
||||
under this Public License.
|
||||
|
||||
k. 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.
|
||||
|
||||
l. 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.
|
||||
|
||||
m. 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. You means the individual or entity exercising the Licensed Rights
|
||||
under this Public License. Your has a corresponding meaning.
|
||||
|
||||
|
||||
Section 2 -- Scope.
|
||||
|
||||
a. License grant.
|
||||
|
||||
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:
|
||||
|
||||
a. reproduce and Share the Licensed Material, in whole or
|
||||
in part, for NonCommercial purposes only; and
|
||||
|
||||
b. produce, reproduce, and Share Adapted Material for
|
||||
NonCommercial purposes only.
|
||||
|
||||
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.
|
||||
|
||||
3. Term. The term of this Public License is specified in Section
|
||||
6(a).
|
||||
|
||||
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.
|
||||
|
||||
5. Downstream recipients.
|
||||
|
||||
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.
|
||||
|
||||
b. Additional offer from the Licensor -- Adapted Material.
|
||||
Every recipient of Adapted Material from You
|
||||
automatically receives an offer from the Licensor to
|
||||
exercise the Licensed Rights in the Adapted Material
|
||||
under the conditions of the Adapter's License You apply.
|
||||
|
||||
c. 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.
|
||||
|
||||
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).
|
||||
|
||||
b. Other rights.
|
||||
|
||||
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.
|
||||
|
||||
2. Patent and trademark rights are not licensed under this
|
||||
Public License.
|
||||
|
||||
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.
|
||||
|
||||
|
||||
Section 3 -- License Conditions.
|
||||
|
||||
Your exercise of the Licensed Rights is expressly made subject to the
|
||||
following conditions.
|
||||
|
||||
a. Attribution.
|
||||
|
||||
1. If You Share the Licensed Material (including in modified
|
||||
form), You must:
|
||||
|
||||
a. retain the following if it is supplied by the Licensor
|
||||
with the Licensed Material:
|
||||
|
||||
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);
|
||||
|
||||
ii. a copyright notice;
|
||||
|
||||
iii. a notice that refers to this Public License;
|
||||
|
||||
iv. a notice that refers to the disclaimer of
|
||||
warranties;
|
||||
|
||||
v. a URI or hyperlink to the Licensed Material to the
|
||||
extent reasonably practicable;
|
||||
|
||||
b. indicate if You modified the Licensed Material and
|
||||
retain an indication of any previous modifications; and
|
||||
|
||||
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.
|
||||
|
||||
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.
|
||||
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.
|
||||
|
||||
b. ShareAlike.
|
||||
|
||||
In addition to the conditions in Section 3(a), if You Share
|
||||
Adapted Material You produce, the following conditions also apply.
|
||||
|
||||
1. The Adapter's License You apply must be a Creative Commons
|
||||
license with the same License Elements, this version or
|
||||
later, or a BY-NC-SA Compatible License.
|
||||
|
||||
2. You must include the text of, or the URI or hyperlink to, the
|
||||
Adapter's License You apply. You may satisfy this condition
|
||||
in any reasonable manner based on the medium, means, and
|
||||
context in which You Share Adapted Material.
|
||||
|
||||
3. You may not offer or impose any additional or different terms
|
||||
or conditions on, or apply any Effective Technological
|
||||
Measures to, Adapted Material that restrict exercise of the
|
||||
rights granted under the Adapter's License You apply.
|
||||
|
||||
|
||||
Section 4 -- Sui Generis Database Rights.
|
||||
|
||||
Where the Licensed Rights include Sui Generis Database Rights that
|
||||
apply to Your use of the Licensed Material:
|
||||
|
||||
a. 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;
|
||||
|
||||
b. 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,
|
||||
including for purposes of Section 3(b); and
|
||||
|
||||
c. You must comply with the conditions in Section 3(a) if You Share
|
||||
all or a substantial portion of the contents of the database.
|
||||
|
||||
For 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.
|
||||
|
||||
|
||||
Section 5 -- Disclaimer of Warranties and Limitation of Liability.
|
||||
|
||||
a. 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.
|
||||
|
||||
b. 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.
|
||||
|
||||
c. 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.
|
||||
|
||||
|
||||
Section 6 -- Term and Termination.
|
||||
|
||||
a. 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.
|
||||
|
||||
b. Where Your right to use the Licensed Material has terminated under
|
||||
Section 6(a), it reinstates:
|
||||
|
||||
1. automatically as of the date the violation is cured, provided
|
||||
it is cured within 30 days of Your discovery of the
|
||||
violation; or
|
||||
|
||||
2. upon express reinstatement by the Licensor.
|
||||
|
||||
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.
|
||||
|
||||
c. 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.
|
||||
|
||||
d. Sections 1, 5, 6, 7, and 8 survive termination of this Public
|
||||
License.
|
||||
|
||||
|
||||
Section 7 -- Other Terms and Conditions.
|
||||
|
||||
a. The Licensor shall not be bound by any additional or different
|
||||
terms or conditions communicated by You unless expressly agreed.
|
||||
|
||||
b. 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.
|
||||
|
||||
|
||||
Section 8 -- Interpretation.
|
||||
|
||||
a. 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.
|
||||
|
||||
b. 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.
|
||||
|
||||
c. 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.
|
||||
|
||||
d. 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.
|
||||
|
||||
=======================================================================
|
||||
|
||||
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.” The text of the Creative Commons
|
||||
public licenses is dedicated to the public domain under the CC0 Public
|
||||
Domain Dedication. 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, 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.
|
||||
|
||||
Creative Commons may be contacted at creativecommons.org.
|
||||
@@ -0,0 +1,123 @@
|
||||
# PSFR-GAN in PyTorch
|
||||
|
||||
[Progressive Semantic-Aware Style Transformation for Blind Face Restoration](https://arxiv.org/abs/2009.08709)
|
||||
[Chaofeng Chen](https://chaofengc.github.io), [Xiaoming Li](https://csxmli2016.github.io/), [Lingbo Yang](https://lotayou.github.io), [Xianhui Lin](https://dblp.org/pid/147/7708.html), [Lei Zhang](https://www4.comp.polyu.edu.hk/~cslzhang/), [Kwan-Yee K. Wong](https://i.cs.hku.hk/~kykwong/)
|
||||
|
||||
|
||||
|
||||
|
||||
### Changelog
|
||||
- **2021.04.26**: Add pytorch vgg19 model to GoogleDrive and remove `--distributed` option which causes training error.
|
||||
- **2021.03.22**: Update new model at 15 epoch (52.5k iterations).
|
||||
- **2021.03.19**: Add train codes for PSFRGAN and FPN.
|
||||
|
||||
## Prerequisites and Installation
|
||||
- Ubuntu 18.04
|
||||
- CUDA 10.1
|
||||
- Clone this repository
|
||||
```
|
||||
git clone https://github.com/chaofengc/PSFR-GAN.git
|
||||
cd PSFR-GAN
|
||||
```
|
||||
- Python 3.7, install required packages by `pip3 install -r requirements.txt`
|
||||
|
||||
## Quick Test
|
||||
|
||||
### Download Pretrain Models and Dataset
|
||||
Download the pretrained models from the following link and put them to `./pretrain_models`
|
||||
- [GoogleDrive](https://drive.google.com/drive/folders/1Ubejhxd2xd4fxGc_M_LWl3Ux6CgQd9rP?usp=sharing)
|
||||
- [BaiduNetDisk](https://pan.baidu.com/s/1cru3uUASEfGX6p6L0_7gWQ), extract code: `gj2r`
|
||||
|
||||
### Test single image
|
||||
Run the following script to enhance face(s) in single input
|
||||
```
|
||||
python test_enhance_single_unalign.py --test_img_path ./test_dir/test_hzgg.jpg --results_dir test_hzgg_results --gpus 1
|
||||
```
|
||||
|
||||
This script do the following things:
|
||||
- Crop and align all the faces from input image, stored at `results_dir/LQ_faces`
|
||||
- Parse these faces and then enhance them, results stored at `results_dir/ParseMaps` and `results_dir/HQ`
|
||||
- Paste then enhanced faces back to the original image `results_dir/hq_final.jpg`
|
||||
- You can use `--gpus` to specify how many GPUs to use, `<=0` means running on CPU. The program will use GPU with the most available memory. Set `CUDA_VISIBLE_DEVICE` to specify the GPU if you do not want automatic GPU selection.
|
||||
|
||||
### Test image folder
|
||||
To test multiple images, we first crop out all the faces and align them use the following script.
|
||||
```
|
||||
python align_and_crop_dir.py --src_dir test_dir --results_dir test_dir_align_results
|
||||
```
|
||||
|
||||
For images (*e.g.* `multiface_test.jpg`) contain multiple faces, the aligned faces will be stored as `multiface_test_{face_index}.jpg`
|
||||
And then parse the aligned faces and enhance them with
|
||||
```
|
||||
python test_enhance_dir_align.py --src_dir test_dir_align_results --results_dir test_dir_enhance_results
|
||||
```
|
||||
Results will be saved to three folders respectively: `results_dir/lq`, `results_dir/parse`, `results_dir/hq`.
|
||||
|
||||
### Additional test script
|
||||
|
||||
For your convenience, we also provide script to test multiple unaligned images and paste the enhance results back. **Note the paste back operation could be quite slow for large size images containing many faces (dlib takes time to detect faces in large image).**
|
||||
```
|
||||
python test_enhance_dir_unalign.py --src_dir test_dir --results_dir test_unalign_results
|
||||
```
|
||||
This script basically do the same thing as `test_enhance_single_unalign.py` for each image in `src_dir`
|
||||
|
||||
## Train the Model
|
||||
|
||||
### Data Preparation
|
||||
|
||||
- Download [FFHQ](https://github.com/NVlabs/ffhq-dataset) and put the images to `../datasets/FFHQ/imgs1024`
|
||||
- Download parsing masks (`512x512`) [HERE](https://drive.google.com/file/d/1eQwO8hKcaluyCnxuZAp0eJVOdgMi30uA/view?usp=sharing) generated by the pretrained FPN and put them to `../datasets/FFHQ/masks512`.
|
||||
|
||||
*Note: you may change `../datasets/FFHQ` to your own path. But images and masks must be stored under `your_own_path/imgs1024` and `your_own_path/masks512` respectively.*
|
||||
|
||||
### Train Script for PSFRGAN
|
||||
|
||||
Here is an example train script for PSFRGAN:
|
||||
|
||||
```
|
||||
python train.py --gpus 2 --model enhance --name PSFRGAN_v001 \
|
||||
--g_lr 0.0001 --d_lr 0.0004 --beta1 0.5 \
|
||||
--gan_mode 'hinge' --lambda_pix 10 --lambda_fm 10 --lambda_ss 1000 \
|
||||
--Dinput_nc 22 --D_num 3 --n_layers_D 4 \
|
||||
--batch_size 2 --dataset ffhq --dataroot ../datasets/FFHQ \
|
||||
--visual_freq 100 --print_freq 10 #--continue_train
|
||||
```
|
||||
- Please change the `--name` option for different experiments. Tensorboard records with the same name will be moved to `check_points/log_archive`, and the weight directory will only store weight history of latest experiment with the same name.
|
||||
- `--gpus` specify number of GPUs used to train. The script will use GPUs with more available memory first. To specify the GPU index, use `export CUDA_VISIBLE_DEVICES=your_gpu_ids` before the script.
|
||||
- Uncomment `--continue_train` to resume training. *Current codes do not resume the optimizer state.*
|
||||
- It needs at least **8GB** memory to train with **batch_size=1**.
|
||||
|
||||
### Scripts for FPN
|
||||
|
||||
You may also train your own FPN and generate masks for the HQ images by yourself with the following steps:
|
||||
|
||||
- Download [CelebAHQ-Mask](https://github.com/switchablenorms/CelebAMask-HQ) dataset. Generate `CelebAMask-HQ-mask` and `CelebAMask-HQ-mask-color` with the provided scripts in `CelebAMask-HQ/face_parsing/Data_preprocessing/`.
|
||||
- Train FPN with the following commmand
|
||||
```
|
||||
python train.py --gpus 1 --model parse --name FPN_v001 \
|
||||
--lr 0.0002 --batch_size 8 \
|
||||
--dataset celebahqmask --dataroot ../datasets/CelebAMask-HQ \
|
||||
--visual_freq 100 --print_freq 10 #--continue_train
|
||||
```
|
||||
- Generate parsing masks with your own FPN using the following command:
|
||||
```
|
||||
python generate_masks.py --save_masks_dir ../datasets/FFHQ/masks512 --batch_size 8 --parse_net_weight path/to/your/own/FPN
|
||||
```
|
||||
|
||||
## Citation
|
||||
```
|
||||
@inproceedings{ChenPSFRGAN,
|
||||
author = {Chen, Chaofeng and Li, Xiaoming and Lingbo, Yang and Lin, Xianhui and Zhang, Lei and Wong, KKY},
|
||||
title = {Progressive Semantic-Aware Style Transformation for Blind Face Restoration},
|
||||
Journal = {IEEE Conference on Computer Vision and Pattern Recognition (CVPR)},
|
||||
year = {2021}
|
||||
}
|
||||
```
|
||||
|
||||
## License
|
||||
|
||||
<a rel="license" href="http://creativecommons.org/licenses/by-nc-sa/4.0/"><img alt="Creative Commons License" style="border-width:0" src="https://i.creativecommons.org/l/by-nc-sa/4.0/88x31.png" /></a><br />This work is licensed under a <a rel="license" href="http://creativecommons.org/licenses/by-nc-sa/4.0/">Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International License</a>.
|
||||
|
||||
## Acknowledgement
|
||||
|
||||
This work is inspired by [SPADE](https://github.com/NVlabs/SPADE), and closed related to [DFDNet](https://github.com/csxmli2016/DFDNet) and [HiFaceGAN](https://github.com/Lotayou/Face-Renovation). Our codes largely benefit from [CycleGAN](https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix).
|
||||
@@ -0,0 +1,86 @@
|
||||
import dlib
|
||||
import os
|
||||
import cv2
|
||||
import numpy as np
|
||||
from tqdm import tqdm
|
||||
from skimage import transform as trans
|
||||
from skimage import io
|
||||
import argparse
|
||||
|
||||
|
||||
def get_points(img, detector, shape_predictor, size_threshold=999):
|
||||
dets = detector(img, 1)
|
||||
if len(dets) == 0:
|
||||
return None
|
||||
|
||||
all_points = []
|
||||
for det in dets:
|
||||
if isinstance(detector, dlib.cnn_face_detection_model_v1):
|
||||
rec = det.rect # for cnn detector
|
||||
else:
|
||||
rec = det
|
||||
if rec.width() > size_threshold or rec.height() > size_threshold:
|
||||
break
|
||||
shape = shape_predictor(img, rec)
|
||||
single_points = []
|
||||
for i in range(5):
|
||||
single_points.append([shape.part(i).x, shape.part(i).y])
|
||||
all_points.append(np.array(single_points))
|
||||
if len(all_points) <= 0:
|
||||
return None
|
||||
else:
|
||||
return all_points
|
||||
|
||||
def align_and_save(img, save_path, src_points, template_path, template_scale=1):
|
||||
out_size = (512, 512)
|
||||
reference = np.load(template_path) / template_scale
|
||||
|
||||
ext = os.path.splitext(save_path)
|
||||
for idx, spoint in enumerate(src_points):
|
||||
tform = trans.SimilarityTransform()
|
||||
tform.estimate(spoint, reference)
|
||||
M = tform.params[0:2,:]
|
||||
|
||||
crop_img = cv2.warpAffine(img, M, out_size)
|
||||
if len(src_points) > 1:
|
||||
save_path = ext[0] + '_{}'.format(idx) + ext[1]
|
||||
dlib.save_image(crop_img.astype(np.uint8), save_path)
|
||||
print('Saving image', save_path)
|
||||
|
||||
def align_and_save_dir(src_dir, save_dir, template_path='./pretrain_models/FFHQ_template.npy', template_scale=2, use_cnn_detector=True):
|
||||
out_size = (512, 512)
|
||||
if use_cnn_detector:
|
||||
detector = dlib.cnn_face_detection_model_v1('./pretrain_models/mmod_human_face_detector.dat')
|
||||
else:
|
||||
detector = dlib.get_frontal_face_detector()
|
||||
sp = dlib.shape_predictor('./pretrain_models/shape_predictor_5_face_landmarks.dat')
|
||||
|
||||
for name in os.listdir(src_dir):
|
||||
img_path = os.path.join(src_dir, name)
|
||||
img = dlib.load_rgb_image(img_path)
|
||||
|
||||
points = get_points(img, detector, sp)
|
||||
if points is not None:
|
||||
save_path = os.path.join(save_dir, name)
|
||||
align_and_save(img, save_path, points, template_path, template_scale)
|
||||
else:
|
||||
print('No face detected in', img_path)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--src_dir', type=str, help='source directory containing images to crop and align.')
|
||||
parser.add_argument('--results_dir', type=str, help='results directory to save the aligned faces.')
|
||||
parser.add_argument('--not_use_cnn_detector', action='store_true', help='do not use cnn face detector in dlib.')
|
||||
args = parser.parse_args()
|
||||
|
||||
src_dir = args.src_dir
|
||||
assert os.path.isdir(src_dir), 'Source path should be a directory containing images'
|
||||
save_dir = args.results_dir
|
||||
if not os.path.exists(save_dir): os.makedirs(save_dir, exist_ok=True)
|
||||
align_and_save_dir(src_dir, save_dir, use_cnn_detector=not args.not_use_cnn_detector)
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,94 @@
|
||||
"""This package includes all the modules related to data loading and preprocessing
|
||||
|
||||
To add a custom dataset class called 'dummy', you need to add a file called 'dummy_dataset.py' and define a subclass 'DummyDataset' inherited from BaseDataset.
|
||||
You need to implement four functions:
|
||||
-- <__init__>: initialize the class, first call BaseDataset.__init__(self, opt).
|
||||
-- <__len__>: return the size of dataset.
|
||||
-- <__getitem__>: get a data point from data loader.
|
||||
-- <modify_commandline_options>: (optionally) add dataset-specific options and set default options.
|
||||
|
||||
Now you can use the dataset class by specifying flag '--dataset_mode dummy'.
|
||||
See our template dataset class 'template_dataset.py' for more details.
|
||||
"""
|
||||
import importlib
|
||||
import torch.utils.data
|
||||
from data.base_dataset import BaseDataset
|
||||
|
||||
|
||||
def find_dataset_using_name(dataset_name):
|
||||
"""Import the module "data/[dataset_name]_dataset.py".
|
||||
|
||||
In the file, the class called DatasetNameDataset() will
|
||||
be instantiated. It has to be a subclass of BaseDataset,
|
||||
and it is case-insensitive.
|
||||
"""
|
||||
dataset_filename = "data." + dataset_name + "_dataset"
|
||||
datasetlib = importlib.import_module(dataset_filename)
|
||||
|
||||
dataset = None
|
||||
target_dataset_name = dataset_name.replace('_', '') + 'dataset'
|
||||
for name, cls in datasetlib.__dict__.items():
|
||||
if name.lower() == target_dataset_name.lower() \
|
||||
and issubclass(cls, BaseDataset):
|
||||
dataset = cls
|
||||
|
||||
if dataset is None:
|
||||
raise NotImplementedError("In %s.py, there should be a subclass of BaseDataset with class name that matches %s in lowercase." % (dataset_filename, target_dataset_name))
|
||||
|
||||
return dataset
|
||||
|
||||
|
||||
def get_option_setter(dataset_name):
|
||||
"""Return the static method <modify_commandline_options> of the dataset class."""
|
||||
dataset_class = find_dataset_using_name(dataset_name)
|
||||
return dataset_class.modify_commandline_options
|
||||
|
||||
|
||||
def create_dataset(opt):
|
||||
"""Create a dataset given the option.
|
||||
|
||||
This function wraps the class CustomDatasetDataLoader.
|
||||
This is the main interface between this package and 'train.py'/'test.py'
|
||||
|
||||
Example:
|
||||
>>> from data import create_dataset
|
||||
>>> dataset = create_dataset(opt)
|
||||
"""
|
||||
data_loader = CustomDatasetDataLoader(opt)
|
||||
dataset = data_loader.load_data()
|
||||
return dataset
|
||||
|
||||
|
||||
class CustomDatasetDataLoader():
|
||||
"""Wrapper class of Dataset class that performs multi-threaded data loading"""
|
||||
|
||||
def __init__(self, opt):
|
||||
"""Initialize this class
|
||||
|
||||
Step 1: create a dataset instance given the name [dataset_mode]
|
||||
Step 2: create a multi-threaded data loader.
|
||||
"""
|
||||
self.opt = opt
|
||||
dataset_class = find_dataset_using_name(opt.dataset_name)
|
||||
self.dataset = dataset_class(opt)
|
||||
print("dataset [%s] was created" % type(self.dataset).__name__)
|
||||
drop_last = True if opt.isTrain else False
|
||||
self.dataloader = torch.utils.data.DataLoader(
|
||||
self.dataset,
|
||||
batch_size=opt.batch_size,
|
||||
shuffle=not opt.serial_batches,
|
||||
num_workers=int(opt.num_threads), drop_last=drop_last)
|
||||
|
||||
def load_data(self):
|
||||
return self
|
||||
|
||||
def __len__(self):
|
||||
"""Return the number of data in the dataset"""
|
||||
return min(len(self.dataset), self.opt.max_dataset_size)
|
||||
|
||||
def __iter__(self):
|
||||
"""Return a batch of data"""
|
||||
for i, data in enumerate(self.dataloader):
|
||||
if i * self.opt.batch_size >= self.opt.max_dataset_size:
|
||||
break
|
||||
yield data
|
||||
@@ -0,0 +1,162 @@
|
||||
"""This module implements an abstract base class (ABC) 'BaseDataset' for datasets.
|
||||
|
||||
It also includes common transformation functions (e.g., get_transform, __scale_width), which can be later used in subclasses.
|
||||
"""
|
||||
import random
|
||||
import numpy as np
|
||||
import torch.utils.data as data
|
||||
from PIL import Image
|
||||
import torchvision.transforms as transforms
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
import imgaug as ia
|
||||
import imgaug.augmenters as iaa
|
||||
|
||||
class BaseDataset(data.Dataset, ABC):
|
||||
"""This class is an abstract base class (ABC) for datasets.
|
||||
|
||||
To create a subclass, you need to implement the following four functions:
|
||||
-- <__init__>: initialize the class, first call BaseDataset.__init__(self, opt).
|
||||
-- <__len__>: return the size of dataset.
|
||||
-- <__getitem__>: get a data point.
|
||||
-- <modify_commandline_options>: (optionally) add dataset-specific options and set default options.
|
||||
"""
|
||||
|
||||
def __init__(self, opt):
|
||||
"""Initialize the class; save the options in the class
|
||||
|
||||
Parameters:
|
||||
opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
|
||||
"""
|
||||
self.opt = opt
|
||||
self.root = opt.dataroot
|
||||
|
||||
@staticmethod
|
||||
def modify_commandline_options(parser, is_train):
|
||||
"""Add new dataset-specific options, and rewrite default values for existing options.
|
||||
|
||||
Parameters:
|
||||
parser -- original option parser
|
||||
is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
|
||||
|
||||
Returns:
|
||||
the modified parser.
|
||||
"""
|
||||
return parser
|
||||
|
||||
@abstractmethod
|
||||
def __len__(self):
|
||||
"""Return the total number of images in the dataset."""
|
||||
return 0
|
||||
|
||||
@abstractmethod
|
||||
def __getitem__(self, index):
|
||||
"""Return a data point and its metadata information.
|
||||
|
||||
Parameters:
|
||||
index - - a random integer for data indexing
|
||||
|
||||
Returns:
|
||||
a dictionary of data with their names. It ususally contains the data itself and its metadata information.
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
def get_params(opt, size):
|
||||
w, h = size
|
||||
new_h = h
|
||||
new_w = w
|
||||
if opt.preprocess == 'resize_and_crop':
|
||||
new_h = new_w = opt.load_size
|
||||
elif opt.preprocess == 'scale_width_and_crop':
|
||||
new_w = opt.load_size
|
||||
new_h = opt.load_size * h // w
|
||||
|
||||
x = random.randint(0, np.maximum(0, new_w - opt.crop_size))
|
||||
y = random.randint(0, np.maximum(0, new_h - opt.crop_size))
|
||||
|
||||
flip = random.random() > 0.5
|
||||
|
||||
return {'crop_pos': (x, y), 'flip': flip}
|
||||
|
||||
def get_transform(opt, params=None, grayscale=False, method=Image.BICUBIC, convert=True):
|
||||
transform_list = []
|
||||
if grayscale:
|
||||
# transform_list.append(transforms.Grayscale(1))
|
||||
from util import util
|
||||
transform_list.append(util.RGBtoY)
|
||||
if 'resize' in opt.preprocess:
|
||||
osize = [opt.load_size, opt.load_size]
|
||||
transform_list.append(transforms.Resize(osize, method))
|
||||
elif 'scale_width' in opt.preprocess:
|
||||
transform_list.append(transforms.Lambda(lambda img: __scale_width(img, opt.load_size, method)))
|
||||
|
||||
if 'crop' in opt.preprocess:
|
||||
if params is None:
|
||||
transform_list.append(transforms.RandomCrop(opt.crop_size))
|
||||
else:
|
||||
if 'crop_size' in params:
|
||||
transform_list.append(transforms.Lambda(lambda img: __crop(img, params['crop_pos'], params['crop_size'])))
|
||||
else:
|
||||
transform_list.append(transforms.Lambda(lambda img: __crop(img, params['crop_pos'], opt.crop_size)))
|
||||
|
||||
if opt.preprocess == 'none':
|
||||
transform_list.append(transforms.Lambda(lambda img: __make_power_2(img, base=4, method=method)))
|
||||
|
||||
if not opt.no_flip:
|
||||
if params is None:
|
||||
transform_list.append(transforms.RandomHorizontalFlip())
|
||||
elif params['flip']:
|
||||
transform_list.append(transforms.Lambda(lambda img: __flip(img, params['flip'])))
|
||||
|
||||
if convert:
|
||||
transform_list += [transforms.ToTensor()]
|
||||
if grayscale:
|
||||
transform_list += [transforms.Normalize((0.5,), (0.5,))]
|
||||
else:
|
||||
transform_list += [transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
|
||||
return transforms.Compose(transform_list)
|
||||
|
||||
def __make_power_2(img, base, method=Image.BICUBIC):
|
||||
ow, oh = img.size
|
||||
h = int(round(oh / base) * base)
|
||||
w = int(round(ow / base) * base)
|
||||
if (h == oh) and (w == ow):
|
||||
return img
|
||||
|
||||
__print_size_warning(ow, oh, w, h)
|
||||
return img.resize((w, h), method)
|
||||
|
||||
|
||||
def __scale_width(img, target_width, method=Image.BICUBIC):
|
||||
ow, oh = img.size
|
||||
if (ow == target_width):
|
||||
return img
|
||||
w = target_width
|
||||
h = int(target_width * oh / ow)
|
||||
return img.resize((w, h), method)
|
||||
|
||||
|
||||
def __crop(img, pos, size):
|
||||
ow, oh = img.size
|
||||
x1, y1 = pos
|
||||
tw = th = size
|
||||
if (ow > tw or oh > th):
|
||||
return img.crop((x1, y1, x1 + tw, y1 + th))
|
||||
return img
|
||||
|
||||
|
||||
def __flip(img, flip):
|
||||
if flip:
|
||||
return img.transpose(Image.FLIP_LEFT_RIGHT)
|
||||
return img
|
||||
|
||||
|
||||
def __print_size_warning(ow, oh, w, h):
|
||||
"""Print warning information about image size(only print once)"""
|
||||
if not hasattr(__print_size_warning, 'has_printed'):
|
||||
print("The image size needs to be a multiple of 4. "
|
||||
"The loaded image size was (%d, %d), so it was adjusted to "
|
||||
"(%d, %d). This adjustment will be done to all images "
|
||||
"whose sizes are not multiples of 4" % (ow, oh, w, h))
|
||||
__print_size_warning.has_printed = True
|
||||
@@ -0,0 +1,60 @@
|
||||
import os
|
||||
import random
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
import imgaug as ia
|
||||
import imgaug.augmenters as iaa
|
||||
|
||||
from data.image_folder import make_dataset
|
||||
|
||||
import torch
|
||||
from torch.utils.data import Dataset
|
||||
from torchvision.transforms import transforms
|
||||
|
||||
from data.base_dataset import BaseDataset
|
||||
from utils.utils import onehot_parse_map
|
||||
|
||||
from data.ffhq_dataset import complex_imgaug, random_gray
|
||||
|
||||
class CelebAHQMaskDataset(BaseDataset):
|
||||
|
||||
def __init__(self, opt):
|
||||
BaseDataset.__init__(self, opt)
|
||||
self.img_size = opt.Pimg_size
|
||||
self.lr_size = opt.Gin_size
|
||||
self.hr_size = opt.Gout_size
|
||||
self.shuffle = True if opt.isTrain else False
|
||||
|
||||
self.img_dataset = sorted(make_dataset(os.path.join(opt.dataroot, 'CelebA-HQ-img')))
|
||||
self.mask_dataset = sorted(make_dataset(os.path.join(opt.dataroot, 'CelebAMask-HQ-mask')))
|
||||
|
||||
self.to_tensor = transforms.Compose([
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
|
||||
])
|
||||
|
||||
def __len__(self,):
|
||||
return len(self.img_dataset)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
sample = {}
|
||||
img_path = self.img_dataset[idx]
|
||||
mask_path = self.mask_dataset[idx]
|
||||
hr_img = Image.open(img_path).convert('RGB')
|
||||
mask_img = Image.open(mask_path)
|
||||
|
||||
hr_img = hr_img.resize((self.hr_size, self.hr_size))
|
||||
hr_img = random_gray(hr_img, p=0.3)
|
||||
scale_size = np.random.randint(32, 256)
|
||||
lr_img = complex_imgaug(hr_img, self.img_size, scale_size)
|
||||
|
||||
mask_img = mask_img.resize((self.hr_size, self.hr_size))
|
||||
mask_label = torch.tensor(np.array(mask_img)).long()
|
||||
|
||||
hr_tensor = self.to_tensor(hr_img)
|
||||
lr_tensor = self.to_tensor(lr_img)
|
||||
|
||||
return {'HR': hr_tensor, 'LR': lr_tensor, 'HR_paths': img_path, 'Mask': mask_label}
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,88 @@
|
||||
import os
|
||||
import random
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
import imgaug as ia
|
||||
import imgaug.augmenters as iaa
|
||||
|
||||
from data.image_folder import make_dataset
|
||||
|
||||
import torch
|
||||
from torch.utils.data import Dataset
|
||||
from torchvision.transforms import transforms
|
||||
|
||||
from data.base_dataset import BaseDataset
|
||||
from utils.utils import onehot_parse_map
|
||||
|
||||
class FFHQDataset(BaseDataset):
|
||||
|
||||
def __init__(self, opt):
|
||||
BaseDataset.__init__(self, opt)
|
||||
self.img_size = opt.Pimg_size
|
||||
self.lr_size = opt.Gin_size
|
||||
self.hr_size = opt.Gout_size
|
||||
self.shuffle = True if opt.isTrain else False
|
||||
|
||||
self.img_dataset = sorted(make_dataset(os.path.join(opt.dataroot, 'imgs1024')))
|
||||
self.mask_dataset = sorted(make_dataset(os.path.join(opt.dataroot, 'masks512')))
|
||||
|
||||
self.to_tensor = transforms.Compose([
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
|
||||
])
|
||||
self.random_crop = transforms.RandomCrop(self.hr_size)
|
||||
|
||||
def __len__(self,):
|
||||
return len(self.img_dataset)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
sample = {}
|
||||
img_path = self.img_dataset[idx]
|
||||
mask_path = self.mask_dataset[idx]
|
||||
hr_img = Image.open(img_path).convert('RGB')
|
||||
mask_img = Image.open(mask_path).convert('RGB')
|
||||
|
||||
hr_img = hr_img.resize((self.hr_size, self.hr_size))
|
||||
hr_img = random_gray(hr_img, p=0.3)
|
||||
scale_size = np.random.randint(32, 256)
|
||||
lr_img = complex_imgaug(hr_img, self.img_size, scale_size)
|
||||
|
||||
mask_img = mask_img.resize((self.hr_size, self.hr_size))
|
||||
mask_label = onehot_parse_map(mask_img)
|
||||
mask_label = torch.tensor(mask_label).float()
|
||||
|
||||
hr_tensor = self.to_tensor(hr_img)
|
||||
lr_tensor = self.to_tensor(lr_img)
|
||||
|
||||
return {'HR': hr_tensor, 'LR': lr_tensor, 'HR_paths': img_path, 'Mask': mask_label}
|
||||
|
||||
|
||||
def complex_imgaug(x, org_size, scale_size):
|
||||
"""input single RGB PIL Image instance"""
|
||||
x = np.array(x)
|
||||
x = x[np.newaxis, :, :, :]
|
||||
aug_seq = iaa.Sequential([
|
||||
iaa.Sometimes(0.5, iaa.OneOf([
|
||||
iaa.GaussianBlur((3, 15)),
|
||||
iaa.AverageBlur(k=(3, 15)),
|
||||
iaa.MedianBlur(k=(3, 15)),
|
||||
iaa.MotionBlur((5, 25))
|
||||
])),
|
||||
iaa.Resize(scale_size, interpolation=ia.ALL),
|
||||
iaa.Sometimes(0.2, iaa.AdditiveGaussianNoise(loc=0, scale=(0.0, 0.1*255), per_channel=0.5)),
|
||||
iaa.Sometimes(0.7, iaa.JpegCompression(compression=(10, 65))),
|
||||
iaa.Resize(org_size),
|
||||
])
|
||||
|
||||
aug_img = aug_seq(images=x)
|
||||
return aug_img[0]
|
||||
|
||||
|
||||
def random_gray(x, p=0.5):
|
||||
"""input single RGB PIL Image instance"""
|
||||
x = np.array(x)
|
||||
x = x[np.newaxis, :, :, :]
|
||||
aug = iaa.Sometimes(p, iaa.Grayscale(alpha=1.0))
|
||||
aug_img = aug(images=x)
|
||||
return aug_img[0]
|
||||
|
||||
@@ -0,0 +1,67 @@
|
||||
"""A modified image folder class
|
||||
|
||||
We modify the official PyTorch image folder (https://github.com/pytorch/vision/blob/master/torchvision/datasets/folder.py)
|
||||
so that this class can load images from both current directory and its subdirectories.
|
||||
"""
|
||||
|
||||
import torch.utils.data as data
|
||||
|
||||
from PIL import Image
|
||||
import os
|
||||
import os.path
|
||||
|
||||
IMG_EXTENSIONS = [
|
||||
'.jpg', '.JPG', '.jpeg', '.JPEG',
|
||||
'.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP',
|
||||
'.tif', '.TIF', '.tiff', '.TIFF',
|
||||
]
|
||||
|
||||
|
||||
def is_image_file(filename):
|
||||
return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
|
||||
|
||||
|
||||
def make_dataset(dir, max_dataset_size=float("inf")):
|
||||
images = []
|
||||
assert os.path.isdir(dir), '%s is not a valid directory' % dir
|
||||
|
||||
for root, _, fnames in sorted(os.walk(dir)):
|
||||
for fname in fnames:
|
||||
if is_image_file(fname):
|
||||
path = os.path.join(root, fname)
|
||||
images.append(path)
|
||||
return images[:min(max_dataset_size, len(images))]
|
||||
|
||||
|
||||
def default_loader(path):
|
||||
return Image.open(path).convert('RGB')
|
||||
|
||||
|
||||
class ImageFolder(data.Dataset):
|
||||
|
||||
def __init__(self, root, transform=None, return_paths=False,
|
||||
loader=default_loader):
|
||||
imgs = make_dataset(root)
|
||||
if len(imgs) == 0:
|
||||
raise(RuntimeError("Found 0 images in: " + root + "\n"
|
||||
"Supported image extensions are: " +
|
||||
",".join(IMG_EXTENSIONS)))
|
||||
|
||||
self.root = root
|
||||
self.imgs = imgs
|
||||
self.transform = transform
|
||||
self.return_paths = return_paths
|
||||
self.loader = loader
|
||||
|
||||
def __getitem__(self, index):
|
||||
path = self.imgs[index]
|
||||
img = self.loader(path)
|
||||
if self.transform is not None:
|
||||
img = self.transform(img)
|
||||
if self.return_paths:
|
||||
return img, path
|
||||
else:
|
||||
return img
|
||||
|
||||
def __len__(self):
|
||||
return len(self.imgs)
|
||||
@@ -0,0 +1,43 @@
|
||||
from data.base_dataset import BaseDataset, get_transform
|
||||
from data.image_folder import make_dataset
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
|
||||
class SingleDataset(BaseDataset):
|
||||
"""This dataset class can load a set of images specified by the path --dataroot /path/to/data.
|
||||
|
||||
It can be used for generating CycleGAN results only for one side with the model option '-model test'.
|
||||
"""
|
||||
|
||||
def __init__(self, opt):
|
||||
"""Initialize this dataset class.
|
||||
|
||||
Parameters:
|
||||
opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
|
||||
"""
|
||||
BaseDataset.__init__(self, opt)
|
||||
self.A_paths = sorted(make_dataset(opt.src_dir, opt.max_dataset_size))
|
||||
input_nc = self.opt.output_nc
|
||||
self.transform = get_transform(opt, grayscale=(input_nc == 1))
|
||||
self.opt = opt
|
||||
|
||||
def __getitem__(self, index):
|
||||
"""Return a data point and its metadata information.
|
||||
|
||||
Parameters:
|
||||
index - - a random integer for data indexing
|
||||
|
||||
Returns a dictionary that contains A and A_paths
|
||||
A(tensor) - - an image in one domain
|
||||
A_paths(str) - - the path of the image
|
||||
"""
|
||||
A_path = self.A_paths[index]
|
||||
A_img = Image.open(A_path).convert('RGB')
|
||||
A_img = A_img.resize((512, 512), Image.BICUBIC)
|
||||
|
||||
A = self.transform(A_img)
|
||||
return {'LR': A, 'LR_paths': A_path}
|
||||
|
||||
def __len__(self):
|
||||
"""Return the total number of images in the dataset."""
|
||||
return len(self.A_paths)
|
||||
@@ -0,0 +1,92 @@
|
||||
import os
|
||||
from options.test_options import TestOptions
|
||||
from data import create_dataset
|
||||
from models import create_model
|
||||
from utils import utils
|
||||
from PIL import Image
|
||||
from tqdm import tqdm
|
||||
import torch
|
||||
import time
|
||||
import numpy as np
|
||||
import cv2
|
||||
import glob
|
||||
from torchvision.transforms import transforms
|
||||
|
||||
if __name__ == '__main__':
|
||||
opt = TestOptions()
|
||||
opt = opt.parse() # get test options
|
||||
opt.num_threads = 0 # test code only supports num_threads = 1
|
||||
opt.serial_batches = True # disable data shuffling; comment this line if results on randomly chosen images are needed.
|
||||
opt.no_flip = True
|
||||
|
||||
dataset = create_dataset(opt) # create a dataset given opt.dataset_mode and other options
|
||||
model = create_model(opt) # create a model given opt.model and other options
|
||||
model.load_pretrain_models()
|
||||
|
||||
netP = model.netP
|
||||
model.eval()
|
||||
|
||||
to_tensor = transforms.Compose([
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
|
||||
])
|
||||
|
||||
image_dir = "G:/VGGFace2-HQ/VGGface2_None_norm_512_true_bygfpgan/"
|
||||
output_dir = "G:/VGGFace2-HQ/VGGface2_HQ_original_aligned_mask"
|
||||
|
||||
temp_path = os.path.join(image_dir,'*/')
|
||||
pathes = glob.glob(temp_path)
|
||||
dataset = []
|
||||
for dir_item in pathes:
|
||||
join_path = glob.glob(os.path.join(dir_item,'*.jpg'))
|
||||
print("processing %s"%dir_item,end='\r')
|
||||
temp_list = []
|
||||
for item in join_path:
|
||||
temp_list.append(item)
|
||||
dataset.append(temp_list)
|
||||
|
||||
# ------------------------ restore ------------------------
|
||||
for i_dir in dataset:
|
||||
path = os.path.dirname(i_dir[0])
|
||||
dir_name = os.path.join(output_dir, os.path.basename(path))
|
||||
if not os.path.exists(dir_name):
|
||||
os.makedirs(dir_name)
|
||||
|
||||
for img_path in i_dir:
|
||||
hr_img = Image.open(img_path).convert('RGB')
|
||||
inp = to_tensor(hr_img).unsqueeze(0)
|
||||
with torch.no_grad():
|
||||
parse_map, _ = netP(inp)
|
||||
parse_map_sm = (parse_map == parse_map.max(dim=1, keepdim=True)[0]).float()
|
||||
ref_parse_img = utils.color_parse_map(parse_map_sm)
|
||||
img_name = os.path.basename(img_path)
|
||||
basename, ext = os.path.splitext(img_name)
|
||||
save_face_name = f'{basename}.png'
|
||||
# print(save_face_name)
|
||||
save_path = os.path.join(dir_name, save_face_name)
|
||||
# os.makedirs(opt.save_masks_dir, exist_ok=True)
|
||||
img = cv2.cvtColor(ref_parse_img[0],cv2.COLOR_RGB2GRAY)
|
||||
cv2.imwrite(save_path,img)
|
||||
|
||||
# for i, data in tqdm(enumerate(dataset), total=len(dataset)//opt.batch_size):
|
||||
# inp = data['LR']
|
||||
# with torch.no_grad():
|
||||
# parse_map, _ = netP(inp)
|
||||
# parse_map_sm = (parse_map == parse_map.max(dim=1, keepdim=True)[0]).float()
|
||||
# img_path = data['LR_paths'] # get image paths
|
||||
# ref_parse_img = utils.color_parse_map(parse_map_sm)
|
||||
# for i in range(len(img_path)):
|
||||
# img_name = os.path.basename(img_path[i])
|
||||
# basename, ext = os.path.splitext(img_name)
|
||||
# save_face_name = f'{basename}.png'
|
||||
# # print(save_face_name)
|
||||
# save_path = os.path.join(opt.save_masks_dir, save_face_name)
|
||||
# os.makedirs(opt.save_masks_dir, exist_ok=True)
|
||||
# img = cv2.cvtColor(ref_parse_img[i],cv2.COLOR_RGB2GRAY)
|
||||
# cv2.imwrite(save_path,img)
|
||||
# save_img = Image.fromarray(ref_parse_img[i])
|
||||
# save_img.save(save_path)
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,67 @@
|
||||
"""This package contains modules related to objective functions, optimizations, and network architectures.
|
||||
|
||||
To add a custom model class called 'dummy', you need to add a file called 'dummy_model.py' and define a subclass DummyModel inherited from BaseModel.
|
||||
You need to implement the following five functions:
|
||||
-- <__init__>: initialize the class; first call BaseModel.__init__(self, opt).
|
||||
-- <set_input>: unpack data from dataset and apply preprocessing.
|
||||
-- <forward>: produce intermediate results.
|
||||
-- <optimize_parameters>: calculate loss, gradients, and update network weights.
|
||||
-- <modify_commandline_options>: (optionally) add model-specific options and set default options.
|
||||
|
||||
In the function <__init__>, you need to define four lists:
|
||||
-- self.loss_names (str list): specify the training losses that you want to plot and save.
|
||||
-- self.model_names (str list): define networks used in our training.
|
||||
-- self.visual_names (str list): specify the images that you want to display and save.
|
||||
-- self.optimizers (optimizer list): define and initialize optimizers. You can define one optimizer for each network. If two networks are updated at the same time, you can use itertools.chain to group them. See cycle_gan_model.py for an usage.
|
||||
|
||||
Now you can use the model class by specifying flag '--model dummy'.
|
||||
See our template model class 'template_model.py' for more details.
|
||||
"""
|
||||
|
||||
import importlib
|
||||
from models.base_model import BaseModel
|
||||
|
||||
|
||||
def find_model_using_name(model_name):
|
||||
"""Import the module "models/[model_name]_model.py".
|
||||
|
||||
In the file, the class called DatasetNameModel() will
|
||||
be instantiated. It has to be a subclass of BaseModel,
|
||||
and it is case-insensitive.
|
||||
"""
|
||||
model_filename = "models." + model_name + "_model"
|
||||
modellib = importlib.import_module(model_filename)
|
||||
model = None
|
||||
target_model_name = model_name.replace('_', '') + 'model'
|
||||
for name, cls in modellib.__dict__.items():
|
||||
if name.lower() == target_model_name.lower() \
|
||||
and issubclass(cls, BaseModel):
|
||||
model = cls
|
||||
|
||||
if model is None:
|
||||
print("In %s.py, there should be a subclass of BaseModel with class name that matches %s in lowercase." % (model_filename, target_model_name))
|
||||
exit(0)
|
||||
|
||||
return model
|
||||
|
||||
|
||||
def get_option_setter(model_name):
|
||||
"""Return the static method <modify_commandline_options> of the model class."""
|
||||
model_class = find_model_using_name(model_name)
|
||||
return model_class.modify_commandline_options
|
||||
|
||||
|
||||
def create_model(opt):
|
||||
"""Create a model given the option.
|
||||
|
||||
This function warps the class CustomDatasetDataLoader.
|
||||
This is the main interface between this package and 'train.py'/'test.py'
|
||||
|
||||
Example:
|
||||
>>> from models import create_model
|
||||
>>> model = create_model(opt)
|
||||
"""
|
||||
model = find_model_using_name(opt.model)
|
||||
instance = model(opt)
|
||||
print("model [%s] was created" % type(instance).__name__)
|
||||
return instance
|
||||
@@ -0,0 +1,248 @@
|
||||
import os
|
||||
import torch
|
||||
from collections import OrderedDict
|
||||
from abc import ABC, abstractmethod
|
||||
from . import networks
|
||||
|
||||
class BaseModel(ABC):
|
||||
"""This class is an abstract base class (ABC) for models.
|
||||
To create a subclass, you need to implement the following five functions:
|
||||
-- <__init__>: initialize the class; first call BaseModel.__init__(self, opt).
|
||||
-- <set_input>: unpack data from dataset and apply preprocessing.
|
||||
-- <forward>: produce intermediate results.
|
||||
-- <optimize_parameters>: calculate losses, gradients, and update network weights.
|
||||
-- <modify_commandline_options>: (optionally) add model-specific options and set default options.
|
||||
"""
|
||||
|
||||
def __init__(self, opt):
|
||||
"""Initialize the BaseModel class.
|
||||
|
||||
Parameters:
|
||||
opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
|
||||
|
||||
When creating your custom class, you need to implement your own initialization.
|
||||
In this fucntion, you should first call <BaseModel.__init__(self, opt)>
|
||||
Then, you need to define four lists:
|
||||
-- self.loss_names (str list): specify the training losses that you want to plot and save.
|
||||
-- self.model_names (str list): specify the images that you want to display and save.
|
||||
-- self.visual_names (str list): define networks used in our training.
|
||||
-- self.optimizers (optimizer list): define and initialize optimizers. You can define one optimizer for each network. If two networks are updated at the same time, you can use itertools.chain to group them. See cycle_gan_model.py for an example.
|
||||
"""
|
||||
self.opt = opt
|
||||
self.gpu_ids = opt.gpu_ids
|
||||
self.isTrain = opt.isTrain
|
||||
self.save_dir = os.path.join(opt.checkpoints_dir, opt.name) # save all the checkpoints to save_dir
|
||||
self.device = torch.device('cuda:{}'.format(self.gpu_ids[0])) if self.gpu_ids else torch.device('cpu') # get device name: CPU or GPU
|
||||
|
||||
self.loss_names = []
|
||||
self.model_names = []
|
||||
self.visual_names = []
|
||||
self.optimizers = []
|
||||
self.image_paths = []
|
||||
self.metric = 0 # used for learning rate policy 'plateau'
|
||||
|
||||
@staticmethod
|
||||
def modify_commandline_options(parser, is_train):
|
||||
"""Add new model-specific options, and rewrite default values for existing options.
|
||||
|
||||
Parameters:
|
||||
parser -- original option parser
|
||||
is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
|
||||
|
||||
Returns:
|
||||
the modified parser.
|
||||
"""
|
||||
return parser
|
||||
|
||||
@abstractmethod
|
||||
def set_input(self, input):
|
||||
"""Unpack input data from the dataloader and perform necessary pre-processing steps.
|
||||
|
||||
Parameters:
|
||||
input (dict): includes the data itself and its metadata information.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def forward(self):
|
||||
"""Run forward pass; called by both functions <optimize_parameters> and <test>."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def optimize_parameters(self):
|
||||
"""Calculate losses, gradients, and update network weights; called in every training iteration"""
|
||||
pass
|
||||
|
||||
def setup(self, opt):
|
||||
"""Load and print networks; create schedulers
|
||||
|
||||
Parameters:
|
||||
opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
|
||||
"""
|
||||
if self.isTrain:
|
||||
self.schedulers = [networks.get_scheduler(optimizer, opt) for optimizer in self.optimizers]
|
||||
if not self.isTrain or opt.continue_train:
|
||||
load_suffix = 'iter_%d' % opt.load_iter if opt.load_iter > 0 else opt.epoch
|
||||
self.load_networks(load_suffix)
|
||||
self.print_networks(opt.verbose)
|
||||
|
||||
def eval(self):
|
||||
"""Make models eval mode during test time"""
|
||||
for name in self.model_names:
|
||||
if isinstance(name, str):
|
||||
net = getattr(self, 'net' + name)
|
||||
net.eval()
|
||||
|
||||
def test(self):
|
||||
"""Forward function used in test time.
|
||||
|
||||
This function wraps <forward> function in no_grad() so we don't save intermediate steps for backprop
|
||||
It also calls <compute_visuals> to produce additional visualization results
|
||||
"""
|
||||
with torch.no_grad():
|
||||
self.forward()
|
||||
self.compute_visuals()
|
||||
|
||||
def compute_visuals(self):
|
||||
"""Calculate additional output images for visdom and HTML visualization"""
|
||||
pass
|
||||
|
||||
def get_image_paths(self):
|
||||
""" Return image paths that are used to load current data"""
|
||||
return self.image_paths
|
||||
|
||||
def update_learning_rate(self):
|
||||
"""Update learning rates for all the networks; called at the end of every epoch"""
|
||||
for scheduler in self.schedulers:
|
||||
if self.opt.lr_policy == 'plateau':
|
||||
scheduler.step(self.metric)
|
||||
else:
|
||||
scheduler.step()
|
||||
|
||||
lr = self.optimizers[0].param_groups[0]['lr']
|
||||
print('learning rate = %.7f' % lr)
|
||||
|
||||
def get_lr(self,):
|
||||
lrs = {}
|
||||
for idx, p in enumerate(self.optimizers):
|
||||
lrs['LR{}'.format(idx)] = p.param_groups[0]['lr']
|
||||
return lrs
|
||||
|
||||
def get_current_visuals(self):
|
||||
"""Return visualization images. train.py will display these images with visdom, and save the images to a HTML"""
|
||||
visual_ret = OrderedDict()
|
||||
for name in self.visual_names:
|
||||
if isinstance(name, str):
|
||||
visual_ret[name] = getattr(self, name)
|
||||
return visual_ret
|
||||
|
||||
def get_current_losses(self):
|
||||
"""Return traning losses / errors. train.py will print out these errors on console, and save them to a file"""
|
||||
errors_ret = OrderedDict()
|
||||
for name in self.loss_names:
|
||||
if isinstance(name, str):
|
||||
errors_ret['Loss_' + name] = float(getattr(self, 'loss_' + name)) # float(...) works for both scalar tensor and float number
|
||||
return errors_ret
|
||||
|
||||
def save_networks(self, epoch, info=None):
|
||||
"""Save all the networks to the disk.
|
||||
|
||||
Parameters:
|
||||
epoch (int) -- current epoch; used in the file name '%s_net_%s.pth' % (epoch, name)
|
||||
"""
|
||||
for name in self.model_names:
|
||||
if isinstance(name, str):
|
||||
save_filename = '%s_net_%s.pth' % (epoch, name)
|
||||
save_path = os.path.join(self.save_dir, save_filename)
|
||||
net = getattr(self, 'net' + name)
|
||||
|
||||
if len(self.gpu_ids) > 0 and torch.cuda.is_available():
|
||||
torch.save(net.module.cpu().state_dict(), save_path)
|
||||
print('Model saved in:', save_filename)
|
||||
net.cuda(self.gpu_ids[0])
|
||||
else:
|
||||
torch.save(net.cpu().state_dict(), save_path)
|
||||
if info is not None:
|
||||
torch.save(info, os.path.join(self.save_dir, '%s.info' % epoch))
|
||||
|
||||
def __patch_instance_norm_state_dict(self, state_dict, module, keys, i=0):
|
||||
"""Fix InstanceNorm checkpoints incompatibility (prior to 0.4)"""
|
||||
key = keys[i]
|
||||
if i + 1 == len(keys): # at the end, pointing to a parameter/buffer
|
||||
if module.__class__.__name__.startswith('InstanceNorm') and \
|
||||
(key == 'running_mean' or key == 'running_var'):
|
||||
if getattr(module, key) is None:
|
||||
state_dict.pop('.'.join(keys))
|
||||
if module.__class__.__name__.startswith('InstanceNorm') and \
|
||||
(key == 'num_batches_tracked'):
|
||||
state_dict.pop('.'.join(keys))
|
||||
else:
|
||||
self.__patch_instance_norm_state_dict(state_dict, getattr(module, key), keys, i + 1)
|
||||
|
||||
def load_networks(self, epoch):
|
||||
"""Load all the networks from the disk.
|
||||
|
||||
Parameters:
|
||||
epoch (int) -- current epoch; used in the file name '%s_net_%s.pth' % (epoch, name)
|
||||
"""
|
||||
for name in self.load_model_names:
|
||||
if isinstance(name, str):
|
||||
load_filename = '%s_net_%s.pth' % (epoch, name)
|
||||
load_path = os.path.join(self.save_dir, load_filename)
|
||||
net = getattr(self, 'net' + name)
|
||||
if isinstance(net, torch.nn.DataParallel) or isinstance(net, torch.nn.parallel.DistributedDataParallel):
|
||||
net = net.module
|
||||
print('loading the model from %s' % load_path)
|
||||
# if you are using PyTorch newer than 0.4 (e.g., built from
|
||||
# GitHub source), you can remove str() on self.device
|
||||
map_location = str(self.device)
|
||||
|
||||
state_dict = torch.load(load_path, map_location=map_location)
|
||||
|
||||
# patch InstanceNorm checkpoints prior to 0.4
|
||||
# for key in list(state_dict.keys()): # need to copy keys here because we mutate in loop
|
||||
# self.__patch_instance_norm_state_dict(state_dict, net, key.split('.'))
|
||||
# net.load_state_dict(state_dict)
|
||||
if not self.opt.no_strict_load:
|
||||
net.load_state_dict(state_dict)
|
||||
# Load partial weights
|
||||
else:
|
||||
model_dict = net.state_dict()
|
||||
pretrained_dict = {k: v for k, v in state_dict.items() if k in model_dict}
|
||||
model_dict.update(pretrained_dict)
|
||||
net.load_state_dict(model_dict, strict=False)
|
||||
|
||||
info_path = os.path.join(self.save_dir, '%s.info' % epoch)
|
||||
if os.path.exists(info_path):
|
||||
info_dict = torch.load(info_path)
|
||||
for k, v in info_dict.items():
|
||||
setattr(self.opt, k, v)
|
||||
|
||||
def print_networks(self, verbose):
|
||||
"""Print the total number of parameters in the network and (if verbose) network architecture
|
||||
|
||||
Parameters:
|
||||
verbose (bool) -- if verbose: print the network architecture
|
||||
"""
|
||||
print('---------- Networks initialized -------------')
|
||||
for name in self.model_names:
|
||||
if isinstance(name, str):
|
||||
net = getattr(self, 'net' + name)
|
||||
num_params = 0
|
||||
for param in net.parameters():
|
||||
num_params += param.numel()
|
||||
print('[Network %s] Total number of parameters : %.3f M' % (name, num_params / 1e6))
|
||||
print('-----------------------------------------------')
|
||||
|
||||
def set_requires_grad(self, nets, requires_grad=False):
|
||||
"""Set requies_grad=Fasle for all the networks to avoid unnecessary computations
|
||||
Parameters:
|
||||
nets (network list) -- a list of networks
|
||||
requires_grad (bool) -- whether the networks require gradients or not
|
||||
"""
|
||||
if not isinstance(nets, list):
|
||||
nets = [nets]
|
||||
for net in nets:
|
||||
if net is not None:
|
||||
for param in net.parameters():
|
||||
param.requires_grad = requires_grad
|
||||
@@ -0,0 +1,131 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.nn.parameter import Parameter
|
||||
from torch.nn import functional as F
|
||||
import numpy as np
|
||||
|
||||
|
||||
class NormLayer(nn.Module):
|
||||
"""Normalization Layers.
|
||||
------------
|
||||
# Arguments
|
||||
- channels: input channels, for batch norm and instance norm.
|
||||
- input_size: input shape without batch size, for layer norm.
|
||||
"""
|
||||
def __init__(self, channels, normalize_shape=None, norm_type='bn'):
|
||||
super(NormLayer, self).__init__()
|
||||
norm_type = norm_type.lower()
|
||||
self.norm_type = norm_type
|
||||
self.channels = channels
|
||||
if norm_type == 'bn':
|
||||
self.norm = nn.BatchNorm2d(channels, affine=True)
|
||||
elif norm_type == 'in':
|
||||
self.norm = nn.InstanceNorm2d(channels, affine=False)
|
||||
elif norm_type == 'gn':
|
||||
self.norm = nn.GroupNorm(32, channels, affine=True)
|
||||
elif norm_type == 'pixel':
|
||||
self.norm = lambda x: F.normalize(x, p=2, dim=1)
|
||||
elif norm_type == 'layer':
|
||||
self.norm = nn.LayerNorm(normalize_shape)
|
||||
elif norm_type == 'none':
|
||||
self.norm = lambda x: x*1.0
|
||||
else:
|
||||
assert 1==0, 'Norm type {} not support.'.format(norm_type)
|
||||
|
||||
def forward(self, x, ref=None):
|
||||
return self.norm(x)
|
||||
|
||||
|
||||
class ReluLayer(nn.Module):
|
||||
"""Relu Layer.
|
||||
------------
|
||||
# Arguments
|
||||
- relu type: type of relu layer, candidates are
|
||||
- ReLU
|
||||
- LeakyReLU: default relu slope 0.2
|
||||
- PRelu
|
||||
- SELU
|
||||
- none: direct pass
|
||||
"""
|
||||
def __init__(self, channels, relu_type='relu'):
|
||||
super(ReluLayer, self).__init__()
|
||||
relu_type = relu_type.lower()
|
||||
if relu_type == 'relu':
|
||||
self.func = nn.ReLU(True)
|
||||
elif relu_type == 'leakyrelu':
|
||||
self.func = nn.LeakyReLU(0.2, inplace=True)
|
||||
elif relu_type == 'prelu':
|
||||
self.func = nn.PReLU(channels)
|
||||
elif relu_type == 'selu':
|
||||
self.func = nn.SELU(True)
|
||||
elif relu_type == 'none':
|
||||
self.func = lambda x: x*1.0
|
||||
else:
|
||||
assert 1==0, 'Relu type {} not support.'.format(relu_type)
|
||||
|
||||
def forward(self, x):
|
||||
return self.func(x)
|
||||
|
||||
|
||||
class ConvLayer(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, kernel_size=3, scale='none', norm_type='none', relu_type='none', use_pad=True, bias=True):
|
||||
super(ConvLayer, self).__init__()
|
||||
self.use_pad = use_pad
|
||||
self.norm_type = norm_type
|
||||
self.in_channels = in_channels
|
||||
if norm_type in ['bn']:
|
||||
bias = False
|
||||
|
||||
stride = 2 if scale == 'down' else 1
|
||||
self.scale = scale
|
||||
|
||||
self.scale_func = lambda x: x
|
||||
if scale == 'up':
|
||||
self.scale_func = lambda x: nn.functional.interpolate(x, scale_factor=2, mode='nearest')
|
||||
|
||||
self.reflection_pad = nn.ReflectionPad2d(int(np.ceil((kernel_size - 1.)/2)))
|
||||
self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride, bias=bias)
|
||||
|
||||
self.avgpool = nn.AvgPool2d(2, 2)
|
||||
self.relu = ReluLayer(out_channels, relu_type)
|
||||
self.norm = NormLayer(out_channels, norm_type=norm_type)
|
||||
|
||||
def forward(self, x):
|
||||
out = self.scale_func(x)
|
||||
if self.use_pad:
|
||||
out = self.reflection_pad(out)
|
||||
out = self.conv2d(out)
|
||||
if self.scale == 'down_avg':
|
||||
out = self.avgpool(out)
|
||||
out = self.norm(out)
|
||||
out = self.relu(out)
|
||||
return out
|
||||
|
||||
|
||||
class ResidualBlock(nn.Module):
|
||||
"""
|
||||
Residual block recommended in: http://torch.ch/blog/2016/02/04/resnets.html
|
||||
"""
|
||||
def __init__(self, c_in, c_out, relu_type='prelu', norm_type='bn', scale='none'):
|
||||
super(ResidualBlock, self).__init__()
|
||||
|
||||
if scale == 'none' and c_in == c_out:
|
||||
self.shortcut_func = lambda x: x
|
||||
else:
|
||||
self.shortcut_func = ConvLayer(c_in, c_out, 3, scale)
|
||||
|
||||
scale_config_dict = {'down': ['none', 'down'], 'up': ['up', 'none'], 'none': ['none', 'none']}
|
||||
scale_conf = scale_config_dict[scale]
|
||||
|
||||
self.conv1 = ConvLayer(c_in, c_out, 3, scale_conf[0], norm_type=norm_type, relu_type=relu_type)
|
||||
self.conv2 = ConvLayer(c_out, c_out, 3, scale_conf[1], norm_type=norm_type, relu_type='none')
|
||||
|
||||
def forward(self, x):
|
||||
identity = self.shortcut_func(x)
|
||||
|
||||
res = self.conv1(x)
|
||||
res = self.conv2(res)
|
||||
return identity + res
|
||||
|
||||
|
||||
@@ -0,0 +1,157 @@
|
||||
import os
|
||||
import numpy as np
|
||||
import collections
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.optim as optim
|
||||
|
||||
from models import loss
|
||||
from models import networks
|
||||
from .base_model import BaseModel
|
||||
from utils import utils
|
||||
|
||||
class EnhanceModel(BaseModel):
|
||||
|
||||
def modify_commandline_options(parser, is_train):
|
||||
if is_train:
|
||||
parser.add_argument('--parse_net_weight', type=str, default='./pretrain_models/parse_multi_iter_90000.pth', help='parse model path')
|
||||
parser.add_argument('--lambda_pix', type=float, default=10.0, help='weight for parsing map')
|
||||
parser.add_argument('--lambda_pcp', type=float, default=0.0, help='weight for vgg perceptual loss')
|
||||
parser.add_argument('--lambda_fm', type=float, default=10.0, help='weight for sr')
|
||||
parser.add_argument('--lambda_g', type=float, default=1.0, help='weight for sr')
|
||||
parser.add_argument('--lambda_ss', type=float, default=1000., help='weight for global style')
|
||||
return parser
|
||||
|
||||
def __init__(self, opt):
|
||||
BaseModel.__init__(self, opt)
|
||||
|
||||
self.netP = networks.define_P(opt, weight_path=opt.parse_net_weight)
|
||||
self.netG = networks.define_G(opt, use_norm='spectral_norm')
|
||||
|
||||
if self.isTrain:
|
||||
self.netD = networks.define_D(opt, opt.Dinput_nc, use_norm='spectral_norm')
|
||||
self.vgg_model = loss.PCPFeat(weight_path='./pretrain_models/vgg19-dcbb9e9d.pth').to(opt.device)
|
||||
if len(opt.gpu_ids) > 0:
|
||||
self.vgg_model = torch.nn.DataParallel(self.vgg_model, opt.gpu_ids, output_device=opt.device)
|
||||
|
||||
self.model_names = ['G']
|
||||
self.loss_names = ['Pix', 'PCP', 'G', 'FM', 'D', 'SS'] # Generator loss, fm loss, parsing loss, discriminator loss
|
||||
self.visual_names = ['img_LR', 'img_HR', 'img_SR', 'ref_Parse', 'hr_mask']
|
||||
self.fm_weights = [1**x for x in range(opt.D_num)]
|
||||
|
||||
if self.isTrain:
|
||||
self.model_names = ['G', 'D']
|
||||
self.load_model_names = ['G', 'D']
|
||||
|
||||
self.criterionParse = torch.nn.CrossEntropyLoss().to(opt.device)
|
||||
self.criterionFM = loss.FMLoss().to(opt.device)
|
||||
self.criterionGAN = loss.GANLoss(opt.gan_mode).to(opt.device)
|
||||
self.criterionPCP = loss.PCPLoss(opt)
|
||||
self.criterionPix= nn.L1Loss()
|
||||
self.criterionRS = loss.RegionStyleLoss()
|
||||
|
||||
self.optimizer_G = optim.Adam([p for p in self.netG.parameters() if p.requires_grad], lr=opt.g_lr, betas=(opt.beta1, 0.999))
|
||||
self.optimizer_D = optim.Adam([p for p in self.netD.parameters() if p.requires_grad], lr=opt.d_lr, betas=(opt.beta1, 0.999))
|
||||
self.optimizers = [self.optimizer_G, self.optimizer_D]
|
||||
|
||||
def eval(self):
|
||||
self.netG.eval()
|
||||
self.netP.eval()
|
||||
|
||||
def load_pretrain_models(self,):
|
||||
self.netP.eval()
|
||||
print('Loading pretrained LQ face parsing network from', self.opt.parse_net_weight)
|
||||
if len(self.opt.gpu_ids) > 0:
|
||||
self.netP.module.load_state_dict(torch.load(self.opt.parse_net_weight))
|
||||
else:
|
||||
self.netP.load_state_dict(torch.load(self.opt.parse_net_weight))
|
||||
self.netG.eval()
|
||||
print('Loading pretrained PSFRGAN from', self.opt.psfr_net_weight)
|
||||
if len(self.opt.gpu_ids) > 0:
|
||||
self.netG.module.load_state_dict(torch.load(self.opt.psfr_net_weight), strict=False)
|
||||
else:
|
||||
self.netG.load_state_dict(torch.load(self.opt.psfr_net_weight), strict=False)
|
||||
|
||||
def set_input(self, input, cur_iters=None):
|
||||
self.cur_iters = cur_iters
|
||||
self.img_LR = input['LR'].to(self.opt.device)
|
||||
self.img_HR = input['HR'].to(self.opt.device)
|
||||
self.hr_mask = input['Mask'].to(self.opt.device)
|
||||
if self.opt.debug:
|
||||
print('SRNet input shape:', self.img_LR.shape, self.img_HR.shape)
|
||||
|
||||
def forward(self):
|
||||
with torch.no_grad():
|
||||
ref_mask, _ = self.netP(self.img_LR)
|
||||
self.ref_mask_onehot = (ref_mask == ref_mask.max(dim=1, keepdim=True)[0]).float().detach()
|
||||
|
||||
if self.opt.debug:
|
||||
print('SRNet reference mask shape:', self.ref_mask_onehot.shape)
|
||||
self.img_SR = self.netG(self.img_LR, self.ref_mask_onehot)
|
||||
|
||||
self.real_D_results = self.netD(torch.cat((self.img_HR, self.hr_mask), dim=1), return_feat=True)
|
||||
self.fake_D_results = self.netD(torch.cat((self.img_SR.detach(), self.hr_mask), dim=1), return_feat=False)
|
||||
self.fake_G_results = self.netD(torch.cat((self.img_SR, self.hr_mask), dim=1), return_feat=True)
|
||||
|
||||
self.img_SR_feats = self.vgg_model(self.img_SR)
|
||||
self.img_HR_feats = self.vgg_model(self.img_HR)
|
||||
|
||||
def backward_G(self):
|
||||
# Pix Loss
|
||||
self.loss_Pix = self.criterionPix(self.img_SR, self.img_HR) * self.opt.lambda_pix
|
||||
|
||||
# semantic style loss
|
||||
self.loss_SS = self.criterionRS(self.img_SR_feats, self.img_HR_feats, self.hr_mask) * self.opt.lambda_ss
|
||||
|
||||
# perceptual loss
|
||||
self.loss_PCP = self.criterionPCP(self.img_SR_feats, self.img_HR_feats) * self.opt.lambda_pcp
|
||||
|
||||
# Feature matching loss
|
||||
tmp_loss = 0
|
||||
for i, w in zip(range(self.opt.D_num), self.fm_weights):
|
||||
tmp_loss = tmp_loss + self.criterionFM(self.fake_G_results[i][1], self.real_D_results[i][1]) * w
|
||||
self.loss_FM = tmp_loss * self.opt.lambda_fm / self.opt.D_num
|
||||
|
||||
# Generator loss
|
||||
tmp_loss = 0
|
||||
for i in range(self.opt.D_num):
|
||||
tmp_loss = tmp_loss + self.criterionGAN(self.fake_G_results[i][0], True, for_discriminator=False)
|
||||
self.loss_G = tmp_loss * self.opt.lambda_g / self.opt.D_num
|
||||
|
||||
total_loss = self.loss_Pix + self.loss_PCP + self.loss_FM + self.loss_G + self.loss_SS
|
||||
total_loss.backward()
|
||||
|
||||
def backward_D(self, ):
|
||||
self.loss_D = 0
|
||||
for i in range(self.opt.D_num):
|
||||
self.loss_D += 0.5 * (self.criterionGAN(self.fake_D_results[i], False) + self.criterionGAN(self.real_D_results[i][0], True))
|
||||
self.loss_D /= self.opt.D_num
|
||||
self.loss_D.backward()
|
||||
|
||||
def optimize_parameters(self, ):
|
||||
# ---- Update G ------------
|
||||
self.optimizer_G.zero_grad()
|
||||
self.backward_G()
|
||||
self.optimizer_G.step()
|
||||
|
||||
# ---- Update D ------------
|
||||
self.optimizer_D.zero_grad()
|
||||
self.backward_D()
|
||||
self.optimizer_D.step()
|
||||
|
||||
def get_current_visuals(self, size=512):
|
||||
out = []
|
||||
visual_imgs = []
|
||||
out.append(utils.tensor_to_numpy(self.img_LR))
|
||||
out.append(utils.tensor_to_numpy(self.img_SR))
|
||||
out.append(utils.tensor_to_numpy(self.img_HR))
|
||||
|
||||
out_imgs = [utils.batch_numpy_to_image(x, size) for x in out]
|
||||
|
||||
visual_imgs += out_imgs
|
||||
visual_imgs.append(utils.color_parse_map(self.ref_mask_onehot, size))
|
||||
visual_imgs.append(utils.color_parse_map(self.hr_mask, size))
|
||||
|
||||
return visual_imgs
|
||||
|
||||
@@ -0,0 +1,224 @@
|
||||
import torch
|
||||
from torchvision import models
|
||||
from utils import utils
|
||||
from torch import nn
|
||||
|
||||
|
||||
def tv_loss(x):
|
||||
"""
|
||||
Total Variation Loss.
|
||||
"""
|
||||
return torch.sum(torch.abs(x[:, :, :, :-1] - x[:, :, :, 1:])
|
||||
) + torch.sum(torch.abs(x[:, :, :-1, :] - x[:, :, 1:, :]))
|
||||
|
||||
|
||||
class PCPFeat(torch.nn.Module):
|
||||
"""
|
||||
Features used to calculate Perceptual Loss based on ResNet50 features.
|
||||
Input: (B, C, H, W), RGB, [0, 1]
|
||||
"""
|
||||
def __init__(self, weight_path, model='vgg'):
|
||||
super(PCPFeat, self).__init__()
|
||||
if model == 'vgg':
|
||||
self.model = models.vgg19(pretrained=False)
|
||||
self.build_vgg_layers()
|
||||
elif model == 'resnet':
|
||||
self.model = models.resnet50(pretrained=False)
|
||||
self.build_resnet_layers()
|
||||
|
||||
self.model.load_state_dict(torch.load(weight_path))
|
||||
self.model.eval()
|
||||
for param in self.model.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
def build_resnet_layers(self):
|
||||
self.layer1 = torch.nn.Sequential(
|
||||
self.model.conv1,
|
||||
self.model.bn1,
|
||||
self.model.relu,
|
||||
self.model.maxpool,
|
||||
self.model.layer1
|
||||
)
|
||||
self.layer2 = self.model.layer2
|
||||
self.layer3 = self.model.layer3
|
||||
self.layer4 = self.model.layer4
|
||||
self.features = torch.nn.ModuleList(
|
||||
[self.layer1, self.layer2, self.layer3, self.layer4]
|
||||
)
|
||||
|
||||
def build_vgg_layers(self):
|
||||
vgg_pretrained_features = self.model.features
|
||||
self.features = []
|
||||
feature_layers = [0, 3, 8, 17, 26, 35]
|
||||
for i in range(len(feature_layers)-1):
|
||||
module_layers = torch.nn.Sequential()
|
||||
for j in range(feature_layers[i], feature_layers[i+1]):
|
||||
module_layers.add_module(str(j), vgg_pretrained_features[j])
|
||||
self.features.append(module_layers)
|
||||
self.features = torch.nn.ModuleList(self.features)
|
||||
|
||||
def preprocess(self, x):
|
||||
x = (x + 1) / 2
|
||||
mean = torch.Tensor([0.485, 0.456, 0.406]).to(x)
|
||||
std = torch.Tensor([0.229, 0.224, 0.225]).to(x)
|
||||
mean = mean.view(1, 3, 1, 1)
|
||||
std = std.view(1, 3, 1, 1)
|
||||
x = (x - mean) / std
|
||||
if x.shape[3] < 224:
|
||||
x = torch.nn.functional.interpolate(x, size=(224, 224), mode='bilinear', align_corners=False)
|
||||
return x
|
||||
|
||||
def forward(self, x):
|
||||
x = self.preprocess(x)
|
||||
|
||||
features = []
|
||||
for m in self.features:
|
||||
x = m(x)
|
||||
features.append(x)
|
||||
return features
|
||||
|
||||
|
||||
class PCPLoss(torch.nn.Module):
|
||||
"""Perceptual Loss.
|
||||
"""
|
||||
def __init__(self,
|
||||
opt,
|
||||
layer=5,
|
||||
model='vgg',
|
||||
):
|
||||
super(PCPLoss, self).__init__()
|
||||
|
||||
self.mse = torch.nn.MSELoss()
|
||||
self.weights = [1.0/32, 1.0/16, 1.0/8, 1.0/4, 1.0]
|
||||
|
||||
def forward(self, x_feats, y_feats):
|
||||
loss = 0
|
||||
for xf, yf, w in zip(x_feats, y_feats, self.weights):
|
||||
loss = loss + self.mse(xf, yf.detach()) * w
|
||||
return loss
|
||||
|
||||
|
||||
class FMLoss(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.mse = torch.nn.MSELoss()
|
||||
|
||||
def forward(self, x_feats, y_feats):
|
||||
loss = 0
|
||||
for xf, yf in zip(x_feats, y_feats):
|
||||
loss = loss + self.mse(xf, yf.detach())
|
||||
return loss
|
||||
|
||||
|
||||
class GANLoss(nn.Module):
|
||||
def __init__(self, gan_mode, target_real_label=1.0, target_fake_label=0.0):
|
||||
""" Initialize the GANLoss class.
|
||||
Parameters:
|
||||
gan_mode (str) - - the type of GAN objective. It currently supports vanilla, lsgan, and wgangp.
|
||||
target_real_label (bool) - - label for a real image
|
||||
target_fake_label (bool) - - label of a fake image
|
||||
Note: Do not use sigmoid as the last layer of Discriminator.
|
||||
LSGAN needs no sigmoid. vanilla GANs will handle it with BCEWithLogitsLoss.
|
||||
"""
|
||||
super(GANLoss, self).__init__()
|
||||
self.register_buffer('real_label', torch.tensor(target_real_label))
|
||||
self.register_buffer('fake_label', torch.tensor(target_fake_label))
|
||||
self.gan_mode = gan_mode
|
||||
if gan_mode == 'lsgan':
|
||||
self.loss = nn.MSELoss()
|
||||
elif gan_mode == 'vanilla':
|
||||
self.loss = nn.BCEWithLogitsLoss()
|
||||
elif gan_mode == 'hinge':
|
||||
pass
|
||||
elif gan_mode in ['wgangp']:
|
||||
self.loss = None
|
||||
else:
|
||||
raise NotImplementedError('gan mode %s not implemented' % gan_mode)
|
||||
|
||||
def get_target_tensor(self, prediction, target_is_real):
|
||||
if target_is_real:
|
||||
target_tensor = self.real_label
|
||||
else:
|
||||
target_tensor = self.fake_label
|
||||
return target_tensor.expand_as(prediction)
|
||||
|
||||
def __call__(self, prediction, target_is_real, for_discriminator=True):
|
||||
"""Calculate loss given Discriminator's output and grount truth labels.
|
||||
Parameters:
|
||||
prediction (tensor) - - tpyically the prediction output from a discriminator
|
||||
target_is_real (bool) - - if the ground truth label is for real images or fake images
|
||||
Returns:
|
||||
the calculated loss.
|
||||
"""
|
||||
if self.gan_mode in ['lsgan', 'vanilla']:
|
||||
target_tensor = self.get_target_tensor(prediction, target_is_real)
|
||||
loss = self.loss(prediction, target_tensor)
|
||||
elif self.gan_mode == 'hinge':
|
||||
if for_discriminator:
|
||||
if target_is_real:
|
||||
loss = nn.ReLU()(1 - prediction).mean()
|
||||
else:
|
||||
loss = nn.ReLU()(1 + prediction).mean()
|
||||
else:
|
||||
assert target_is_real, "The generator's hinge loss must be aiming for real"
|
||||
loss = - prediction.mean()
|
||||
return loss
|
||||
|
||||
elif self.gan_mode == 'wgangp':
|
||||
if target_is_real:
|
||||
loss = -prediction.mean()
|
||||
else:
|
||||
loss = prediction.mean()
|
||||
return loss
|
||||
|
||||
|
||||
class RegionStyleLoss(nn.Module):
|
||||
def __init__(self, reg_num=19, eps=1e-8):
|
||||
super().__init__()
|
||||
self.reg_num = reg_num
|
||||
self.eps = eps
|
||||
self.mse = nn.MSELoss()
|
||||
|
||||
def __masked_gram_matrix(self, x, m):
|
||||
b, c, h, w = x.shape
|
||||
m = m.view(b, -1, h*w)
|
||||
x = x.view(b, -1, h*w)
|
||||
total_elements = m.sum(2) + self.eps
|
||||
|
||||
x = x * m
|
||||
G = torch.bmm(x, x.transpose(1, 2))
|
||||
return G / (c * total_elements.view(b, 1, 1))
|
||||
|
||||
def __layer_gram_matrix(self, x, mask):
|
||||
b, c, h, w = x.shape
|
||||
all_gm = []
|
||||
for i in range(self.reg_num):
|
||||
sub_mask = mask[:, i].unsqueeze(1)
|
||||
gram_matrix = self.__masked_gram_matrix(x, sub_mask)
|
||||
all_gm.append(gram_matrix)
|
||||
return torch.stack(all_gm, dim=1)
|
||||
|
||||
def forward(self, x_feats, y_feats, mask):
|
||||
loss = 0
|
||||
for xf, yf in zip(x_feats[2:], y_feats[2:]):
|
||||
tmp_mask = torch.nn.functional.interpolate(mask, xf.shape[2:])
|
||||
xf_gm = self.__layer_gram_matrix(xf, tmp_mask)
|
||||
yf_gm = self.__layer_gram_matrix(yf, tmp_mask)
|
||||
tmp_loss = self.mse(xf_gm, yf_gm.detach())
|
||||
loss = loss + tmp_loss
|
||||
return loss
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
x = [
|
||||
torch.randn(2, 64, 512, 512),
|
||||
torch.randn(2, 128, 256, 256),
|
||||
torch.randn(2, 256, 128, 128),
|
||||
torch.randn(2, 512, 64, 64),
|
||||
torch.randn(2, 512, 32, 32),
|
||||
]
|
||||
|
||||
y = torch.randint(10, (2, 19, 512, 512)).float()
|
||||
loss = RegionStyleLoss()
|
||||
print(loss(x, x, y))
|
||||
|
||||
@@ -0,0 +1,254 @@
|
||||
from models.blocks import *
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.nn import init
|
||||
from torch.optim import lr_scheduler
|
||||
from utils import utils
|
||||
import numpy as np
|
||||
|
||||
from models import psfrnet
|
||||
import torch.nn.utils as tutils
|
||||
from models.loss import PCPFeat
|
||||
|
||||
|
||||
def apply_norm(net, weight_norm_type):
|
||||
for m in net.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
if weight_norm_type.lower() == 'spectral_norm':
|
||||
tutils.spectral_norm(m)
|
||||
elif weight_norm_type.lower() == 'weight_norm':
|
||||
tutils.weight_norm(m)
|
||||
else:
|
||||
pass
|
||||
|
||||
|
||||
def init_weights(net, init_type='normal', init_gain=0.02):
|
||||
"""Initialize network weights.
|
||||
Parameters:
|
||||
net (network) -- network to be initialized
|
||||
init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal
|
||||
init_gain (float) -- scaling factor for normal, xavier and orthogonal.
|
||||
We use 'normal' in the original pix2pix and CycleGAN paper. But xavier and kaiming might
|
||||
work better for some applications. Feel free to try yourself.
|
||||
"""
|
||||
def init_func(m): # define the initialization function
|
||||
classname = m.__class__.__name__
|
||||
if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):
|
||||
if init_type == 'normal':
|
||||
init.normal_(m.weight.data, 0.0, init_gain)
|
||||
elif init_type == 'xavier':
|
||||
init.xavier_normal_(m.weight.data, gain=init_gain)
|
||||
elif init_type == 'kaiming':
|
||||
init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
|
||||
elif init_type == 'orthogonal':
|
||||
init.orthogonal_(m.weight.data, gain=init_gain)
|
||||
else:
|
||||
raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
|
||||
if hasattr(m, 'bias') and m.bias is not None:
|
||||
init.constant_(m.bias.data, 0.0)
|
||||
elif classname.find('BatchNorm2d') != -1: # BatchNorm Layer's weight is not a matrix; only normal distribution applies.
|
||||
init.normal_(m.weight.data, 1.0, init_gain)
|
||||
init.constant_(m.bias.data, 0.0)
|
||||
|
||||
print('initialize network with %s' % init_type)
|
||||
net.apply(init_func) # apply the initialization function <init_func>
|
||||
|
||||
|
||||
def init_net(net, init_type='normal', init_gain=0.02, gpu_ids=[]):
|
||||
"""Initialize a network: 1. register CPU/GPU device (with multi-GPU support); 2. initialize the network weights
|
||||
Parameters:
|
||||
net (network) -- the network to be initialized
|
||||
init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal
|
||||
gain (float) -- scaling factor for normal, xavier and orthogonal.
|
||||
gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2
|
||||
Return an initialized network.
|
||||
"""
|
||||
if len(gpu_ids) > 0:
|
||||
assert(torch.cuda.is_available())
|
||||
net.to(gpu_ids[0])
|
||||
net = torch.nn.DataParallel(net, gpu_ids) # multi-GPUs
|
||||
init_weights(net, init_type, init_gain=init_gain)
|
||||
return net
|
||||
|
||||
|
||||
def get_scheduler(optimizer, opt):
|
||||
"""Return a learning rate scheduler
|
||||
Parameters:
|
||||
optimizer -- the optimizer of the network
|
||||
opt (option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions．
|
||||
opt.lr_policy is the name of learning rate policy: linear | step | plateau | cosine
|
||||
For 'linear', we keep the same learning rate for the first <opt.n_epochs> epochs
|
||||
and linearly decay the rate to zero over the next <opt.n_epochs_decay> epochs.
|
||||
For other schedulers (step, plateau, and cosine), we use the default PyTorch schedulers.
|
||||
See https://pytorch.org/docs/stable/optim.html for more details.
|
||||
"""
|
||||
if opt.lr_policy == 'linear':
|
||||
def lambda_rule(epoch):
|
||||
lr_l = 1.0 - max(0, epoch + opt.epoch_count - opt.n_epochs) / float(opt.n_epochs_decay + 1)
|
||||
return lr_l
|
||||
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule)
|
||||
elif opt.lr_policy == 'step':
|
||||
scheduler = lr_scheduler.StepLR(optimizer, step_size=opt.lr_decay_iters, gamma=0.1)
|
||||
elif opt.lr_policy == 'plateau':
|
||||
scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.2, threshold=0.01, patience=5)
|
||||
elif opt.lr_policy == 'cosine':
|
||||
scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=opt.n_epochs, eta_min=0)
|
||||
else:
|
||||
return NotImplementedError('learning rate policy [%s] is not implemented', opt.lr_policy)
|
||||
return scheduler
|
||||
|
||||
|
||||
def define_P(opt, in_size=512, out_size=512, min_feat_size=32, relu_type='LeakyReLU', isTrain=True, weight_path=None):
|
||||
net = ParseNet(in_size, out_size, min_feat_size, 64, 19, norm_type=opt.Pnorm, relu_type=relu_type, ch_range=[32, 256])
|
||||
if not isTrain:
|
||||
net.eval()
|
||||
if weight_path is not None:
|
||||
net.load_state_dict(torch.load(weight_path))
|
||||
if len(opt.gpu_ids) > 0:
|
||||
assert(torch.cuda.is_available())
|
||||
net.to(opt.device)
|
||||
net = torch.nn.DataParallel(net, opt.gpu_ids, output_device=opt.device)
|
||||
return net
|
||||
|
||||
|
||||
def define_G(opt, isTrain=True, use_norm='none', relu_type='LeakyReLU'):
|
||||
net = psfrnet.PSFRGenerator(3, 3, in_size=opt.Gin_size, out_size=opt.Gout_size, relu_type=relu_type, parse_ch=19, norm_type=opt.Gnorm)
|
||||
apply_norm(net, use_norm)
|
||||
if not isTrain:
|
||||
net.eval()
|
||||
if len(opt.gpu_ids) > 0:
|
||||
assert(torch.cuda.is_available())
|
||||
net.to(opt.device)
|
||||
net = torch.nn.DataParallel(net, opt.gpu_ids, output_device=opt.device)
|
||||
# init_weights(net, init_type='normal', init_gain=0.02)
|
||||
return net
|
||||
|
||||
|
||||
def define_D(opt, in_channel=3, isTrain=True, use_norm='none'):
|
||||
net = MultiScaleDiscriminator(in_channel, opt.ndf, opt.n_layers_D, opt.Dnorm, num_D=opt.D_num)
|
||||
apply_norm(net, use_norm)
|
||||
if not isTrain:
|
||||
net.eval()
|
||||
if len(opt.gpu_ids) > 0:
|
||||
assert(torch.cuda.is_available())
|
||||
net.to(opt.device)
|
||||
net = torch.nn.DataParallel(net, opt.gpu_ids, output_device=opt.device)
|
||||
init_weights(net, init_type='normal', init_gain=0.02)
|
||||
return net
|
||||
|
||||
|
||||
class ParseNet(nn.Module):
|
||||
def __init__(self,
|
||||
in_size=128,
|
||||
out_size=128,
|
||||
min_feat_size=32,
|
||||
base_ch=64,
|
||||
parsing_ch=19,
|
||||
res_depth=10,
|
||||
relu_type='prelu',
|
||||
norm_type='bn',
|
||||
ch_range=[32, 512],
|
||||
):
|
||||
super().__init__()
|
||||
self.res_depth = res_depth
|
||||
act_args = {'norm_type': norm_type, 'relu_type': relu_type}
|
||||
min_ch, max_ch = ch_range
|
||||
|
||||
ch_clip = lambda x: max(min_ch, min(x, max_ch))
|
||||
min_feat_size = min(in_size, min_feat_size)
|
||||
|
||||
down_steps = int(np.log2(in_size//min_feat_size))
|
||||
up_steps = int(np.log2(out_size//min_feat_size))
|
||||
|
||||
# =============== define encoder-body-decoder ====================
|
||||
self.encoder = []
|
||||
self.encoder.append(ConvLayer(3, base_ch, 3, 1))
|
||||
head_ch = base_ch
|
||||
for i in range(down_steps):
|
||||
cin, cout = ch_clip(head_ch), ch_clip(head_ch * 2)
|
||||
self.encoder.append(ResidualBlock(cin, cout, scale='down', **act_args))
|
||||
head_ch = head_ch * 2
|
||||
|
||||
self.body = []
|
||||
for i in range(res_depth):
|
||||
self.body.append(ResidualBlock(ch_clip(head_ch), ch_clip(head_ch), **act_args))
|
||||
|
||||
self.decoder = []
|
||||
for i in range(up_steps):
|
||||
cin, cout = ch_clip(head_ch), ch_clip(head_ch // 2)
|
||||
self.decoder.append(ResidualBlock(cin, cout, scale='up', **act_args))
|
||||
head_ch = head_ch // 2
|
||||
|
||||
self.encoder = nn.Sequential(*self.encoder)
|
||||
self.body = nn.Sequential(*self.body)
|
||||
self.decoder = nn.Sequential(*self.decoder)
|
||||
self.out_img_conv = ConvLayer(ch_clip(head_ch), 3)
|
||||
self.out_mask_conv = ConvLayer(ch_clip(head_ch), parsing_ch)
|
||||
|
||||
def forward(self, x):
|
||||
feat = self.encoder(x)
|
||||
x = feat + self.body(feat)
|
||||
x = self.decoder(x)
|
||||
out_img = self.out_img_conv(x)
|
||||
out_mask = self.out_mask_conv(x)
|
||||
return out_mask, out_img
|
||||
|
||||
|
||||
class MultiScaleDiscriminator(nn.Module):
|
||||
def __init__(self, input_ch, base_ch=64, n_layers=3, norm_type='none', relu_type='LeakyReLU', num_D=4):
|
||||
super().__init__()
|
||||
|
||||
self.D_pool = nn.ModuleList()
|
||||
for i in range(num_D):
|
||||
netD = NLayerDiscriminator(input_ch, base_ch, depth=n_layers, norm_type=norm_type, relu_type=relu_type)
|
||||
self.D_pool.append(netD)
|
||||
|
||||
self.downsample = nn.AvgPool2d(3, stride=2, padding=[1, 1], count_include_pad=False)
|
||||
|
||||
def forward(self, input, return_feat=False):
|
||||
results = []
|
||||
for netd in self.D_pool:
|
||||
output = netd(input, return_feat)
|
||||
results.append(output)
|
||||
# Downsample input
|
||||
input = self.downsample(input)
|
||||
return results
|
||||
|
||||
|
||||
class NLayerDiscriminator(nn.Module):
|
||||
def __init__(self,
|
||||
input_ch = 3,
|
||||
base_ch = 64,
|
||||
max_ch = 1024,
|
||||
depth = 4,
|
||||
norm_type = 'none',
|
||||
relu_type = 'LeakyReLU',
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
nargs = {'norm_type': norm_type, 'relu_type': relu_type}
|
||||
self.norm_type = norm_type
|
||||
self.input_ch = input_ch
|
||||
|
||||
self.model = []
|
||||
self.model.append(ConvLayer(input_ch, base_ch, norm_type='none', relu_type=relu_type))
|
||||
for i in range(depth):
|
||||
cin = min(base_ch * 2**(i), max_ch)
|
||||
cout = min(base_ch * 2**(i+1), max_ch)
|
||||
self.model.append(ConvLayer(cin, cout, scale='down_avg', **nargs))
|
||||
self.model = nn.Sequential(*self.model)
|
||||
self.score_out = ConvLayer(cout, 1, use_pad=False)
|
||||
|
||||
def forward(self, x, return_feat=False):
|
||||
ret_feats = []
|
||||
for idx, m in enumerate(self.model):
|
||||
x = m(x)
|
||||
ret_feats.append(x)
|
||||
x = self.score_out(x)
|
||||
if return_feat:
|
||||
return x, ret_feats
|
||||
else:
|
||||
return x
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,80 @@
|
||||
import torch
|
||||
from .base_model import BaseModel
|
||||
from . import networks
|
||||
from utils import utils
|
||||
|
||||
class ParseModel(BaseModel):
|
||||
def modify_commandline_options(parser, is_train):
|
||||
if is_train:
|
||||
parser.add_argument('--parse_map', type=float, default=1.0, help='weight for parsing map')
|
||||
parser.add_argument('--parse_sr', type=float, default=1.0, help='weight for sr')
|
||||
return parser
|
||||
|
||||
def __init__(self, opt):
|
||||
"""Initialize this model class.
|
||||
|
||||
Parameters:
|
||||
opt -- training/test options
|
||||
|
||||
A few things can be done here.
|
||||
- (required) call the initialization function of BaseModel
|
||||
- define loss function, visualization images, model names, and optimizers
|
||||
"""
|
||||
BaseModel.__init__(self, opt) # call the initialization method of BaseModel
|
||||
self.loss_names = ['P', 'SR']
|
||||
self.visual_names = ['img_LR', 'img_HR', 'gt_Parse', 'img_SR', 'pred_Parse']
|
||||
|
||||
self.model_names = ['P']
|
||||
self.netP = networks.define_P(opt)
|
||||
|
||||
if self.isTrain: # only defined during training time
|
||||
self.criterionParse = torch.nn.CrossEntropyLoss()
|
||||
self.criterionSR = torch.nn.L1Loss()
|
||||
self.optimizer = torch.optim.Adam(self.netP.parameters(), lr=opt.lr, betas=(0.9, 0.999))
|
||||
self.optimizers = [self.optimizer]
|
||||
|
||||
def set_input(self, input, cur_iters=None):
|
||||
self.img_LR = input['LR'].to(self.opt.device)
|
||||
self.img_HR = input['HR'].to(self.opt.device)
|
||||
self.gt_Parse = input['Mask'].to(self.opt.device)
|
||||
if self.opt.debug:
|
||||
print('ParseNet input shape:', self.img_LR.shape, self.img_HR.shape, self.gt_Parse.shape)
|
||||
|
||||
def load_pretrain_models(self,):
|
||||
self.netP.eval()
|
||||
print('Loading pretrained LQ face parsing network from', self.opt.parse_net_weight)
|
||||
self.netP.load_state_dict(torch.load(self.opt.parse_net_weight))
|
||||
|
||||
def forward(self):
|
||||
self.pred_Parse, self.img_SR = self.netP(self.img_LR)
|
||||
if self.opt.debug:
|
||||
print('ParseNet output shape', self.pred_Parse.shape, self.img_SR.shape)
|
||||
|
||||
def backward(self):
|
||||
self.loss_P = self.criterionParse(self.pred_Parse, self.gt_Parse) * self.opt.parse_map
|
||||
self.loss_SR = self.criterionSR(self.img_SR, self.img_HR) * self.opt.parse_sr
|
||||
|
||||
loss = self.loss_P + self.loss_SR
|
||||
loss.backward()
|
||||
|
||||
def optimize_parameters(self):
|
||||
self.optimizer.zero_grad() # clear network G's existing gradients
|
||||
self.backward() # calculate gradients for network G
|
||||
self.optimizer.step()
|
||||
|
||||
def get_current_visuals(self, size=512):
|
||||
out = []
|
||||
visual_imgs = []
|
||||
out.append(utils.tensor_to_numpy(self.img_LR))
|
||||
out.append(utils.tensor_to_numpy(self.img_SR))
|
||||
out.append(utils.tensor_to_numpy(self.img_HR))
|
||||
out_imgs = [utils.batch_numpy_to_image(x, size) for x in out]
|
||||
|
||||
visual_imgs.append(out_imgs[0])
|
||||
visual_imgs.append(out_imgs[1])
|
||||
visual_imgs.append(utils.color_parse_map(self.pred_Parse))
|
||||
visual_imgs.append(utils.color_parse_map(self.gt_Parse.unsqueeze(1)))
|
||||
visual_imgs.append(out_imgs[2])
|
||||
|
||||
return visual_imgs
|
||||
|
||||
@@ -0,0 +1,130 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.nn import init
|
||||
import numpy as np
|
||||
from models.blocks import *
|
||||
|
||||
|
||||
class SPADENorm(nn.Module):
|
||||
def __init__(self, norm_nc, ref_nc, norm_type='spade', ksz=3):
|
||||
super().__init__()
|
||||
|
||||
self.param_free_norm = nn.InstanceNorm2d(norm_nc, affine=False)
|
||||
mid_c = 64
|
||||
|
||||
self.norm_type = norm_type
|
||||
if norm_type == 'spade':
|
||||
self.conv1 = nn.Sequential(
|
||||
nn.Conv2d(ref_nc, mid_c, ksz, 1, ksz//2),
|
||||
nn.LeakyReLU(0.2, True),
|
||||
)
|
||||
self.gamma_conv = nn.Conv2d(mid_c, norm_nc, ksz, 1, ksz//2)
|
||||
self.beta_conv = nn.Conv2d(mid_c, norm_nc, ksz, 1, ksz//2)
|
||||
|
||||
def get_gamma_beta(self, x, conv, gamma_conv, beta_conv):
|
||||
act = conv(x)
|
||||
gamma = gamma_conv(act)
|
||||
beta = beta_conv(act)
|
||||
return gamma, beta
|
||||
|
||||
def forward(self, x, ref):
|
||||
normalized_input = self.param_free_norm(x)
|
||||
if x.shape[-1] != ref.shape[-1]:
|
||||
ref = nn.functional.interpolate(ref, x.shape[2:], mode='bicubic', align_corners=False)
|
||||
if self.norm_type == 'spade':
|
||||
gamma, beta = self.get_gamma_beta(ref, self.conv1, self.gamma_conv, self.beta_conv)
|
||||
return normalized_input * gamma + beta
|
||||
elif self.norm_type == 'in':
|
||||
return normalized_input
|
||||
|
||||
|
||||
class SPADEResBlock(nn.Module):
|
||||
def __init__(self, fin, fout, ref_nc, relu_type, norm_type='spade'):
|
||||
super().__init__()
|
||||
|
||||
fmiddle = min(fin, fout)
|
||||
self.conv_0 = nn.Conv2d(fin, fmiddle, kernel_size=3, padding=1)
|
||||
self.conv_1 = nn.Conv2d(fmiddle, fout, kernel_size=3, padding=1)
|
||||
|
||||
# define normalization layers
|
||||
self.norm_0 = SPADENorm(fmiddle, ref_nc, norm_type)
|
||||
self.norm_1 = SPADENorm(fmiddle, ref_nc, norm_type)
|
||||
self.relu = ReluLayer(fmiddle, relu_type)
|
||||
|
||||
def forward(self, x, ref):
|
||||
res = self.conv_0(self.relu(self.norm_0(x, ref)))
|
||||
res = self.conv_1(self.relu(self.norm_1(res, ref)))
|
||||
out = x + res
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class PSFRGenerator(nn.Module):
|
||||
def __init__(self, input_nc, output_nc, in_size=512, out_size=512, min_feat_size=16, ngf=64, n_blocks=9, parse_ch=19, relu_type='relu',
|
||||
ch_range=[32, 1024], norm_type='spade'):
|
||||
super().__init__()
|
||||
|
||||
min_ch, max_ch = ch_range
|
||||
ch_clip = lambda x: max(min_ch, min(x, max_ch))
|
||||
get_ch = lambda size: ch_clip(1024*16//size)
|
||||
|
||||
self.const_input = nn.Parameter(torch.randn(1, get_ch(min_feat_size), min_feat_size, min_feat_size))
|
||||
up_steps = int(np.log2(out_size//min_feat_size))
|
||||
self.up_steps = up_steps
|
||||
|
||||
ref_ch = 19+3
|
||||
|
||||
head_ch = get_ch(min_feat_size)
|
||||
head = [
|
||||
nn.Conv2d(head_ch, head_ch, kernel_size=3, padding=1),
|
||||
SPADEResBlock(head_ch, head_ch, ref_ch, relu_type, norm_type),
|
||||
]
|
||||
|
||||
body = []
|
||||
for i in range(up_steps):
|
||||
cin, cout = ch_clip(head_ch), ch_clip(head_ch // 2)
|
||||
body += [
|
||||
nn.Sequential(
|
||||
nn.Upsample(scale_factor=2),
|
||||
nn.Conv2d(cin, cout, kernel_size=3, padding=1),
|
||||
SPADEResBlock(cout, cout, ref_ch, relu_type, norm_type)
|
||||
)
|
||||
]
|
||||
head_ch = head_ch // 2
|
||||
|
||||
self.img_out = nn.Conv2d(ch_clip(head_ch), output_nc, kernel_size=3, padding=1)
|
||||
|
||||
self.head = nn.Sequential(*head)
|
||||
self.body = nn.Sequential(*body)
|
||||
self.upsample = nn.Upsample(scale_factor=2)
|
||||
|
||||
def forward_spade(self, net, x, ref):
|
||||
for m in net:
|
||||
x = self.forward_spade_m(m, x, ref)
|
||||
return x
|
||||
|
||||
def forward_spade_m(self, m, x, ref):
|
||||
if isinstance(m, SPADENorm) or isinstance(m, SPADEResBlock):
|
||||
x = m(x, ref)
|
||||
else:
|
||||
x = m(x)
|
||||
return x
|
||||
|
||||
def forward(self, x, ref):
|
||||
b, c, h, w = x.shape
|
||||
const_input = self.const_input.repeat(b, 1, 1, 1)
|
||||
ref_input = torch.cat((x, ref), dim=1)
|
||||
|
||||
feat = self.forward_spade(self.head, const_input, ref_input)
|
||||
|
||||
for idx, m in enumerate(self.body):
|
||||
feat = self.forward_spade(m, feat, ref_input)
|
||||
|
||||
out_img = self.img_out(feat)
|
||||
|
||||
return out_img
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
x = torch.randn(2, 16, 567, 234)
|
||||
nearest_interpolate(x)
|
||||
@@ -0,0 +1 @@
|
||||
"""This package options includes option modules: training options, test options, and basic options (used in both training and test)."""
|
||||
@@ -0,0 +1,165 @@
|
||||
import argparse
|
||||
import os
|
||||
import numpy as np
|
||||
import random
|
||||
from utils import utils
|
||||
import torch
|
||||
import models
|
||||
import data
|
||||
from utils import utils
|
||||
|
||||
|
||||
class BaseOptions():
|
||||
"""This class defines options used during both training and test time.
|
||||
|
||||
It also implements several helper functions such as parsing, printing, and saving the options.
|
||||
It also gathers additional options defined in <modify_commandline_options> functions in both dataset class and model class.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""Reset the class; indicates the class hasn't been initailized"""
|
||||
self.initialized = False
|
||||
|
||||
def initialize(self, parser):
|
||||
"""Define the common options that are used in both training and test."""
|
||||
# basic parameters
|
||||
parser.add_argument('--dataroot', required=False, help='path to images')
|
||||
parser.add_argument('--name', type=str, default='experiment_name', help='name of the experiment. It decides where to store samples and models')
|
||||
parser.add_argument('--gpus', type=int, default=1, help='how many gpus to use')
|
||||
parser.add_argument('--seed', type=int, default=123, help='Random seed for training')
|
||||
parser.add_argument('--checkpoints_dir', type=str, default='./check_points', help='models are saved here')
|
||||
# model parameters
|
||||
parser.add_argument('--model', type=str, default='enhance', help='chooses which model to train [parse|enhance]')
|
||||
parser.add_argument('--input_nc', type=int, default=3, help='# of input image channels: 3 for RGB and 1 for grayscale')
|
||||
parser.add_argument('--Dinput_nc', type=int, default=3, help='# of input image channels: 3 for RGB and 1 for grayscale')
|
||||
parser.add_argument('--output_nc', type=int, default=3, help='# of output image channels: 3 for RGB and 1 for grayscale')
|
||||
parser.add_argument('--ngf', type=int, default=64, help='# of gen filters in the last conv layer')
|
||||
parser.add_argument('--ndf', type=int, default=64, help='# of discrim filters in the first conv layer')
|
||||
parser.add_argument('--n_layers_D', type=int, default=4, help='downsampling layers in discriminator')
|
||||
parser.add_argument('--D_num', type=int, default=3, help='numbers of discriminators')
|
||||
|
||||
parser.add_argument('--Pnorm', type=str, default='bn', help='parsing net norm [in | bn| none]')
|
||||
parser.add_argument('--Gnorm', type=str, default='spade', help='generator norm [in | bn | none]')
|
||||
parser.add_argument('--Dnorm', type=str, default='in', help='discriminator norm [in | bn | none]')
|
||||
parser.add_argument('--init_type', type=str, default='normal', help='network initialization [normal | xavier | kaiming | orthogonal]')
|
||||
parser.add_argument('--init_gain', type=float, default=0.02, help='scaling factor for normal, xavier and orthogonal.')
|
||||
# dataset parameters
|
||||
parser.add_argument('--dataset_name', type=str, default='single', help='dataset name')
|
||||
parser.add_argument('--Pimg_size', type=int, default='512', help='image size for face parse net')
|
||||
parser.add_argument('--Gin_size', type=int, default='512', help='image size for face parse net')
|
||||
parser.add_argument('--Gout_size', type=int, default='512', help='image size for face parse net')
|
||||
parser.add_argument('--serial_batches', action='store_true', help='if true, takes images in order to make batches, otherwise takes them randomly')
|
||||
parser.add_argument('--num_threads', default=8, type=int, help='# threads for loading data')
|
||||
parser.add_argument('--batch_size', type=int, default=16, help='input batch size')
|
||||
parser.add_argument('--load_size', type=int, default=512, help='scale images to this size')
|
||||
parser.add_argument('--crop_size', type=int, default=256, help='then crop to this size')
|
||||
parser.add_argument('--max_dataset_size', type=int, default=float("inf"), help='Maximum number of samples allowed per dataset. If the dataset directory contains more than max_dataset_size, only a subset is loaded.')
|
||||
parser.add_argument('--preprocess', type=str, default='none', help='scaling and cropping of images at load time [resize_and_crop | crop | scale_width | scale_width_and_crop | none]')
|
||||
parser.add_argument('--no_flip', action='store_true', help='if specified, do not flip the images for data augmentation')
|
||||
# additional parameters
|
||||
parser.add_argument('--epoch', type=str, default='latest', help='which epoch to load? set to latest to use latest cached model')
|
||||
parser.add_argument('--load_iter', type=int, default='0', help='which iteration to load? if load_iter > 0, the code will load models by iter_[load_iter]; otherwise, the code will load models by [epoch]')
|
||||
parser.add_argument('--verbose', action='store_true', help='if specified, print more debugging information')
|
||||
parser.add_argument('--suffix', default='', type=str, help='customized suffix: opt.name = opt.name + suffix: e.g., {model}_{netG}_size{load_size}')
|
||||
|
||||
parser.add_argument('--debug', action='store_true', help='if specified, set to debug mode')
|
||||
self.initialized = True
|
||||
return parser
|
||||
|
||||
def gather_options(self):
|
||||
"""Initialize our parser with basic options(only once).
|
||||
Add additional model-specific and dataset-specific options.
|
||||
These options are defined in the <modify_commandline_options> function
|
||||
in model and dataset classes.
|
||||
"""
|
||||
if not self.initialized: # check if it has been initialized
|
||||
parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
||||
parser = self.initialize(parser)
|
||||
|
||||
# get the basic options
|
||||
opt, _ = parser.parse_known_args()
|
||||
|
||||
# modify model-related parser options
|
||||
model_name = opt.model
|
||||
model_option_setter = models.get_option_setter(model_name)
|
||||
parser = model_option_setter(parser, self.isTrain)
|
||||
opt, _ = parser.parse_known_args() # parse again with new defaults
|
||||
|
||||
# modify dataset-related parser options
|
||||
dataset_name = opt.dataset_name
|
||||
dataset_option_setter = data.get_option_setter(dataset_name)
|
||||
parser = dataset_option_setter(parser, self.isTrain)
|
||||
|
||||
# save and return the parser
|
||||
self.parser = parser
|
||||
return parser.parse_args()
|
||||
|
||||
def print_options(self, opt):
|
||||
"""Print and save options
|
||||
|
||||
It will print both current options and default values(if different).
|
||||
It will save options into a text file / [checkpoints_dir] / opt.txt
|
||||
"""
|
||||
message = ''
|
||||
message += '----------------- Options ---------------\n'
|
||||
for k, v in sorted(vars(opt).items()):
|
||||
comment = ''
|
||||
default = self.parser.get_default(k)
|
||||
if v != default:
|
||||
comment = '\t[default: %s]' % str(default)
|
||||
message += '{:>25}: {:<30}{}\n'.format(str(k), str(v), comment)
|
||||
message += '----------------- End -------------------'
|
||||
print(message)
|
||||
|
||||
# save to the disk
|
||||
opt.expr_dir = os.path.join(opt.checkpoints_dir, opt.name)
|
||||
utils.mkdirs(opt.expr_dir)
|
||||
file_name = os.path.join(opt.expr_dir, '{}_opt.txt'.format(opt.phase))
|
||||
with open(file_name, 'wt') as opt_file:
|
||||
opt_file.write(message)
|
||||
opt_file.write('\n')
|
||||
|
||||
opt.log_dir = os.path.join(opt.checkpoints_dir, 'log_dir')
|
||||
utils.mkdirs(opt.log_dir)
|
||||
opt.log_archive = os.path.join(opt.checkpoints_dir, 'log_archive')
|
||||
utils.mkdirs(opt.log_archive)
|
||||
|
||||
def parse(self):
|
||||
"""Parse our options, create checkpoints directory suffix, and set up gpu device."""
|
||||
opt = self.gather_options()
|
||||
opt.isTrain = self.isTrain # train or test
|
||||
|
||||
if opt.debug:
|
||||
opt.name = 'debug'
|
||||
opt.save_iter_freq = 1
|
||||
opt.save_latest_freq = 1
|
||||
opt.visual_freq = 1
|
||||
opt.print_freq = 1
|
||||
|
||||
# Find avaliable GPUs automatically
|
||||
if opt.gpus > 0:
|
||||
opt.gpu_ids = utils.get_gpu_memory_map()[1][:opt.gpus]
|
||||
if not isinstance(opt.gpu_ids, list):
|
||||
opt.gpu_ids = [opt.gpu_ids]
|
||||
torch.cuda.set_device(opt.gpu_ids[0])
|
||||
opt.device = torch.device('cuda:{}'.format(opt.gpu_ids[0 % opt.gpus]))
|
||||
opt.data_device = torch.device('cuda:{}'.format(opt.gpu_ids[1 % opt.gpus]))
|
||||
else:
|
||||
opt.gpu_ids = []
|
||||
opt.device = torch.device('cpu')
|
||||
|
||||
# set random seeds to ensure reproducibility
|
||||
np.random.seed(opt.seed)
|
||||
random.seed(opt.seed)
|
||||
torch.manual_seed(opt.seed)
|
||||
torch.cuda.manual_seed_all(opt.seed)
|
||||
|
||||
# process opt.suffix
|
||||
if opt.suffix:
|
||||
suffix = ('_' + opt.suffix.format(**vars(opt))) if opt.suffix != '' else ''
|
||||
opt.name = opt.name + suffix
|
||||
|
||||
self.print_options(opt)
|
||||
|
||||
self.opt = opt
|
||||
return self.opt
|
||||
@@ -0,0 +1,30 @@
|
||||
from .base_options import BaseOptions
|
||||
|
||||
|
||||
class TestOptions(BaseOptions):
|
||||
"""This class includes test options.
|
||||
|
||||
It also includes shared options defined in BaseOptions.
|
||||
"""
|
||||
|
||||
def initialize(self, parser):
|
||||
parser = BaseOptions.initialize(self, parser) # define shared options
|
||||
parser.add_argument('--src_dir', type=str, default='G:/VGGFace2-HQ/VGGface2_None_norm_512_true_bygfpgan/n000002', help='source directory containing test images')
|
||||
parser.add_argument('--save_masks_dir', type=str, default='../datasets/FFHQ/masks512', help='path to save parsing masks for FFHQ')
|
||||
parser.add_argument('--test_img_path', type=str, default='', help='path for single image test')
|
||||
parser.add_argument('--test_upscale', type=float, default=1, help='upsample scale for single image test')
|
||||
parser.add_argument('--ntest', type=int, default=float("inf"), help='# of test examples.')
|
||||
parser.add_argument('--results_dir', type=str, default='./results/', help='saves results here.')
|
||||
parser.add_argument('--aspect_ratio', type=float, default=1.0, help='aspect ratio of result images')
|
||||
parser.add_argument('--phase', type=str, default='test', help='train, val, test, etc')
|
||||
# Dropout and Batchnorm has different behavioir during training and test.
|
||||
parser.add_argument('--eval', action='store_true', help='use eval mode during test time.')
|
||||
parser.add_argument('--num_test', type=int, default=50, help='how many test images to run')
|
||||
parser.add_argument('--parse_net_weight', type=str, default='./pretrain_models/parse_multi_iter_90000.pth', help='parse model path')
|
||||
parser.add_argument('--psfr_net_weight', type=str, default='./pretrain_models/psfrgan_epoch15_net_G.pth', help='parse model path')
|
||||
# rewrite devalue values
|
||||
# To avoid cropping, the load_size should be the same as crop_size
|
||||
parser.set_defaults(load_size=parser.get_default('crop_size'))
|
||||
self.isTrain = False
|
||||
|
||||
return parser
|
||||
@@ -0,0 +1,41 @@
|
||||
from .base_options import BaseOptions
|
||||
|
||||
class TrainOptions(BaseOptions):
|
||||
"""This class includes training options.
|
||||
|
||||
It also includes shared options defined in BaseOptions.
|
||||
"""
|
||||
|
||||
def initialize(self, parser):
|
||||
parser = BaseOptions.initialize(self, parser)
|
||||
# visdom and HTML visualization parameters
|
||||
parser.add_argument('--visual_freq', type=int, default=400, help='frequency of show training images in tensorboard')
|
||||
parser.add_argument('--print_freq', type=int, default=100, help='frequency of showing training results on console')
|
||||
# network saving and loading parameters
|
||||
parser.add_argument('--save_iter_freq', type=int, default=5000, help='frequency of saving the models')
|
||||
parser.add_argument('--save_latest_freq', type=int, default=500, help='save latest freq')
|
||||
parser.add_argument('--save_epoch_freq', type=int, default=5, help='frequency of saving checkpoints at the end of epochs')
|
||||
parser.add_argument('--save_by_iter', action='store_true', help='whether saves model by iteration')
|
||||
parser.add_argument('--continue_train', action='store_true', help='continue training: load the latest model')
|
||||
parser.add_argument('--no_strict_load', action='store_true', help='set strict load to false')
|
||||
parser.add_argument('--epoch_count', type=int, default=1, help='the starting epoch count, we save the model by <epoch_count>, <epoch_count>+<save_latest_freq>, ...')
|
||||
parser.add_argument('--phase', type=str, default='train', help='train, val, test, etc')
|
||||
# training parameters
|
||||
parser.add_argument('--resume_epoch', type=int, default=0, help='training resume epoch')
|
||||
parser.add_argument('--resume_iter', type=int, default=0, help='training resume iter')
|
||||
parser.add_argument('--total_epochs', type=int, default=50, help='# of epochs to train')
|
||||
parser.add_argument('--n_epochs', type=int, default=100, help='number of epochs with the initial learning rate')
|
||||
parser.add_argument('--n_epochs_decay', type=int, default=100, help='number of epochs to linearly decay learning rate to zero')
|
||||
parser.add_argument('--niter_decay', type=int, default=100, help='# of iter to linearly decay learning rate to zero')
|
||||
parser.add_argument('--beta1', type=float, default=0.5, help='momentum term of adam')
|
||||
parser.add_argument('--lr', type=float, default=0.0002, help='initial learning rate for adam')
|
||||
parser.add_argument('--g_lr', type=float, default=0.0001, help='generator learning rate')
|
||||
parser.add_argument('--d_lr', type=float, default=0.0004, help='discriminator learning rate')
|
||||
parser.add_argument('--gan_mode', type=str, default='hinge', help='the type of GAN objective. [vanilla| lsgan | wgangp]. vanilla GAN loss is the cross-entropy objective used in the original GAN paper.')
|
||||
parser.add_argument('--lr_policy', type=str, default='step', help='learning rate policy. [linear | step | plateau | cosine]')
|
||||
parser.add_argument('--lr_decay_iters', type=int, default=50, help='multiply by a gamma every lr_decay_iters iterations')
|
||||
parser.add_argument('--lr_decay_gamma', type=float, default=1, help='multiply by a gamma every lr_decay_iters iterations')
|
||||
|
||||
self.isTrain = True
|
||||
|
||||
return parser
|
||||
@@ -0,0 +1,11 @@
|
||||
torch==1.5.1
|
||||
torchvision==0.6.1
|
||||
tensorflow>=1.15.4
|
||||
tensorboard==1.15.0
|
||||
tensorboardX==2.1
|
||||
opencv-python
|
||||
dlib
|
||||
scikit-image==0.17.2
|
||||
scipy==1.4.1
|
||||
tqdm
|
||||
imgaug
|
||||
@@ -0,0 +1,62 @@
|
||||
import os
|
||||
from options.test_options import TestOptions
|
||||
from data import create_dataset
|
||||
from models import create_model
|
||||
from utils import utils
|
||||
from PIL import Image
|
||||
from tqdm import tqdm
|
||||
import torch
|
||||
import time
|
||||
import numpy as np
|
||||
|
||||
if __name__ == '__main__':
|
||||
opt = TestOptions().parse() # get test options
|
||||
opt.num_threads = 0 # test code only supports num_threads = 1
|
||||
opt.batch_size = 4 # test code only supports batch_size = 1
|
||||
opt.serial_batches = True # disable data shuffling; comment this line if results on randomly chosen images are needed.
|
||||
opt.no_flip = True
|
||||
|
||||
dataset = create_dataset(opt) # create a dataset given opt.dataset_mode and other options
|
||||
model = create_model(opt) # create a model given opt.model and other options
|
||||
model.load_pretrain_models()
|
||||
|
||||
save_dir = opt.results_dir
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
|
||||
print('creating result directory', save_dir)
|
||||
netP = model.netP
|
||||
netG = model.netG
|
||||
model.eval()
|
||||
max_size = 9999
|
||||
os.makedirs(os.path.join(save_dir, 'sr'), exist_ok=True)
|
||||
for i, data in tqdm(enumerate(dataset), total=len(dataset)//opt.batch_size):
|
||||
inp = data['LR']
|
||||
with torch.no_grad():
|
||||
parse_map, _ = netP(inp)
|
||||
parse_map_sm = (parse_map == parse_map.max(dim=1, keepdim=True)[0]).float()
|
||||
output_SR = netG(inp, parse_map_sm)
|
||||
img_path = data['LR_paths'] # get image paths
|
||||
for i in tqdm(range(len(img_path))):
|
||||
inp_img = utils.batch_tensor_to_img(inp)
|
||||
output_sr_img = utils.batch_tensor_to_img(output_SR)
|
||||
ref_parse_img = utils.color_parse_map(parse_map_sm)
|
||||
|
||||
save_path = os.path.join(save_dir, 'lq', os.path.basename(img_path[i]))
|
||||
os.makedirs(os.path.join(save_dir, 'lq'), exist_ok=True)
|
||||
save_img = Image.fromarray(inp_img[i])
|
||||
save_img.save(save_path)
|
||||
|
||||
save_path = os.path.join(save_dir, 'hq', os.path.basename(img_path[i]))
|
||||
os.makedirs(os.path.join(save_dir, 'hq'), exist_ok=True)
|
||||
save_img = Image.fromarray(output_sr_img[i])
|
||||
save_img.save(save_path)
|
||||
|
||||
save_path = os.path.join(save_dir, 'parse', os.path.basename(img_path[i]))
|
||||
os.makedirs(os.path.join(save_dir, 'parse'), exist_ok=True)
|
||||
save_img = Image.fromarray(ref_parse_img[i])
|
||||
save_img.save(save_path)
|
||||
|
||||
if i > max_size: break
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,57 @@
|
||||
'''
|
||||
This script enhance images with unaligned faces in a folder and paste it back to the original place.
|
||||
'''
|
||||
import dlib
|
||||
import os
|
||||
import cv2
|
||||
import numpy as np
|
||||
from tqdm import tqdm
|
||||
from skimage import transform as trans
|
||||
from skimage import io
|
||||
|
||||
import torch
|
||||
from utils import utils
|
||||
from options.test_options import TestOptions
|
||||
from models import create_model
|
||||
|
||||
from test_enhance_single_unalign import *
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
opt = TestOptions().parse()
|
||||
# face_detector = dlib.get_frontal_face_detector()
|
||||
face_detector = dlib.cnn_face_detection_model_v1('./pretrain_models/mmod_human_face_detector.dat')
|
||||
lmk_predictor = dlib.shape_predictor('./pretrain_models/shape_predictor_5_face_landmarks.dat')
|
||||
template_path = './pretrain_models/FFHQ_template.npy'
|
||||
enhance_model = def_models(opt)
|
||||
|
||||
for img_name in os.listdir(opt.src_dir):
|
||||
img_path = os.path.join(opt.src_dir, img_name)
|
||||
save_current_dir = os.path.join(opt.results_dir, os.path.splitext(img_name)[0])
|
||||
os.makedirs(save_current_dir, exist_ok=True)
|
||||
print('======> Loading image', img_path)
|
||||
img = dlib.load_rgb_image(img_path)
|
||||
aligned_faces, tform_params = detect_and_align_faces(img, face_detector, lmk_predictor, template_path)
|
||||
# Save aligned LQ faces
|
||||
save_lq_dir = os.path.join(save_current_dir, 'LQ_faces')
|
||||
os.makedirs(save_lq_dir, exist_ok=True)
|
||||
print('======> Saving aligned LQ faces to', save_lq_dir)
|
||||
save_imgs(aligned_faces, save_lq_dir)
|
||||
|
||||
hq_faces, lq_parse_maps = enhance_faces(aligned_faces, enhance_model)
|
||||
# Save LQ parsing maps and enhanced faces
|
||||
save_parse_dir = os.path.join(save_current_dir, 'ParseMaps')
|
||||
save_hq_dir = os.path.join(save_current_dir, 'HQ')
|
||||
os.makedirs(save_parse_dir, exist_ok=True)
|
||||
os.makedirs(save_hq_dir, exist_ok=True)
|
||||
print('======> Save parsing map and the enhanced faces.')
|
||||
save_imgs(lq_parse_maps, save_parse_dir)
|
||||
save_imgs(hq_faces, save_hq_dir)
|
||||
|
||||
print('======> Paste the enhanced faces back to the original image.')
|
||||
hq_img = past_faces_back(img, hq_faces, tform_params, upscale=opt.test_upscale)
|
||||
final_save_path = os.path.join(save_current_dir, 'hq_final.jpg')
|
||||
print('======> Save final result to', final_save_path)
|
||||
io.imsave(final_save_path, hq_img)
|
||||
|
||||
|
||||
@@ -0,0 +1,126 @@
|
||||
'''
|
||||
This script enhance all faces in one image with PSFR-GAN and paste it back to the original place.
|
||||
'''
|
||||
import dlib
|
||||
import os
|
||||
import cv2
|
||||
import numpy as np
|
||||
from tqdm import tqdm
|
||||
from skimage import transform as trans
|
||||
from skimage import io
|
||||
|
||||
import torch
|
||||
from utils import utils
|
||||
from options.test_options import TestOptions
|
||||
from models import create_model
|
||||
|
||||
|
||||
def detect_and_align_faces(img, face_detector, lmk_predictor, template_path, template_scale=2, size_threshold=999):
|
||||
align_out_size = (512, 512)
|
||||
ref_points = np.load(template_path) / template_scale
|
||||
|
||||
# Detect landmark points
|
||||
face_dets = face_detector(img, 1)
|
||||
assert len(face_dets) > 0, 'No faces detected'
|
||||
|
||||
aligned_faces = []
|
||||
tform_params = []
|
||||
for det in face_dets:
|
||||
if isinstance(face_detector, dlib.cnn_face_detection_model_v1):
|
||||
rec = det.rect # for cnn detector
|
||||
else:
|
||||
rec = det
|
||||
if rec.width() > size_threshold or rec.height() > size_threshold:
|
||||
print('Face is too large')
|
||||
break
|
||||
landmark_points = lmk_predictor(img, rec)
|
||||
single_points = []
|
||||
for i in range(5):
|
||||
single_points.append([landmark_points.part(i).x, landmark_points.part(i).y])
|
||||
single_points = np.array(single_points)
|
||||
tform = trans.SimilarityTransform()
|
||||
tform.estimate(single_points, ref_points)
|
||||
tmp_face = trans.warp(img, tform.inverse, output_shape=align_out_size, order=3)
|
||||
aligned_faces.append(tmp_face*255)
|
||||
tform_params.append(tform)
|
||||
return [aligned_faces, tform_params]
|
||||
|
||||
|
||||
def def_models(opt):
|
||||
model = create_model(opt)
|
||||
model.load_pretrain_models()
|
||||
model.netP.to(opt.device)
|
||||
model.netG.to(opt.device)
|
||||
return model
|
||||
|
||||
|
||||
def enhance_faces(LQ_faces, model):
|
||||
hq_faces = []
|
||||
lq_parse_maps = []
|
||||
for lq_face in tqdm(LQ_faces):
|
||||
with torch.no_grad():
|
||||
lq_tensor = torch.tensor(lq_face.transpose(2, 0, 1)) / 255. * 2 - 1
|
||||
lq_tensor = lq_tensor.unsqueeze(0).float().to(model.device)
|
||||
parse_map, _ = model.netP(lq_tensor)
|
||||
parse_map_onehot = (parse_map == parse_map.max(dim=1, keepdim=True)[0]).float()
|
||||
output_SR = model.netG(lq_tensor, parse_map_onehot)
|
||||
hq_faces.append(utils.tensor_to_img(output_SR))
|
||||
lq_parse_maps.append(utils.color_parse_map(parse_map_onehot)[0])
|
||||
return hq_faces, lq_parse_maps
|
||||
|
||||
|
||||
def past_faces_back(img, hq_faces, tform_params, upscale=1):
|
||||
h, w = img.shape[:2]
|
||||
img = cv2.resize(img, (int(w*upscale), int(h*upscale)), interpolation=cv2.INTER_CUBIC)
|
||||
for hq_img, tform in tqdm(zip(hq_faces, tform_params), total=len(hq_faces)):
|
||||
tform.params[0:2,0:2] /= upscale
|
||||
back_img = trans.warp(hq_img/255., tform, output_shape=[int(h*upscale), int(w*upscale)], order=3) * 255
|
||||
|
||||
# blur mask to avoid border artifacts
|
||||
mask = (back_img == 0)
|
||||
mask = cv2.blur(mask.astype(np.float32), (5,5))
|
||||
mask = (mask > 0)
|
||||
img = img * mask + (1 - mask) * back_img
|
||||
return img.astype(np.uint8)
|
||||
|
||||
|
||||
def save_imgs(img_list, save_dir):
|
||||
for idx, img in enumerate(img_list):
|
||||
save_path = os.path.join(save_dir, '{:03d}.jpg'.format(idx))
|
||||
io.imsave(save_path, img.astype(np.uint8))
|
||||
|
||||
if __name__ == '__main__':
|
||||
opt = TestOptions().parse()
|
||||
# face_detector = dlib.get_frontal_face_detector()
|
||||
face_detector = dlib.cnn_face_detection_model_v1('./pretrain_models/mmod_human_face_detector.dat')
|
||||
lmk_predictor = dlib.shape_predictor('./pretrain_models/shape_predictor_5_face_landmarks.dat')
|
||||
template_path = './pretrain_models/FFHQ_template.npy'
|
||||
|
||||
print('======> Loading images, crop and align faces.')
|
||||
img_path = opt.test_img_path
|
||||
img = dlib.load_rgb_image(img_path)
|
||||
aligned_faces, tform_params = detect_and_align_faces(img, face_detector, lmk_predictor, template_path)
|
||||
# Save aligned LQ faces
|
||||
save_lq_dir = os.path.join(opt.results_dir, 'LQ_faces')
|
||||
os.makedirs(save_lq_dir, exist_ok=True)
|
||||
print('======> Saving aligned LQ faces to', save_lq_dir)
|
||||
save_imgs(aligned_faces, save_lq_dir)
|
||||
|
||||
enhance_model = def_models(opt)
|
||||
hq_faces, lq_parse_maps = enhance_faces(aligned_faces, enhance_model)
|
||||
# Save LQ parsing maps and enhanced faces
|
||||
save_parse_dir = os.path.join(opt.results_dir, 'ParseMaps')
|
||||
save_hq_dir = os.path.join(opt.results_dir, 'HQ')
|
||||
os.makedirs(save_parse_dir, exist_ok=True)
|
||||
os.makedirs(save_hq_dir, exist_ok=True)
|
||||
print('======> Save parsing map and the enhanced faces.')
|
||||
save_imgs(lq_parse_maps, save_parse_dir)
|
||||
save_imgs(hq_faces, save_hq_dir)
|
||||
|
||||
print('======> Paste the enhanced faces back to the original image.')
|
||||
hq_img = past_faces_back(img, hq_faces, tform_params, upscale=opt.test_upscale)
|
||||
final_save_path = os.path.join(opt.results_dir, 'hq_final.jpg')
|
||||
print('======> Save final result to', final_save_path)
|
||||
io.imsave(final_save_path, hq_img)
|
||||
|
||||
|
||||
@@ -0,0 +1,79 @@
|
||||
from utils.timer import Timer
|
||||
from utils.logger import Logger
|
||||
from utils import utils
|
||||
|
||||
from options.train_options import TrainOptions
|
||||
from data import create_dataset
|
||||
from models import create_model
|
||||
|
||||
import torch
|
||||
import os
|
||||
import torch.multiprocessing as mp
|
||||
|
||||
def train(opt):
|
||||
|
||||
dataset = create_dataset(opt) # create a dataset given opt.dataset_mode and other options
|
||||
dataset_size = len(dataset) # get the number of images in the dataset.
|
||||
print('The number of training images = %d' % dataset_size)
|
||||
|
||||
model = create_model(opt)
|
||||
model.setup(opt)
|
||||
|
||||
logger = Logger(opt)
|
||||
timer = Timer()
|
||||
|
||||
single_epoch_iters = (dataset_size // opt.batch_size)
|
||||
total_iters = opt.total_epochs * single_epoch_iters
|
||||
cur_iters = opt.resume_iter + opt.resume_epoch * single_epoch_iters
|
||||
start_iter = opt.resume_iter
|
||||
print('Start training from epoch: {:05d}; iter: {:07d}'.format(opt.resume_epoch, opt.resume_iter))
|
||||
for epoch in range(opt.resume_epoch, opt.total_epochs + 1):
|
||||
for i, data in enumerate(dataset, start=start_iter):
|
||||
cur_iters += 1
|
||||
logger.set_current_iter(cur_iters)
|
||||
# =================== load data ===============
|
||||
model.set_input(data, cur_iters)
|
||||
timer.update_time('DataTime')
|
||||
|
||||
# =================== model train ===============
|
||||
model.forward(), timer.update_time('Forward')
|
||||
model.optimize_parameters()
|
||||
loss = model.get_current_losses()
|
||||
loss.update(model.get_lr())
|
||||
logger.record_losses(loss)
|
||||
timer.update_time('Backward')
|
||||
|
||||
# =================== save model and visualize ===============
|
||||
if cur_iters % opt.print_freq == 0:
|
||||
print('Model log directory: {}'.format(opt.expr_dir))
|
||||
epoch_progress = '{:03d}|{:05d}/{:05d}'.format(epoch, i, single_epoch_iters)
|
||||
logger.printIterSummary(epoch_progress, cur_iters, total_iters, timer)
|
||||
|
||||
if cur_iters % opt.visual_freq == 0:
|
||||
visual_imgs = model.get_current_visuals()
|
||||
logger.record_images(visual_imgs)
|
||||
|
||||
if cur_iters % opt.save_iter_freq == 0:
|
||||
print('saving current model (epoch %d, iters %d)' % (epoch, cur_iters))
|
||||
save_suffix = 'iter_%d' % cur_iters
|
||||
info = {'resume_epoch': epoch, 'resume_iter': i+1}
|
||||
model.save_networks(save_suffix, info)
|
||||
|
||||
if cur_iters % opt.save_latest_freq == 0:
|
||||
print('saving the latest model (epoch %d, iters %d)' % (epoch, cur_iters))
|
||||
info = {'resume_epoch': epoch, 'resume_iter': i+1}
|
||||
model.save_networks('latest', info)
|
||||
|
||||
if i >= single_epoch_iters - 1:
|
||||
start_iter = 0
|
||||
break
|
||||
|
||||
# model.update_learning_rate()
|
||||
if opt.debug: break
|
||||
if opt.debug and epoch >= 0: break
|
||||
logger.close()
|
||||
|
||||
if __name__ == '__main__':
|
||||
opt = TrainOptions().parse()
|
||||
train(opt)
|
||||
|
||||
@@ -0,0 +1,91 @@
|
||||
import os
|
||||
from collections import OrderedDict
|
||||
import numpy as np
|
||||
from .utils import mkdirs
|
||||
from tensorboardX import SummaryWriter
|
||||
from datetime import datetime
|
||||
import socket
|
||||
import shutil
|
||||
|
||||
class Logger():
|
||||
def __init__(self, opts):
|
||||
time_stamp = '_{}'.format(datetime.now().strftime('%Y-%m-%d_%H:%M'))
|
||||
self.opts = opts
|
||||
self.log_dir = os.path.join(opts.log_dir, opts.name+time_stamp)
|
||||
self.phase_keys = ['train', 'val', 'test']
|
||||
self.iter_log = []
|
||||
self.epoch_log = OrderedDict()
|
||||
self.set_mode(opts.phase)
|
||||
|
||||
# check if exist previous log belong to the same experiment name
|
||||
exist_log = None
|
||||
for log_name in os.listdir(opts.log_dir):
|
||||
if opts.name in log_name:
|
||||
exist_log = log_name
|
||||
if exist_log is not None:
|
||||
old_dir = os.path.join(opts.log_dir, exist_log)
|
||||
archive_dir = os.path.join(opts.log_archive, exist_log)
|
||||
shutil.move(old_dir, archive_dir)
|
||||
|
||||
self.mk_log_file()
|
||||
|
||||
self.writer = SummaryWriter(self.log_dir)
|
||||
|
||||
def mk_log_file(self):
|
||||
mkdirs(self.log_dir)
|
||||
self.txt_files = OrderedDict()
|
||||
for i in self.phase_keys:
|
||||
self.txt_files[i] = os.path.join(self.log_dir, 'log_{}'.format(i))
|
||||
|
||||
def set_mode(self, mode):
|
||||
self.mode = mode
|
||||
self.epoch_log[mode] = []
|
||||
|
||||
def set_current_iter(self, cur_iter):
|
||||
self.cur_iter = cur_iter
|
||||
|
||||
def record_losses(self, items):
|
||||
"""
|
||||
iteration log: [iter][{key: value}]
|
||||
"""
|
||||
self.iter_log.append(items)
|
||||
for k, v in items.items():
|
||||
if 'loss' in k.lower():
|
||||
self.writer.add_scalar('loss/{}'.format(k), v, self.cur_iter)
|
||||
|
||||
def record_scalar(self, items):
|
||||
"""
|
||||
Add scalar records. item, {key: value}
|
||||
"""
|
||||
for i in items.keys():
|
||||
self.writer.add_scalar('{}'.format(i), items[i], self.cur_iter)
|
||||
|
||||
def record_images(self, visuals, nrow=6, tag='ckpt_image'):
|
||||
imgs = []
|
||||
max_len = visuals[0].shape[0]
|
||||
for i in range(nrow):
|
||||
if i >= max_len: continue
|
||||
tmp_imgs = [x[i] for x in visuals]
|
||||
imgs.append(np.hstack(tmp_imgs))
|
||||
imgs = np.vstack(imgs).astype(np.uint8)
|
||||
self.writer.add_image(tag, imgs, self.cur_iter, dataformats='HWC')
|
||||
|
||||
def record_text(self, tag, text):
|
||||
self.writer.add_text(tag, text)
|
||||
|
||||
def printIterSummary(self, epoch, cur_iters, total_it, timer):
|
||||
msg = '{}\nIter: [{}]{:03d}/{:03d}\t\t'.format(
|
||||
timer.to_string(total_it - cur_iters), epoch, cur_iters, total_it)
|
||||
for k, v in self.iter_log[-1].items():
|
||||
msg += '{}: {:.6f}\t'.format(k, v)
|
||||
print(msg + '\n')
|
||||
with open(self.txt_files[self.mode], 'a+') as f:
|
||||
f.write(msg + '\n')
|
||||
|
||||
def close(self):
|
||||
self.writer.export_scalars_to_json(os.path.join(self.log_dir, 'all_scalars.json'))
|
||||
self.writer.close()
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,34 @@
|
||||
import time
|
||||
import datetime
|
||||
from collections import OrderedDict
|
||||
|
||||
class Timer():
|
||||
def __init__(self):
|
||||
self.reset_timer()
|
||||
self.start = time.time()
|
||||
|
||||
def reset_timer(self):
|
||||
self.before = time.time()
|
||||
self.timer = OrderedDict()
|
||||
|
||||
def restart(self):
|
||||
self.before = time.time()
|
||||
|
||||
def update_time(self, key):
|
||||
self.timer[key] = time.time() - self.before
|
||||
self.before = time.time()
|
||||
|
||||
def to_string(self, iters_left, short=False):
|
||||
iter_total = sum(self.timer.values())
|
||||
msg = "{:%Y-%m-%d %H:%M:%S}\tElapse: {}\tTimeLeft: {}\t".format(
|
||||
datetime.datetime.now(),
|
||||
datetime.timedelta(seconds=round(time.time() - self.start)),
|
||||
datetime.timedelta(seconds=round(iter_total*iters_left))
|
||||
)
|
||||
if short:
|
||||
msg += '{}: {:.2f}s'.format('|'.join(self.timer.keys()), iter_total)
|
||||
else:
|
||||
msg += '\tIterTotal: {:.2f}s\t{}: {} '.format(iter_total,
|
||||
'|'.join(self.timer.keys()), ' '.join('{:.2f}s'.format(x) for x in self.timer.values()))
|
||||
return msg
|
||||
|
||||
@@ -0,0 +1,169 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
import cv2 as cv
|
||||
from skimage import io
|
||||
from PIL import Image
|
||||
import os
|
||||
import subprocess
|
||||
|
||||
# MASK_COLORMAP = [[0, 0, 0], [204, 0, 0], [76, 153, 0], [204, 204, 0], [51, 51, 255], [204, 0, 204], [0, 255, 255], [255, 204, 204], [102, 51, 0], [255, 0, 0], [102, 204, 0], [255, 255, 0], [0, 0, 153], [0, 0, 204], [255, 51, 153], [0, 204, 204], [0, 51, 0], [255, 153, 51], [0, 204, 0]]
|
||||
MASK_COLORMAP = [[0, 0, 0], [255, 255, 255], [255, 255, 255], [255, 255, 255], [255, 255, 255], [255, 255, 255], [255, 255, 255], [255, 255, 255], [0,0, 0], [0, 0, 0], [255, 255, 255], [255, 255, 255], [255, 255, 255], [0, 0, 0], [0, 0, 0], [0, 0, 0], [0, 0, 0], [0, 0, 0], [0, 0, 0]]
|
||||
label_list = ['skin', 'nose', 'eye_g', 'l_eye', 'r_eye', 'l_brow', 'r_brow', 'l_ear', 'r_ear', 'mouth', 'u_lip', 'l_lip', 'hair', 'hat', 'ear_r', 'neck_l', 'neck', 'cloth']
|
||||
|
||||
def array_to_heatmap(x):
|
||||
x = (x - x.min()) / (x.max() - x.min()) * 255
|
||||
x = x.astype(np.uint8)
|
||||
return cv.applyColorMap(x.astype(np.uint8), cv.COLORMAP_RAINBOW)
|
||||
|
||||
def img_to_tensor(img_path, device, size=None, mode='rgb'):
|
||||
"""
|
||||
Read image from img_path, and convert to (C, H, W) tensor in range [-1, 1]
|
||||
"""
|
||||
img = Image.open(img_path).convert('RGB')
|
||||
img = np.array(img)
|
||||
if mode=='bgr':
|
||||
img = img[..., ::-1]
|
||||
if size:
|
||||
img = cv.resize(img, size)
|
||||
img = img / 255 * 2 - 1
|
||||
img_tensor = torch.from_numpy(img.transpose(2, 0, 1)).unsqueeze(0).to(device)
|
||||
return img_tensor.float()
|
||||
|
||||
def tensor_to_img(tensor, save_path=None, size=None, mode='RGB', normal=[-1, 1]):
|
||||
"""
|
||||
mode: RGB or L (gray image)
|
||||
Input: tensor with shape (C, H, W)
|
||||
Output: PIL Image
|
||||
"""
|
||||
if isinstance(size, int):
|
||||
size = (size, size)
|
||||
img_array = tensor.squeeze().data.cpu().numpy()
|
||||
if mode == 'RGB':
|
||||
img_array = img_array.transpose(1, 2, 0)
|
||||
|
||||
if size is not None:
|
||||
img_array = cv.resize(img_array, size, interpolation=cv.INTER_LINEAR)
|
||||
|
||||
if len(normal):
|
||||
img_array = (img_array - normal[0]) / (normal[1] - normal[0]) * 255
|
||||
img_array = img_array.clip(0, 255)
|
||||
|
||||
img_array = img_array.astype(np.uint8)
|
||||
if save_path:
|
||||
img = Image.fromarray(img_array, mode)
|
||||
img.save(save_path)
|
||||
|
||||
return img_array
|
||||
|
||||
def tensor_to_numpy(tensor):
|
||||
return tensor.data.cpu().numpy()
|
||||
|
||||
def batch_numpy_to_image(array, size=None):
|
||||
"""
|
||||
Input: numpy array (B, C, H, W) in [-1, 1]
|
||||
"""
|
||||
if isinstance(size, int):
|
||||
size = (size, size)
|
||||
|
||||
out_imgs = []
|
||||
array = np.clip((array + 1)/2 * 255, 0, 255)
|
||||
array = np.transpose(array, (0, 2, 3, 1))
|
||||
for i in range(array.shape[0]):
|
||||
if size is not None:
|
||||
tmp_array = cv.resize(array[i], size)
|
||||
else:
|
||||
tmp_array = array[i]
|
||||
out_imgs.append(tmp_array)
|
||||
return np.array(out_imgs).astype(np.uint8)
|
||||
|
||||
def batch_tensor_to_img(tensor, size=None):
|
||||
"""
|
||||
Input: (B, C, H, W)
|
||||
Return: RGB image, [0, 255]
|
||||
"""
|
||||
arrays = tensor_to_numpy(tensor)
|
||||
out_imgs = batch_numpy_to_image(arrays, size)
|
||||
return out_imgs
|
||||
|
||||
def color_parse_map(tensor, size=None):
|
||||
"""
|
||||
input: tensor or batch tensor
|
||||
return: colorized parsing maps
|
||||
"""
|
||||
if len(tensor.shape) < 4:
|
||||
tensor = tensor.unsqueeze(0)
|
||||
if tensor.shape[1] > 1:
|
||||
tensor = tensor.argmax(dim=1)
|
||||
|
||||
tensor = tensor.squeeze(1).data.cpu().numpy()
|
||||
color_maps = []
|
||||
for t in tensor:
|
||||
tmp_img = np.zeros(tensor.shape[1:] + (3,))
|
||||
for idx, color in enumerate(MASK_COLORMAP):
|
||||
tmp_img[t == idx] = color
|
||||
if size is not None:
|
||||
tmp_img = cv.resize(tmp_img, (size, size))
|
||||
color_maps.append(tmp_img.astype(np.uint8))
|
||||
return color_maps
|
||||
|
||||
def onehot_parse_map(img):
|
||||
"""
|
||||
input: RGB color parse map
|
||||
output: one hot encoding of parse map
|
||||
"""
|
||||
n_label = len(MASK_COLORMAP)
|
||||
img = np.array(img, dtype=np.uint8)
|
||||
h, w = img.shape[:2]
|
||||
onehot_label = np.zeros((n_label, h, w))
|
||||
colormap = np.array(MASK_COLORMAP).reshape(n_label, 1, 1, 3)
|
||||
colormap = np.tile(colormap, (1, h, w, 1))
|
||||
for idx, color in enumerate(MASK_COLORMAP):
|
||||
tmp_label = colormap[idx] == img
|
||||
onehot_label[idx] = tmp_label[..., 0] * tmp_label[..., 1] * tmp_label[..., 2]
|
||||
return onehot_label
|
||||
|
||||
|
||||
def mkdirs(paths):
|
||||
if isinstance(paths, list) and not isinstance(paths, str):
|
||||
for path in paths:
|
||||
if not os.path.exists(path):
|
||||
os.makedirs(path)
|
||||
else:
|
||||
if not os.path.exists(paths):
|
||||
os.makedirs(paths)
|
||||
|
||||
|
||||
def get_gpu_memory_map():
|
||||
"""Get the current gpu usage within visible cuda devices.
|
||||
|
||||
Returns
|
||||
-------
|
||||
Memory Map: dict
|
||||
Keys are device ids as integers.
|
||||
Values are memory usage as integers in MB.
|
||||
Device Ids: gpu ids sorted in descending order according to the available memory.
|
||||
"""
|
||||
result = subprocess.check_output(
|
||||
[
|
||||
'nvidia-smi', '--query-gpu=memory.used',
|
||||
'--format=csv,nounits,noheader'
|
||||
]).decode('utf-8')
|
||||
# Convert lines into a dictionary
|
||||
gpu_memory = np.array([int(x) for x in result.strip().split('\n')])
|
||||
if 'CUDA_VISIBLE_DEVICES' in os.environ:
|
||||
visible_devices = sorted([int(x) for x in os.environ['CUDA_VISIBLE_DEVICES'].split(',')])
|
||||
else:
|
||||
visible_devices = range(len(gpu_memory))
|
||||
gpu_memory_map = dict(zip(range(len(visible_devices)), gpu_memory[visible_devices]))
|
||||
return gpu_memory_map, sorted(gpu_memory_map, key=gpu_memory_map.get)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
hm = torch.randn(32, 68, 128, 128).cuda()
|
||||
flip(hm, 2)
|
||||
x = torch.ones(32, 68)
|
||||
y = torch.ones(32, 68)
|
||||
print(get_gpu_memory_map())
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user