vsr v1.0.0
This commit is contained in:
YaoFANGUK
@@ -0,0 +1,14 @@
|
||||
channels:
|
||||
- defaults
|
||||
show_channel_urls: true
|
||||
default_channels:
|
||||
- http://mirrors.aliyun.com/anaconda/pkgs/main
|
||||
- http://mirrors.aliyun.com/anaconda/pkgs/r
|
||||
- http://mirrors.aliyun.com/anaconda/pkgs/msys2
|
||||
custom_channels:
|
||||
conda-forge: http://mirrors.aliyun.com/anaconda/cloud
|
||||
msys2: http://mirrors.aliyun.com/anaconda/cloud
|
||||
bioconda: http://mirrors.aliyun.com/anaconda/cloud
|
||||
menpo: http://mirrors.aliyun.com/anaconda/cloud
|
||||
pytorch: http://mirrors.aliyun.com/anaconda/cloud
|
||||
simpleitk: http://mirrors.aliyun.com/anaconda/cloud
|
||||
+373
@@ -0,0 +1,373 @@
|
||||
|
||||
# Created by https://www.toptal.com/developers/gitignore/api/intellij+all,python,pycharm+all,macos,windows
|
||||
# Edit at https://www.toptal.com/developers/gitignore?templates=intellij+all,python,pycharm+all,macos,windows
|
||||
|
||||
### Intellij+all ###
|
||||
# Covers JetBrains IDEs: IntelliJ, RubyMine, PhpStorm, AppCode, PyCharm, CLion, Android Studio, WebStorm and Rider
|
||||
# Reference: https://intellij-support.jetbrains.com/hc/en-us/articles/206544839
|
||||
|
||||
# User-specific stuff
|
||||
.idea/**/workspace.xml
|
||||
.idea/**/tasks.xml
|
||||
.idea/**/usage.statistics.xml
|
||||
.idea/**/dictionaries
|
||||
.idea/**/shelf
|
||||
|
||||
# Generated files
|
||||
.idea/**/contentModel.xml
|
||||
|
||||
# Sensitive or high-churn files
|
||||
.idea/**/dataSources/
|
||||
.idea/**/dataSources.ids
|
||||
.idea/**/dataSources.local.xml
|
||||
.idea/**/sqlDataSources.xml
|
||||
.idea/**/dynamic.xml
|
||||
.idea/**/uiDesigner.xml
|
||||
.idea/**/dbnavigator.xml
|
||||
|
||||
# Gradle
|
||||
.idea/**/gradle.xml
|
||||
.idea/**/libraries
|
||||
|
||||
# Gradle and Maven with auto-import
|
||||
# When using Gradle or Maven with auto-import, you should exclude module files,
|
||||
# since they will be recreated, and may cause churn. Uncomment if using
|
||||
# auto-import.
|
||||
# .idea/artifacts
|
||||
# .idea/compiler.xml
|
||||
# .idea/jarRepositories.xml
|
||||
# .idea/modules.xml
|
||||
# .idea/*.iml
|
||||
# .idea/modules
|
||||
# *.iml
|
||||
# *.ipr
|
||||
|
||||
# CMake
|
||||
cmake-build-*/
|
||||
|
||||
# Mongo Explorer plugin
|
||||
.idea/**/mongoSettings.xml
|
||||
|
||||
# File-based project format
|
||||
*.iws
|
||||
|
||||
# IntelliJ
|
||||
out/
|
||||
|
||||
# mpeltonen/sbt-idea plugin
|
||||
.idea_modules/
|
||||
|
||||
# JIRA plugin
|
||||
atlassian-ide-plugin.xml
|
||||
|
||||
# Cursive Clojure plugin
|
||||
.idea/replstate.xml
|
||||
|
||||
# Crashlytics plugin (for Android Studio and IntelliJ)
|
||||
com_crashlytics_export_strings.xml
|
||||
crashlytics.properties
|
||||
crashlytics-build.properties
|
||||
fabric.properties
|
||||
|
||||
# Editor-based Rest Client
|
||||
.idea/httpRequests
|
||||
|
||||
# Android studio 3.1+ serialized cache file
|
||||
.idea/caches/build_file_checksums.ser
|
||||
|
||||
### Intellij+all Patch ###
|
||||
# Ignores the whole .idea folder and all .iml files
|
||||
# See https://github.com/joeblau/gitignore.io/issues/186 and https://github.com/joeblau/gitignore.io/issues/360
|
||||
|
||||
.idea/
|
||||
|
||||
# Reason: https://github.com/joeblau/gitignore.io/issues/186#issuecomment-249601023
|
||||
|
||||
*.iml
|
||||
modules.xml
|
||||
.idea/misc.xml
|
||||
*.ipr
|
||||
|
||||
# Sonarlint plugin
|
||||
.idea/sonarlint
|
||||
|
||||
### macOS ###
|
||||
# General
|
||||
.DS_Store
|
||||
.AppleDouble
|
||||
.LSOverride
|
||||
|
||||
# Icon must end with two \r
|
||||
Icon
|
||||
|
||||
|
||||
# Thumbnails
|
||||
._*
|
||||
|
||||
# Files that might appear in the root of a volume
|
||||
.DocumentRevisions-V100
|
||||
.fseventsd
|
||||
.Spotlight-V100
|
||||
.TemporaryItems
|
||||
.Trashes
|
||||
.VolumeIcon.icns
|
||||
.com.apple.timemachine.donotpresent
|
||||
|
||||
# Directories potentially created on remote AFP share
|
||||
.AppleDB
|
||||
.AppleDesktop
|
||||
Network Trash Folder
|
||||
Temporary Items
|
||||
.apdisk
|
||||
|
||||
### PyCharm+all ###
|
||||
# Covers JetBrains IDEs: IntelliJ, RubyMine, PhpStorm, AppCode, PyCharm, CLion, Android Studio, WebStorm and Rider
|
||||
# Reference: https://intellij-support.jetbrains.com/hc/en-us/articles/206544839
|
||||
|
||||
# User-specific stuff
|
||||
|
||||
# Generated files
|
||||
|
||||
# Sensitive or high-churn files
|
||||
|
||||
# Gradle
|
||||
|
||||
# Gradle and Maven with auto-import
|
||||
# When using Gradle or Maven with auto-import, you should exclude module files,
|
||||
# since they will be recreated, and may cause churn. Uncomment if using
|
||||
# auto-import.
|
||||
# .idea/artifacts
|
||||
# .idea/compiler.xml
|
||||
# .idea/jarRepositories.xml
|
||||
# .idea/modules.xml
|
||||
# .idea/*.iml
|
||||
# .idea/modules
|
||||
# *.iml
|
||||
# *.ipr
|
||||
|
||||
# CMake
|
||||
|
||||
# Mongo Explorer plugin
|
||||
|
||||
# File-based project format
|
||||
|
||||
# IntelliJ
|
||||
|
||||
# mpeltonen/sbt-idea plugin
|
||||
|
||||
# JIRA plugin
|
||||
|
||||
# Cursive Clojure plugin
|
||||
|
||||
# Crashlytics plugin (for Android Studio and IntelliJ)
|
||||
|
||||
# Editor-based Rest Client
|
||||
|
||||
# Android studio 3.1+ serialized cache file
|
||||
|
||||
### PyCharm+all Patch ###
|
||||
# Ignores the whole .idea folder and all .iml files
|
||||
# See https://github.com/joeblau/gitignore.io/issues/186 and https://github.com/joeblau/gitignore.io/issues/360
|
||||
|
||||
|
||||
# Reason: https://github.com/joeblau/gitignore.io/issues/186#issuecomment-249601023
|
||||
|
||||
|
||||
# Sonarlint plugin
|
||||
|
||||
### Python ###
|
||||
# Byte-compiled / optimized / DLL files
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*$py.class
|
||||
|
||||
# C extensions
|
||||
*.so
|
||||
|
||||
# Distribution / packaging
|
||||
.Python
|
||||
build/
|
||||
develop-eggs/
|
||||
dist/
|
||||
downloads/
|
||||
eggs/
|
||||
.eggs/
|
||||
parts/
|
||||
sdist/
|
||||
var/
|
||||
wheels/
|
||||
pip-wheel-metadata/
|
||||
share/python-wheels/
|
||||
*.egg-info/
|
||||
.installed.cfg
|
||||
*.egg
|
||||
MANIFEST
|
||||
|
||||
# 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/
|
||||
.nox/
|
||||
.coverage
|
||||
.coverage.*
|
||||
.cache
|
||||
nosetests.xml
|
||||
coverage.xml
|
||||
*.cover
|
||||
*.py,cover
|
||||
.hypothesis/
|
||||
.pytest_cache/
|
||||
pytestdebug.log
|
||||
|
||||
# Translations
|
||||
*.mo
|
||||
*.pot
|
||||
|
||||
# Django stuff:
|
||||
*.log
|
||||
local_settings.py
|
||||
db.sqlite3
|
||||
db.sqlite3-journal
|
||||
|
||||
# Flask stuff:
|
||||
instance/
|
||||
.webassets-cache
|
||||
|
||||
# Scrapy stuff:
|
||||
.scrapy
|
||||
|
||||
# Sphinx documentation
|
||||
docs/_build/
|
||||
doc/_build/
|
||||
|
||||
# PyBuilder
|
||||
target/
|
||||
|
||||
# Jupyter Notebook
|
||||
.ipynb_checkpoints
|
||||
|
||||
# IPython
|
||||
profile_default/
|
||||
ipython_config.py
|
||||
|
||||
# pyenv
|
||||
.python-version
|
||||
|
||||
# pipenv
|
||||
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
|
||||
# However, in case of collaboration, if having platform-specific dependencies or dependencies
|
||||
# having no cross-platform support, pipenv may install dependencies that don't work, or not
|
||||
# install all needed dependencies.
|
||||
#Pipfile.lock
|
||||
|
||||
# poetry
|
||||
#poetry.lock
|
||||
|
||||
# PEP 582; used by e.g. github.com/David-OConnor/pyflow
|
||||
__pypackages__/
|
||||
|
||||
# Celery stuff
|
||||
celerybeat-schedule
|
||||
celerybeat.pid
|
||||
|
||||
# SageMath parsed files
|
||||
*.sage.py
|
||||
|
||||
# Environments
|
||||
# .env
|
||||
.env/
|
||||
.venv/
|
||||
env/
|
||||
venv/
|
||||
ENV/
|
||||
env.bak/
|
||||
venv.bak/
|
||||
pythonenv*
|
||||
|
||||
# Spyder project settings
|
||||
.spyderproject
|
||||
.spyproject
|
||||
|
||||
# Rope project settings
|
||||
.ropeproject
|
||||
|
||||
# mkdocs documentation
|
||||
/site
|
||||
|
||||
# mypy
|
||||
.mypy_cache/
|
||||
.dmypy.json
|
||||
dmypy.json
|
||||
|
||||
# Pyre type checker
|
||||
.pyre/
|
||||
|
||||
# pytype static type analyzer
|
||||
.pytype/
|
||||
|
||||
# operating system-related files
|
||||
*.DS_Store #file properties cache/storage on macOS
|
||||
Thumbs.db #thumbnail cache on Windows
|
||||
|
||||
# profiling data
|
||||
.prof
|
||||
|
||||
|
||||
### Windows ###
|
||||
# Windows thumbnail cache files
|
||||
Thumbs.db
|
||||
Thumbs.db:encryptable
|
||||
ehthumbs.db
|
||||
ehthumbs_vista.db
|
||||
|
||||
# Dump file
|
||||
*.stackdump
|
||||
|
||||
# Folder config file
|
||||
[Dd]esktop.ini
|
||||
|
||||
# Recycle Bin used on file shares
|
||||
$RECYCLE.BIN/
|
||||
|
||||
# Windows Installer files
|
||||
*.cab
|
||||
*.msi
|
||||
*.msix
|
||||
*.msm
|
||||
*.msp
|
||||
|
||||
# Windows shortcuts
|
||||
*.lnk
|
||||
|
||||
# End of https://www.toptal.com/developers/gitignore/api/intellij+all,python,pycharm+all,macos,windows
|
||||
/backend/models/V2/ch_rec/inference.pdiparams
|
||||
/backend/models/V4/ch_det/inference.pdiparams
|
||||
/backend/models/big-lama/models/best.ckpt
|
||||
/backend/models/sam/sam.pth
|
||||
/output/
|
||||
/backend/test.py
|
||||
/dylib/
|
||||
/settings.ini
|
||||
/test*.py
|
||||
/subtitle.ini
|
||||
*_no_sub.mp4
|
||||
/backend/ffmpeg/win_x64/ffmpeg.exe
|
||||
/test/pic_test/
|
||||
demo*.mp4
|
||||
out*.mp4
|
||||
test*.py
|
||||
test_*.mp4
|
||||
test*_no_sub*.mp4
|
||||
/test/coods/
|
||||
/local_test/
|
||||
/backend/models/video/ProPainter.pth
|
||||
/backend/models/big-lama/big-lama.pt
|
||||
@@ -0,0 +1,201 @@
|
||||
Apache License
|
||||
Version 2.0, January 2004
|
||||
http://www.apache.org/licenses/
|
||||
|
||||
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
||||
|
||||
1. Definitions.
|
||||
|
||||
"License" shall mean the terms and conditions for use, reproduction,
|
||||
and distribution as defined by Sections 1 through 9 of this document.
|
||||
|
||||
"Licensor" shall mean the copyright owner or entity authorized by
|
||||
the copyright owner that is granting the License.
|
||||
|
||||
"Legal Entity" shall mean the union of the acting entity and all
|
||||
other entities that control, are controlled by, or are under common
|
||||
control with that entity. For the purposes of this definition,
|
||||
"control" means (i) the power, direct or indirect, to cause the
|
||||
direction or management of such entity, whether by contract or
|
||||
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
||||
outstanding shares, or (iii) beneficial ownership of such entity.
|
||||
|
||||
"You" (or "Your") shall mean an individual or Legal Entity
|
||||
exercising permissions granted by this License.
|
||||
|
||||
"Source" form shall mean the preferred form for making modifications,
|
||||
including but not limited to software source code, documentation
|
||||
source, and configuration files.
|
||||
|
||||
"Object" form shall mean any form resulting from mechanical
|
||||
transformation or translation of a Source form, including but
|
||||
not limited to compiled object code, generated documentation,
|
||||
and conversions to other media types.
|
||||
|
||||
"Work" shall mean the work of authorship, whether in Source or
|
||||
Object form, made available under the License, as indicated by a
|
||||
copyright notice that is included in or attached to the work
|
||||
(an example is provided in the Appendix below).
|
||||
|
||||
"Derivative Works" shall mean any work, whether in Source or Object
|
||||
form, that is based on (or derived from) the Work and for which the
|
||||
editorial revisions, annotations, elaborations, or other modifications
|
||||
represent, as a whole, an original work of authorship. For the purposes
|
||||
of this License, Derivative Works shall not include works that remain
|
||||
separable from, or merely link (or bind by name) to the interfaces of,
|
||||
the Work and Derivative Works thereof.
|
||||
|
||||
"Contribution" shall mean any work of authorship, including
|
||||
the original version of the Work and any modifications or additions
|
||||
to that Work or Derivative Works thereof, that is intentionally
|
||||
submitted to Licensor for inclusion in the Work by the copyright owner
|
||||
or by an individual or Legal Entity authorized to submit on behalf of
|
||||
the copyright owner. For the purposes of this definition, "submitted"
|
||||
means any form of electronic, verbal, or written communication sent
|
||||
to the Licensor or its representatives, including but not limited to
|
||||
communication on electronic mailing lists, source code control systems,
|
||||
and issue tracking systems that are managed by, or on behalf of, the
|
||||
Licensor for the purpose of discussing and improving the Work, but
|
||||
excluding communication that is conspicuously marked or otherwise
|
||||
designated in writing by the copyright owner as "Not a Contribution."
|
||||
|
||||
"Contributor" shall mean Licensor and any individual or Legal Entity
|
||||
on behalf of whom a Contribution has been received by Licensor and
|
||||
subsequently incorporated within the Work.
|
||||
|
||||
2. Grant of Copyright License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
copyright license to reproduce, prepare Derivative Works of,
|
||||
publicly display, publicly perform, sublicense, and distribute the
|
||||
Work and such Derivative Works in Source or Object form.
|
||||
|
||||
3. Grant of Patent License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
(except as stated in this section) patent license to make, have made,
|
||||
use, offer to sell, sell, import, and otherwise transfer the Work,
|
||||
where such license applies only to those patent claims licensable
|
||||
by such Contributor that are necessarily infringed by their
|
||||
Contribution(s) alone or by combination of their Contribution(s)
|
||||
with the Work to which such Contribution(s) was submitted. If You
|
||||
institute patent litigation against any entity (including a
|
||||
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
||||
or a Contribution incorporated within the Work constitutes direct
|
||||
or contributory patent infringement, then any patent licenses
|
||||
granted to You under this License for that Work shall terminate
|
||||
as of the date such litigation is filed.
|
||||
|
||||
4. Redistribution. You may reproduce and distribute copies of the
|
||||
Work or Derivative Works thereof in any medium, with or without
|
||||
modifications, and in Source or Object form, provided that You
|
||||
meet the following conditions:
|
||||
|
||||
(a) You must give any other recipients of the Work or
|
||||
Derivative Works a copy of this License; and
|
||||
|
||||
(b) You must cause any modified files to carry prominent notices
|
||||
stating that You changed the files; and
|
||||
|
||||
(c) You must retain, in the Source form of any Derivative Works
|
||||
that You distribute, all copyright, patent, trademark, and
|
||||
attribution notices from the Source form of the Work,
|
||||
excluding those notices that do not pertain to any part of
|
||||
the Derivative Works; and
|
||||
|
||||
(d) If the Work includes a "NOTICE" text file as part of its
|
||||
distribution, then any Derivative Works that You distribute must
|
||||
include a readable copy of the attribution notices contained
|
||||
within such NOTICE file, excluding those notices that do not
|
||||
pertain to any part of the Derivative Works, in at least one
|
||||
of the following places: within a NOTICE text file distributed
|
||||
as part of the Derivative Works; within the Source form or
|
||||
documentation, if provided along with the Derivative Works; or,
|
||||
within a display generated by the Derivative Works, if and
|
||||
wherever such third-party notices normally appear. The contents
|
||||
of the NOTICE file are for informational purposes only and
|
||||
do not modify the License. You may add Your own attribution
|
||||
notices within Derivative Works that You distribute, alongside
|
||||
or as an addendum to the NOTICE text from the Work, provided
|
||||
that such additional attribution notices cannot be construed
|
||||
as modifying the License.
|
||||
|
||||
You may add Your own copyright statement to Your modifications and
|
||||
may provide additional or different license terms and conditions
|
||||
for use, reproduction, or distribution of Your modifications, or
|
||||
for any such Derivative Works as a whole, provided Your use,
|
||||
reproduction, and distribution of the Work otherwise complies with
|
||||
the conditions stated in this License.
|
||||
|
||||
5. Submission of Contributions. Unless You explicitly state otherwise,
|
||||
any Contribution intentionally submitted for inclusion in the Work
|
||||
by You to the Licensor shall be under the terms and conditions of
|
||||
this License, without any additional terms or conditions.
|
||||
Notwithstanding the above, nothing herein shall supersede or modify
|
||||
the terms of any separate license agreement you may have executed
|
||||
with Licensor regarding such Contributions.
|
||||
|
||||
6. Trademarks. This License does not grant permission to use the trade
|
||||
names, trademarks, service marks, or product names of the Licensor,
|
||||
except as required for reasonable and customary use in describing the
|
||||
origin of the Work and reproducing the content of the NOTICE file.
|
||||
|
||||
7. Disclaimer of Warranty. Unless required by applicable law or
|
||||
agreed to in writing, Licensor provides the Work (and each
|
||||
Contributor provides its Contributions) on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
||||
implied, including, without limitation, any warranties or conditions
|
||||
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
||||
PARTICULAR PURPOSE. You are solely responsible for determining the
|
||||
appropriateness of using or redistributing the Work and assume any
|
||||
risks associated with Your exercise of permissions under this License.
|
||||
|
||||
8. Limitation of Liability. In no event and under no legal theory,
|
||||
whether in tort (including negligence), contract, or otherwise,
|
||||
unless required by applicable law (such as deliberate and grossly
|
||||
negligent acts) or agreed to in writing, shall any Contributor be
|
||||
liable to You for damages, including any direct, indirect, special,
|
||||
incidental, or consequential damages of any character arising as a
|
||||
result of this License or out of the use or inability to use the
|
||||
Work (including but not limited to damages for loss of goodwill,
|
||||
work stoppage, computer failure or malfunction, or any and all
|
||||
other commercial damages or losses), even if such Contributor
|
||||
has been advised of the possibility of such damages.
|
||||
|
||||
9. Accepting Warranty or Additional Liability. While redistributing
|
||||
the Work or Derivative Works thereof, You may choose to offer,
|
||||
and charge a fee for, acceptance of support, warranty, indemnity,
|
||||
or other liability obligations and/or rights consistent with this
|
||||
License. However, in accepting such obligations, You may act only
|
||||
on Your own behalf and on Your sole responsibility, not on behalf
|
||||
of any other Contributor, and only if You agree to indemnify,
|
||||
defend, and hold each Contributor harmless for any liability
|
||||
incurred by, or claims asserted against, such Contributor by reason
|
||||
of your accepting any such warranty or additional liability.
|
||||
|
||||
END OF TERMS AND CONDITIONS
|
||||
|
||||
APPENDIX: How to apply the Apache License to your work.
|
||||
|
||||
To apply the Apache License to your work, attach the following
|
||||
boilerplate notice, with the fields enclosed by brackets "[]"
|
||||
replaced with your own identifying information. (Don't include
|
||||
the brackets!) The text should be enclosed in the appropriate
|
||||
comment syntax for the file format. We also recommend that a
|
||||
file or class name and description of purpose be included on the
|
||||
same "printed page" as the copyright notice for easier
|
||||
identification within third-party archives.
|
||||
|
||||
Copyright [yyyy] [name of copyright owner]
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
@@ -0,0 +1,187 @@
|
||||
简体中文 | [English](README_en.md)
|
||||
|
||||
## 项目简介
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
Video-subtitle-remover (VSR) 是一款基于AI技术，将视频中的硬字幕去除的软件。
|
||||
主要实现了以下功能：
|
||||
- **无损分辨率**将视频中的硬字幕去除，生成去除字幕后的文件
|
||||
- 通过超强AI算法模型，对去除字幕文本的区域进行填充（非相邻像素填充与马赛克去除）
|
||||
- 支持自定义字幕位置，仅去除定义位置中的字幕（传入位置）
|
||||
- 支持全视频自动去除所有文本（不传入位置）
|
||||
|
||||
<p style="text-align:center;"><img src="https://github.com/YaoFANGUK/video-subtitle-remover/raw/main/design/demo.png" alt="demo.png"/></p>
|
||||
|
||||
**使用说明：**
|
||||
|
||||
- 有使用问题请加群讨论，QQ群：806152575
|
||||
- 直接下载压缩包解压运行，如果不能运行再按照下面的教程，尝试源码安装conda环境运行
|
||||
|
||||
**下载地址：**
|
||||
|
||||
Windows GPU版本v1.0.0（GPU）：
|
||||
|
||||
- 百度网盘: <a href="https://pan.baidu.com/s/1t0djUPjQ96X8U9HisyoiYA?pwd=vsr1">vsr_windows_gpu_v1.0.0.7z</a> 提取码：**vsr1**
|
||||
|
||||
- Google Drive: <a href="https://drive.google.com/drive/folders/1NRgLNoHHOmdO4GxLhkPbHsYfMOB_3Elr?usp=sharing">vsr_windows_gpu_v1.0.0.7z</a>
|
||||
|
||||
> 仅供具有Nvidia显卡的用户使用(AMD的显卡不行)
|
||||
|
||||
## 演示
|
||||
|
||||
- GUI版：
|
||||
|
||||
<p style="text-align:center;"><img src="https://github.com/YaoFANGUK/video-subtitle-remover/raw/main/design/demo2.gif" alt="demo2.gif"/></p>
|
||||
|
||||
- <a href="https://b23.tv/guEbl9C">点击查看演示视频👇</a>
|
||||
|
||||
<p style="text-align:center;"><a href="https://b23.tv/guEbl9C"><img src="https://github.com/YaoFANGUK/video-subtitle-remover/raw/main/design/demo.gif" alt="demo.gif"/></a></p>
|
||||
|
||||
## 源码使用说明
|
||||
|
||||
> **无Nvidia显卡请勿使用本项目**，最低配置：
|
||||
>
|
||||
> **GPU**：GTX 1060或以上显卡
|
||||
>
|
||||
> CPU: 支持AVX指令集
|
||||
|
||||
#### 1. 下载安装Miniconda
|
||||
|
||||
- Windows: <a href="https://repo.anaconda.com/miniconda/Miniconda3-py38_4.11.0-Windows-x86_64.exe">Miniconda3-py38_4.11.0-Windows-x86_64.exe</a>
|
||||
|
||||
- Linux: <a href="https://repo.anaconda.com/miniconda/Miniconda3-py38_4.11.0-Linux-x86_64.sh">Miniconda3-py38_4.11.0-Linux-x86_64.sh</a>
|
||||
|
||||
#### 2. 创建并激活虚机环境
|
||||
|
||||
（1）切换到源码所在目录：
|
||||
```shell
|
||||
cd <源码所在目录>
|
||||
```
|
||||
> 例如：如果你的源代码放在D盘的tools文件下，并且源代码的文件夹名为video-subtitle-remover，就输入 ```cd D:/tools/video-subtitle-remover-main```
|
||||
|
||||
（2）创建激活conda环境
|
||||
```shell
|
||||
conda create -n videoEnv python=3.8
|
||||
```
|
||||
|
||||
```shell
|
||||
conda activate videoEnv
|
||||
```
|
||||
|
||||
#### 3. 安装依赖文件
|
||||
|
||||
请确保你已经安装 python 3.8+，使用conda创建项目虚拟环境并激活环境 (建议创建虚拟环境运行，以免后续出现问题)
|
||||
|
||||
- 安装CUDA和cuDNN
|
||||
|
||||
<details>
|
||||
<summary>Linux用户</summary>
|
||||
<h5>(1) 下载CUDA 11.7</h5>
|
||||
<pre><code>wget https://developer.download.nvidia.com/compute/cuda/11.7.0/local_installers/cuda_11.7.0_515.43.04_linux.run</code></pre>
|
||||
<h5>(2) 安装CUDA 11.7</h5>
|
||||
<pre><code>sudo sh cuda_11.7.0_515.43.04_linux.run</code></pre>
|
||||
<p>1. 输入accept</p>
|
||||
<img src="https://i.328888.xyz/2023/03/31/iwVoeH.png" width="500" alt="">
|
||||
<p>2. 选中CUDA Toolkit 11.7（如果你没有安装nvidia驱动则选中Driver，如果你已经安装了nvidia驱动请不要选中driver），之后选中install，回车</p>
|
||||
<img src="https://i.328888.xyz/2023/03/31/iwVThJ.png" width="500" alt="">
|
||||
<p>3. 添加环境变量</p>
|
||||
<p>在 ~/.bashrc 加入以下内容</p>
|
||||
<pre><code># CUDA
|
||||
export PATH=/usr/local/cuda-11.7/bin${PATH:+:${PATH}}
|
||||
export LD_LIBRARY_PATH=/usr/local/cuda-11.7/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}</code></pre>
|
||||
<p>使其生效</p>
|
||||
<pre><code>source ~/.bashrc</code></pre>
|
||||
<h5>(3) 下载cuDNN 8.4.1</h5>
|
||||
<p>国内：<a href="https://pan.baidu.com/s/1Gd_pSVzWfX1G7zCuqz6YYA">cudnn-linux-x86_64-8.4.1.50_cuda11.6-archive.tar.xz</a> 提取码：57mg</p>
|
||||
<p>国外：<a href="https://github.com/YaoFANGUK/video-subtitle-extractor/releases/download/1.0.0/cudnn-linux-x86_64-8.4.1.50_cuda11.6-archive.tar.xz">cudnn-linux-x86_64-8.4.1.50_cuda11.6-archive.tar.xz</a></p>
|
||||
<h5>(4) 安装cuDNN 8.4.1</h5>
|
||||
<pre><code> tar -xf cudnn-linux-x86_64-8.4.1.50_cuda11.6-archive.tar.xz
|
||||
mv cudnn-linux-x86_64-8.4.1.50_cuda11.6-archive cuda
|
||||
sudo cp ./cuda/include/* /usr/local/cuda-11.7/include/
|
||||
sudo cp ./cuda/lib/* /usr/local/cuda-11.7/lib64/
|
||||
sudo chmod a+r /usr/local/cuda-11.7/lib64/*
|
||||
sudo chmod a+r /usr/local/cuda-11.7/include/*</code></pre>
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary>Windows用户</summary>
|
||||
<h5>(1) 下载CUDA 11.7</h5>
|
||||
<a href="https://developer.download.nvidia.com/compute/cuda/11.7.0/local_installers/cuda_11.7.0_516.01_windows.exe">cuda_11.7.0_516.01_windows.exe</a>
|
||||
<h5>(2) 安装CUDA 11.7</h5>
|
||||
<h5>(3) 下载cuDNN 8.2.4</h5>
|
||||
<p><a href="https://github.com/YaoFANGUK/video-subtitle-extractor/releases/download/1.0.0/cudnn-windows-x64-v8.2.4.15.zip">cudnn-windows-x64-v8.2.4.15.zip</a></p>
|
||||
<h5>(4) 安装cuDNN 8.2.4</h5>
|
||||
<p>
|
||||
将cuDNN解压后的cuda文件夹中的bin, include, lib目录下的文件复制到C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.7\对应目录下
|
||||
</p>
|
||||
</details>
|
||||
|
||||
|
||||
- 安装GPU版本Paddlepaddle:
|
||||
|
||||
- windows:
|
||||
|
||||
```shell
|
||||
python -m pip install paddlepaddle-gpu==2.4.2.post117 -f https://www.paddlepaddle.org.cn/whl/windows/mkl/avx/stable.html
|
||||
```
|
||||
|
||||
- Linux:
|
||||
|
||||
```shell
|
||||
python -m pip install paddlepaddle-gpu==2.4.2.post117 -f https://www.paddlepaddle.org.cn/whl/linux/mkl/avx/stable.html
|
||||
```
|
||||
|
||||
- 安装GPU版本Pytorch:
|
||||
|
||||
```shell
|
||||
conda install pytorch==2.0.1 torchvision==0.15.2 pytorch-cuda=11.7 -c pytorch -c nvidia
|
||||
```
|
||||
或者使用
|
||||
```shell
|
||||
pip install torch==2.0.1 torchvision==0.15.2 --index-url https://download.pytorch.org/whl/cu117
|
||||
```
|
||||
|
||||
- 安装其他依赖:
|
||||
|
||||
```shell
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
|
||||
#### 4. 运行程序
|
||||
|
||||
- 运行图形化界面
|
||||
|
||||
```shell
|
||||
python gui.py
|
||||
```
|
||||
|
||||
- 运行命令行版本(CLI)
|
||||
|
||||
```shell
|
||||
python ./backend/main.py
|
||||
```
|
||||
|
||||
## 常见问题
|
||||
1. CondaHTTPError
|
||||
|
||||
将项目中的.condarc放在用户目录下(C:/Users/<你的用户名>)，如果用户目录已经存在该文件则覆盖
|
||||
|
||||
解决方案：https://zhuanlan.zhihu.com/p/260034241
|
||||
|
||||
2. 7z文件解压错误
|
||||
|
||||
解决方案：升级7-zip解压程序到最新版本
|
||||
|
||||
## 赞助
|
||||
<img src="https://i.imgur.com/THE81jM.jpg" width="300">
|
||||
|
||||
| 捐赠者 | 累计捐赠金额 | 赞助席位 |
|
||||
| --- | --- | --- |
|
||||
| 很奇异| 15.00 RMB | 金牌赞助席位 |
|
||||
| Leo| 1.00 RMB | 银牌赞助席位 |
|
||||
| 暂无| 暂无 | 铜牌赞助席位 |
|
||||
|
||||
Executable
+177
@@ -0,0 +1,177 @@
|
||||
[简体中文](README.md) | English
|
||||
|
||||
## Project Introduction
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
Video-subtitle-remover (VSR) is an AI-based software that removes hardcoded subtitles from videos. It mainly implements the following functionalities:
|
||||
|
||||
- **Lossless resolution**: Removes hardcoded subtitles from videos and generates files without subtitles.
|
||||
- Fills in the removed subtitle text area using a powerful AI algorithm model (non-adjacent pixel filling and mosaic removal).
|
||||
- Supports custom subtitle positions by only removing subtitles in the defined location (input position).
|
||||
- Supports automatic removal of all text throughout the entire video (without inputting a position).
|
||||
|
||||
<p style="text-align:center;"><img src="https://github.com/YaoFANGUK/video-subtitle-remover/raw/main/design/demo.png" alt="demo.png"/></p>
|
||||
|
||||
> Download the .7z package directly, extract, and run it. If it cannot run, follow the tutorial below to try installing the conda environment and running the source code.
|
||||
|
||||
**Download Links:**
|
||||
|
||||
Windows GPU Version v1.0.0 (GPU):
|
||||
|
||||
- Baidu Cloud Disk: <a href="https://pan.baidu.com/s/1t0djUPjQ96X8U9HisyoiYA?pwd=vsr1">vsr_windows_gpu_v1.0.0.7z</a> Extraction Code: **vsr1**
|
||||
|
||||
- Google Drive: <a href="https://drive.google.com/drive/folders/1NRgLNoHHOmdO4GxLhkPbHsYfMOB_3Elr?usp=sharing">vsr_windows_gpu_v1.0.0.7z</a>
|
||||
|
||||
> For use only by users with Nvidia graphics cards (AMD graphics cards are not supported).
|
||||
|
||||
## Demonstration
|
||||
|
||||
- GUI:
|
||||
|
||||
<p style="text-align:center;"><img src="https://github.com/YaoFANGUK/video-subtitle-remover/raw/main/design/demo2.gif" alt="demo2.gif"/></p>
|
||||
|
||||
- <a href="https://b23.tv/guEbl9C">Click to view demo video👇</a>
|
||||
|
||||
<p style="text-align:center;"><a href="https://b23.tv/guEbl9C"><img src="https://github.com/YaoFANGUK/video-subtitle-remover/raw/main/design/demo.gif" alt="demo.gif"/></a></p>
|
||||
|
||||
## Source Code Usage Instructions
|
||||
|
||||
> **Do not use this project without an Nvidia graphics card**. The minimum requirements are:
|
||||
>
|
||||
> **GPU**: GTX 1060 or higher graphics card
|
||||
>
|
||||
> CPU: Supports AVX instruction set
|
||||
|
||||
#### 1. Download and install Miniconda
|
||||
|
||||
- Windows: <a href="https://repo.anaconda.com/miniconda/Miniconda3-py38_4.11.0-Windows-x86_64.exe">Miniconda3-py38_4.11.0-Windows-x86_64.exe</a>
|
||||
|
||||
- Linux: <a href="https://repo.anaconda.com/miniconda/Miniconda3-py38_4.11.0-Linux-x86_64.sh">Miniconda3-py38_4.11.0-Linux-x86_64.sh</a>
|
||||
|
||||
#### 2. Create and activate a virtual environment
|
||||
|
||||
(1) Switch to the source code directory:
|
||||
|
||||
```shell
|
||||
cd <source_code_directory>
|
||||
```
|
||||
|
||||
> For example, if your source code is in the `tools` folder on drive D, and the source code folder name is `video-subtitle-remover`, enter `cd D:/tools/video-subtitle-remover-main`.
|
||||
|
||||
(2) Create and activate the conda environment:
|
||||
|
||||
```shell
|
||||
conda create -n videoEnv python=3.8
|
||||
```
|
||||
|
||||
```shell
|
||||
conda activate videoEnv
|
||||
```
|
||||
|
||||
#### 3. Install dependencies
|
||||
|
||||
Please make sure you have already installed Python 3.8+, use conda to create a project virtual environment and activate the environment (it is recommended to create a virtual environment to run to avoid subsequent problems).
|
||||
|
||||
- Install **CUDA** and **cuDNN**
|
||||
|
||||
<details>
|
||||
<summary>Linux</summary>
|
||||
<h5>(1) Download CUDA 11.7</h5>
|
||||
<pre><code>wget https://developer.download.nvidia.com/compute/cuda/11.7.0/local_installers/cuda_11.7.0_515.43.04_linux.run</code></pre>
|
||||
<h5>(2) Install CUDA 11.7</h5>
|
||||
<pre><code>sudo sh cuda_11.7.0_515.43.04_linux.run</code></pre>
|
||||
<p>1. Input accept</p>
|
||||
<img src="https://i.328888.xyz/2023/03/31/iwVoeH.png" width="500" alt="">
|
||||
<p>2. make sure CUDA Toolkit 11.7 is chosen (If you have already installed driver, do not select Driver)</p>
|
||||
<img src="https://i.328888.xyz/2023/03/31/iwVThJ.png" width="500" alt="">
|
||||
<p>3. Add environment variables</p>
|
||||
<p>add the following content in <strong>~/.bashrc</strong></p>
|
||||
<pre><code># CUDA
|
||||
export PATH=/usr/local/cuda-11.7/bin${PATH:+:${PATH}}
|
||||
export LD_LIBRARY_PATH=/usr/local/cuda-11.7/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}</code></pre>
|
||||
<p>Make sure it works</p>
|
||||
<pre><code>source ~/.bashrc</code></pre>
|
||||
<h5>(3) Download cuDNN 8.4.1</h5>
|
||||
<p><a href="https://github.com/YaoFANGUK/video-subtitle-extractor/releases/download/1.0.0/cudnn-linux-x86_64-8.4.1.50_cuda11.6-archive.tar.xz">cudnn-linux-x86_64-8.4.1.50_cuda11.6-archive.tar.xz</a></p>
|
||||
<h5>(4) Install cuDNN 8.4.1</h5>
|
||||
<pre><code> tar -xf cudnn-linux-x86_64-8.4.1.50_cuda11.6-archive.tar.xz
|
||||
mv cudnn-linux-x86_64-8.4.1.50_cuda11.6-archive cuda
|
||||
sudo cp ./cuda/include/* /usr/local/cuda-11.7/include/
|
||||
sudo cp ./cuda/lib/* /usr/local/cuda-11.7/lib64/
|
||||
sudo chmod a+r /usr/local/cuda-11.7/lib64/*
|
||||
sudo chmod a+r /usr/local/cuda-11.7/include/*</code></pre>
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary>Windows</summary>
|
||||
<h5>(1) Download CUDA 11.7</h5>
|
||||
<a href="https://developer.download.nvidia.com/compute/cuda/11.7.0/local_installers/cuda_11.7.0_516.01_windows.exe">cuda_11.7.0_516.01_windows.exe</a>
|
||||
<h5>(2) Install CUDA 11.7</h5>
|
||||
<h5>(3) Download cuDNN 8.2.4</h5>
|
||||
<p><a href="https://github.com/YaoFANGUK/video-subtitle-extractor/releases/download/1.0.0/cudnn-windows-x64-v8.2.4.15.zip">cudnn-windows-x64-v8.2.4.15.zip</a></p>
|
||||
<h5>(4) Install cuDNN 8.2.4</h5>
|
||||
<p>
|
||||
unzip "cudnn-windows-x64-v8.2.4.15.zip", then move all files in "bin, include, lib" in cuda
|
||||
directory to C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.7\
|
||||
</p>
|
||||
</details>
|
||||
|
||||
|
||||
- Install GPU version of Paddlepaddle:
|
||||
- windows:
|
||||
|
||||
```shell
|
||||
python -m pip install paddlepaddle-gpu==2.4.2.post117 -f https://www.paddlepaddle.org.cn/whl/windows/mkl/avx/stable.html
|
||||
```
|
||||
|
||||
- Linux:
|
||||
|
||||
```shell
|
||||
python -m pip install paddlepaddle-gpu==2.4.2.post117 -f https://www.paddlepaddle.org.cn/whl/linux/mkl/avx/stable.html
|
||||
```
|
||||
|
||||
- Install GPU version of Pytorch:
|
||||
|
||||
```shell
|
||||
conda install pytorch==2.0.1 torchvision==0.15.2 pytorch-cuda=11.7 -c pytorch -c nvidia
|
||||
```
|
||||
or use
|
||||
|
||||
```shell
|
||||
pip install torch==2.0.1 torchvision==0.15.2 --index-url https://download.pytorch.org/whl/cu117
|
||||
```
|
||||
|
||||
- Install other dependencies:
|
||||
|
||||
```shell
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
|
||||
#### 4. Run the program
|
||||
|
||||
- Run the graphical interface
|
||||
|
||||
```shell
|
||||
python gui.py
|
||||
```
|
||||
|
||||
- Run the command line version (CLI)
|
||||
|
||||
```shell
|
||||
python ./backend/main.py
|
||||
```
|
||||
|
||||
## Common Issues
|
||||
1. CondaHTTPError
|
||||
|
||||
Place the .condarc file from the project in the user directory (C:/Users/<your_username>). If the file already exists in the user directory, overwrite it.
|
||||
|
||||
Solution: https://zhuanlan.zhihu.com/p/260034241
|
||||
|
||||
2. 7z file extraction error
|
||||
|
||||
Solution: Upgrade the 7-zip extraction program to the latest version.
|
||||
@@ -0,0 +1,3 @@
|
||||
import warnings
|
||||
# 忽略所有的 DeprecationWarning
|
||||
warnings.filterwarnings("ignore", category=DeprecationWarning)
|
||||
@@ -0,0 +1,79 @@
|
||||
import warnings
|
||||
warnings.filterwarnings('ignore')
|
||||
import os
|
||||
import torch
|
||||
import logging
|
||||
import platform
|
||||
import stat
|
||||
from fsplit.filesplit import Filesplit
|
||||
import paddle
|
||||
paddle.disable_signal_handler()
|
||||
logging.disable(logging.DEBUG) # 关闭DEBUG日志的打印
|
||||
logging.disable(logging.WARNING) # 关闭WARNING日志的打印
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
||||
LAMA_MODEL_PATH = os.path.join(BASE_DIR, 'models', 'big-lama')
|
||||
VIDEO_INPAINT_MODEL_PATH = os.path.join(BASE_DIR, 'models', 'video')
|
||||
MODEL_VERSION = 'V4'
|
||||
DET_MODEL_BASE = os.path.join(BASE_DIR, 'models')
|
||||
DET_MODEL_PATH = os.path.join(DET_MODEL_BASE, MODEL_VERSION, 'ch_det')
|
||||
|
||||
# ×××××××××××××××××××× [可以改] start ××××××××××××××××××××
|
||||
# 容忍的像素点偏差
|
||||
PIXEL_TOLERANCE_Y = 20 # 允许检测框纵向偏差50个像素点
|
||||
PIXEL_TOLERANCE_X = 20 # 允许检测框横向偏差100个像素点
|
||||
# 字幕区域偏移量， 放大诗歌像素点，防止字幕偏移
|
||||
SUBTITLE_AREA_DEVIATION_PIXEL = 10
|
||||
# 20个像素点以内的差距认为是同一行
|
||||
TOLERANCE_Y = 20
|
||||
# 高度差阈值
|
||||
THRESHOLD_HEIGHT_DIFFERENCE = 20
|
||||
# 【根据自己的GPU显存大小设置】最大同时处理的图片数量，设置越大处理效果越好，但是要求显存越高
|
||||
# 1280x720p视频设置80需要25G显存，设置50需要19G显存
|
||||
# 720x480p视频设置80需要8G显存，设置50需要7G显存
|
||||
MAX_PROCESS_NUM = 70
|
||||
# 【根据自己内存大小设置，应该大于等于MAX_PROCESS_NUM】
|
||||
MAX_LOAD_NUM = 70
|
||||
# 是否开启精细模式，开启精细模式将消耗大量GPU显存，如果您的显卡显存较少，建议设置为False
|
||||
ACCURATE_MODE = True
|
||||
# 是否开启快速模型，不保证inpaint效果
|
||||
FAST_MODE = False
|
||||
# ×××××××××××××××××××× [可以改] start ××××××××××××××××××××
|
||||
|
||||
|
||||
# ×××××××××××××××××××× [不要改] start ××××××××××××××××××××
|
||||
# 查看该路径下是否有模型完整文件，没有的话合并小文件生成完整文件
|
||||
if 'big-lama.pt' not in (os.listdir(LAMA_MODEL_PATH)):
|
||||
fs = Filesplit()
|
||||
fs.merge(input_dir=LAMA_MODEL_PATH)
|
||||
|
||||
if 'inference.pdiparams' not in os.listdir(DET_MODEL_PATH):
|
||||
fs = Filesplit()
|
||||
fs.merge(input_dir=DET_MODEL_PATH)
|
||||
|
||||
if 'ProPainter.pth' not in os.listdir(VIDEO_INPAINT_MODEL_PATH):
|
||||
fs = Filesplit()
|
||||
fs.merge(input_dir=VIDEO_INPAINT_MODEL_PATH)
|
||||
|
||||
# 指定ffmpeg可执行程序路径
|
||||
sys_str = platform.system()
|
||||
if sys_str == "Windows":
|
||||
ffmpeg_bin = os.path.join('win_x64', 'ffmpeg.exe')
|
||||
elif sys_str == "Linux":
|
||||
ffmpeg_bin = os.path.join('linux_x64', 'ffmpeg')
|
||||
else:
|
||||
ffmpeg_bin = os.path.join('macos', 'ffmpeg')
|
||||
FFMPEG_PATH = os.path.join(BASE_DIR, '', 'ffmpeg', ffmpeg_bin)
|
||||
|
||||
if 'ffmpeg.exe' not in os.listdir(os.path.join(BASE_DIR, '', 'ffmpeg', 'win_x64')):
|
||||
fs = Filesplit()
|
||||
fs.merge(input_dir=os.path.join(BASE_DIR, '', 'ffmpeg', 'win_x64'))
|
||||
# 将ffmpeg添加可执行权限
|
||||
os.chmod(FFMPEG_PATH, stat.S_IRWXU+stat.S_IRWXG+stat.S_IRWXO)
|
||||
|
||||
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
|
||||
|
||||
# 如果开启了快速模式，则强制关闭ACCURATE_MODE
|
||||
if FAST_MODE:
|
||||
ACCURATE_MODE = False
|
||||
# ×××××××××××××××××××× [不要改] end ××××××××××××××××××××
|
||||
Executable
BIN
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,4 @@
|
||||
filename,filesize,encoding,header
|
||||
ffmpeg_1.exe,50000000,,
|
||||
ffmpeg_2.exe,50000000,,
|
||||
ffmpeg_3.exe,13721856,,
|
||||
|
@@ -0,0 +1,27 @@
|
||||
import os
|
||||
from typing import Union
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from backend.inpaint.utils.lama_util import prepare_img_and_mask
|
||||
from backend import config
|
||||
|
||||
|
||||
class LamaInpaint:
|
||||
def __init__(self, device: torch.device = torch.device("cuda" if torch.cuda.is_available() else "cpu"), model_path=None) -> None:
|
||||
if model_path is None:
|
||||
model_path = os.path.join(config.LAMA_MODEL_PATH, 'big-lama.pt')
|
||||
self.model = torch.jit.load(model_path, map_location=device)
|
||||
self.model.eval()
|
||||
self.model.to(device)
|
||||
self.device = device
|
||||
|
||||
def __call__(self, image: Union[Image.Image, np.ndarray], mask: Union[Image.Image, np.ndarray]):
|
||||
image, mask = prepare_img_and_mask(image, mask, self.device)
|
||||
with torch.inference_mode():
|
||||
inpainted = self.model(image, mask)
|
||||
cur_res = inpainted[0].permute(1, 2, 0).detach().cpu().numpy()
|
||||
cur_res = np.clip(cur_res * 255, 0, 255).astype(np.uint8)
|
||||
return cur_res
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
from .utils import *
|
||||
@@ -0,0 +1,80 @@
|
||||
import os
|
||||
import sys
|
||||
import torch
|
||||
import numpy as np
|
||||
import cv2
|
||||
from PIL import Image
|
||||
from torch.hub import download_url_to_file, get_dir
|
||||
from urllib.parse import urlparse
|
||||
|
||||
|
||||
# Source https://github.com/advimman/lama
|
||||
def get_image(image):
|
||||
if isinstance(image, Image.Image):
|
||||
img = np.array(image)
|
||||
elif isinstance(image, np.ndarray):
|
||||
img = image.copy()
|
||||
else:
|
||||
raise Exception("Input image should be either PIL Image or numpy array!")
|
||||
|
||||
if img.ndim == 3:
|
||||
img = np.transpose(img, (2, 0, 1)) # chw
|
||||
elif img.ndim == 2:
|
||||
img = img[np.newaxis, ...]
|
||||
|
||||
assert img.ndim == 3
|
||||
|
||||
img = img.astype(np.float32) / 255
|
||||
return img
|
||||
|
||||
|
||||
def ceil_modulo(x, mod):
|
||||
if x % mod == 0:
|
||||
return x
|
||||
return (x // mod + 1) * mod
|
||||
|
||||
|
||||
def scale_image(img, factor, interpolation=cv2.INTER_AREA):
|
||||
if img.shape[0] == 1:
|
||||
img = img[0]
|
||||
else:
|
||||
img = np.transpose(img, (1, 2, 0))
|
||||
|
||||
img = cv2.resize(img, dsize=None, fx=factor, fy=factor, interpolation=interpolation)
|
||||
|
||||
if img.ndim == 2:
|
||||
img = img[None, ...]
|
||||
else:
|
||||
img = np.transpose(img, (2, 0, 1))
|
||||
return img
|
||||
|
||||
|
||||
def pad_img_to_modulo(img, mod):
|
||||
channels, height, width = img.shape
|
||||
out_height = ceil_modulo(height, mod)
|
||||
out_width = ceil_modulo(width, mod)
|
||||
return np.pad(
|
||||
img,
|
||||
((0, 0), (0, out_height - height), (0, out_width - width)),
|
||||
mode="symmetric",
|
||||
)
|
||||
|
||||
|
||||
def prepare_img_and_mask(image, mask, device, pad_out_to_modulo=8, scale_factor=None):
|
||||
out_image = get_image(image)
|
||||
out_mask = get_image(mask)
|
||||
|
||||
if scale_factor is not None:
|
||||
out_image = scale_image(out_image, scale_factor)
|
||||
out_mask = scale_image(out_mask, scale_factor, interpolation=cv2.INTER_NEAREST)
|
||||
|
||||
if pad_out_to_modulo is not None and pad_out_to_modulo > 1:
|
||||
out_image = pad_img_to_modulo(out_image, pad_out_to_modulo)
|
||||
out_mask = pad_img_to_modulo(out_mask, pad_out_to_modulo)
|
||||
|
||||
out_image = torch.from_numpy(out_image).unsqueeze(0).to(device)
|
||||
out_mask = torch.from_numpy(out_mask).unsqueeze(0).to(device)
|
||||
|
||||
out_mask = (out_mask > 0) * 1
|
||||
|
||||
return out_image, out_mask
|
||||
@@ -0,0 +1,85 @@
|
||||
import cv2
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from typing import Any, Dict, List
|
||||
|
||||
|
||||
def load_img_to_array(img_p):
|
||||
img = Image.open(img_p)
|
||||
if img.mode == "RGBA":
|
||||
img = img.convert("RGB")
|
||||
return np.array(img)
|
||||
|
||||
|
||||
def save_array_to_img(img_arr, img_p):
|
||||
Image.fromarray(img_arr.astype(np.uint8)).save(img_p)
|
||||
|
||||
|
||||
def dilate_mask(mask, dilate_factor=15):
|
||||
mask = mask.astype(np.uint8)
|
||||
mask = cv2.dilate(
|
||||
mask,
|
||||
np.ones((dilate_factor, dilate_factor), np.uint8),
|
||||
iterations=1
|
||||
)
|
||||
return mask
|
||||
|
||||
def erode_mask(mask, dilate_factor=15):
|
||||
mask = mask.astype(np.uint8)
|
||||
mask = cv2.erode(
|
||||
mask,
|
||||
np.ones((dilate_factor, dilate_factor), np.uint8),
|
||||
iterations=1
|
||||
)
|
||||
return mask
|
||||
|
||||
def show_mask(ax, mask: np.ndarray, random_color=False):
|
||||
mask = mask.astype(np.uint8)
|
||||
if np.max(mask) == 255:
|
||||
mask = mask / 255
|
||||
if random_color:
|
||||
color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
|
||||
else:
|
||||
color = np.array([30 / 255, 144 / 255, 255 / 255, 0.6])
|
||||
h, w = mask.shape[-2:]
|
||||
mask_img = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
|
||||
ax.imshow(mask_img)
|
||||
|
||||
|
||||
def show_points(ax, coords: List[List[float]], labels: List[int], size=375):
|
||||
coords = np.array(coords)
|
||||
labels = np.array(labels)
|
||||
color_table = {0: 'red', 1: 'green'}
|
||||
for label_value, color in color_table.items():
|
||||
points = coords[labels == label_value]
|
||||
ax.scatter(points[:, 0], points[:, 1], color=color, marker='*',
|
||||
s=size, edgecolor='white', linewidth=1.25)
|
||||
|
||||
def get_clicked_point(img_path):
|
||||
img = cv2.imread(img_path)
|
||||
cv2.namedWindow("image")
|
||||
cv2.imshow("image", img)
|
||||
|
||||
last_point = []
|
||||
keep_looping = True
|
||||
|
||||
def mouse_callback(event, x, y, flags, param):
|
||||
nonlocal last_point, keep_looping, img
|
||||
|
||||
if event == cv2.EVENT_LBUTTONDOWN:
|
||||
if last_point:
|
||||
cv2.circle(img, tuple(last_point), 5, (0, 0, 0), -1)
|
||||
last_point = [x, y]
|
||||
cv2.circle(img, tuple(last_point), 5, (0, 0, 255), -1)
|
||||
cv2.imshow("image", img)
|
||||
elif event == cv2.EVENT_RBUTTONDOWN:
|
||||
keep_looping = False
|
||||
|
||||
cv2.setMouseCallback("image", mouse_callback)
|
||||
|
||||
while keep_looping:
|
||||
cv2.waitKey(1)
|
||||
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
return last_point
|
||||
@@ -0,0 +1,232 @@
|
||||
import os
|
||||
import json
|
||||
import random
|
||||
|
||||
import cv2
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
|
||||
import torch
|
||||
import torchvision.transforms as transforms
|
||||
|
||||
from utils.file_client import FileClient
|
||||
from utils.img_util import imfrombytes
|
||||
from utils.flow_util import resize_flow, flowread
|
||||
from core.utils import (create_random_shape_with_random_motion, Stack,
|
||||
ToTorchFormatTensor, GroupRandomHorizontalFlip,GroupRandomHorizontalFlowFlip)
|
||||
|
||||
|
||||
class TrainDataset(torch.utils.data.Dataset):
|
||||
def __init__(self, args: dict):
|
||||
self.args = args
|
||||
self.video_root = args['video_root']
|
||||
self.flow_root = args['flow_root']
|
||||
self.num_local_frames = args['num_local_frames']
|
||||
self.num_ref_frames = args['num_ref_frames']
|
||||
self.size = self.w, self.h = (args['w'], args['h'])
|
||||
|
||||
self.load_flow = args['load_flow']
|
||||
if self.load_flow:
|
||||
assert os.path.exists(self.flow_root)
|
||||
|
||||
json_path = os.path.join('./datasets', args['name'], 'train.json')
|
||||
|
||||
with open(json_path, 'r') as f:
|
||||
self.video_train_dict = json.load(f)
|
||||
self.video_names = sorted(list(self.video_train_dict.keys()))
|
||||
|
||||
# self.video_names = sorted(os.listdir(self.video_root))
|
||||
self.video_dict = {}
|
||||
self.frame_dict = {}
|
||||
|
||||
for v in self.video_names:
|
||||
frame_list = sorted(os.listdir(os.path.join(self.video_root, v)))
|
||||
v_len = len(frame_list)
|
||||
if v_len > self.num_local_frames + self.num_ref_frames:
|
||||
self.video_dict[v] = v_len
|
||||
self.frame_dict[v] = frame_list
|
||||
|
||||
|
||||
self.video_names = list(self.video_dict.keys()) # update names
|
||||
|
||||
self._to_tensors = transforms.Compose([
|
||||
Stack(),
|
||||
ToTorchFormatTensor(),
|
||||
])
|
||||
self.file_client = FileClient('disk')
|
||||
|
||||
def __len__(self):
|
||||
return len(self.video_names)
|
||||
|
||||
def _sample_index(self, length, sample_length, num_ref_frame=3):
|
||||
complete_idx_set = list(range(length))
|
||||
pivot = random.randint(0, length - sample_length)
|
||||
local_idx = complete_idx_set[pivot:pivot + sample_length]
|
||||
remain_idx = list(set(complete_idx_set) - set(local_idx))
|
||||
ref_index = sorted(random.sample(remain_idx, num_ref_frame))
|
||||
|
||||
return local_idx + ref_index
|
||||
|
||||
def __getitem__(self, index):
|
||||
video_name = self.video_names[index]
|
||||
# create masks
|
||||
all_masks = create_random_shape_with_random_motion(
|
||||
self.video_dict[video_name], imageHeight=self.h, imageWidth=self.w)
|
||||
|
||||
# create sample index
|
||||
selected_index = self._sample_index(self.video_dict[video_name],
|
||||
self.num_local_frames,
|
||||
self.num_ref_frames)
|
||||
|
||||
# read video frames
|
||||
frames = []
|
||||
masks = []
|
||||
flows_f, flows_b = [], []
|
||||
for idx in selected_index:
|
||||
frame_list = self.frame_dict[video_name]
|
||||
img_path = os.path.join(self.video_root, video_name, frame_list[idx])
|
||||
img_bytes = self.file_client.get(img_path, 'img')
|
||||
img = imfrombytes(img_bytes, float32=False)
|
||||
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
||||
img = cv2.resize(img, self.size, interpolation=cv2.INTER_LINEAR)
|
||||
img = Image.fromarray(img)
|
||||
|
||||
frames.append(img)
|
||||
masks.append(all_masks[idx])
|
||||
|
||||
if len(frames) <= self.num_local_frames-1 and self.load_flow:
|
||||
current_n = frame_list[idx][:-4]
|
||||
next_n = frame_list[idx+1][:-4]
|
||||
flow_f_path = os.path.join(self.flow_root, video_name, f'{current_n}_{next_n}_f.flo')
|
||||
flow_b_path = os.path.join(self.flow_root, video_name, f'{next_n}_{current_n}_b.flo')
|
||||
flow_f = flowread(flow_f_path, quantize=False)
|
||||
flow_b = flowread(flow_b_path, quantize=False)
|
||||
flow_f = resize_flow(flow_f, self.h, self.w)
|
||||
flow_b = resize_flow(flow_b, self.h, self.w)
|
||||
flows_f.append(flow_f)
|
||||
flows_b.append(flow_b)
|
||||
|
||||
if len(frames) == self.num_local_frames: # random reverse
|
||||
if random.random() < 0.5:
|
||||
frames.reverse()
|
||||
masks.reverse()
|
||||
if self.load_flow:
|
||||
flows_f.reverse()
|
||||
flows_b.reverse()
|
||||
flows_ = flows_f
|
||||
flows_f = flows_b
|
||||
flows_b = flows_
|
||||
|
||||
if self.load_flow:
|
||||
frames, flows_f, flows_b = GroupRandomHorizontalFlowFlip()(frames, flows_f, flows_b)
|
||||
else:
|
||||
frames = GroupRandomHorizontalFlip()(frames)
|
||||
|
||||
# normalizate, to tensors
|
||||
frame_tensors = self._to_tensors(frames) * 2.0 - 1.0
|
||||
mask_tensors = self._to_tensors(masks)
|
||||
if self.load_flow:
|
||||
flows_f = np.stack(flows_f, axis=-1) # H W 2 T-1
|
||||
flows_b = np.stack(flows_b, axis=-1)
|
||||
flows_f = torch.from_numpy(flows_f).permute(3, 2, 0, 1).contiguous().float()
|
||||
flows_b = torch.from_numpy(flows_b).permute(3, 2, 0, 1).contiguous().float()
|
||||
|
||||
# img [-1,1] mask [0,1]
|
||||
if self.load_flow:
|
||||
return frame_tensors, mask_tensors, flows_f, flows_b, video_name
|
||||
else:
|
||||
return frame_tensors, mask_tensors, 'None', 'None', video_name
|
||||
|
||||
|
||||
class TestDataset(torch.utils.data.Dataset):
|
||||
def __init__(self, args):
|
||||
self.args = args
|
||||
self.size = self.w, self.h = args['size']
|
||||
|
||||
self.video_root = args['video_root']
|
||||
self.mask_root = args['mask_root']
|
||||
self.flow_root = args['flow_root']
|
||||
|
||||
self.load_flow = args['load_flow']
|
||||
if self.load_flow:
|
||||
assert os.path.exists(self.flow_root)
|
||||
self.video_names = sorted(os.listdir(self.mask_root))
|
||||
|
||||
self.video_dict = {}
|
||||
self.frame_dict = {}
|
||||
|
||||
for v in self.video_names:
|
||||
frame_list = sorted(os.listdir(os.path.join(self.video_root, v)))
|
||||
v_len = len(frame_list)
|
||||
self.video_dict[v] = v_len
|
||||
self.frame_dict[v] = frame_list
|
||||
|
||||
self._to_tensors = transforms.Compose([
|
||||
Stack(),
|
||||
ToTorchFormatTensor(),
|
||||
])
|
||||
self.file_client = FileClient('disk')
|
||||
|
||||
def __len__(self):
|
||||
return len(self.video_names)
|
||||
|
||||
def __getitem__(self, index):
|
||||
video_name = self.video_names[index]
|
||||
selected_index = list(range(self.video_dict[video_name]))
|
||||
|
||||
# read video frames
|
||||
frames = []
|
||||
masks = []
|
||||
flows_f, flows_b = [], []
|
||||
for idx in selected_index:
|
||||
frame_list = self.frame_dict[video_name]
|
||||
frame_path = os.path.join(self.video_root, video_name, frame_list[idx])
|
||||
|
||||
img_bytes = self.file_client.get(frame_path, 'input')
|
||||
img = imfrombytes(img_bytes, float32=False)
|
||||
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
||||
img = cv2.resize(img, self.size, interpolation=cv2.INTER_LINEAR)
|
||||
img = Image.fromarray(img)
|
||||
|
||||
frames.append(img)
|
||||
|
||||
mask_path = os.path.join(self.mask_root, video_name, str(idx).zfill(5) + '.png')
|
||||
mask = Image.open(mask_path).resize(self.size, Image.NEAREST).convert('L')
|
||||
|
||||
# origin: 0 indicates missing. now: 1 indicates missing
|
||||
mask = np.asarray(mask)
|
||||
m = np.array(mask > 0).astype(np.uint8)
|
||||
|
||||
m = cv2.dilate(m,
|
||||
cv2.getStructuringElement(cv2.MORPH_CROSS, (3, 3)),
|
||||
iterations=4)
|
||||
mask = Image.fromarray(m * 255)
|
||||
masks.append(mask)
|
||||
|
||||
if len(frames) <= len(selected_index)-1 and self.load_flow:
|
||||
current_n = frame_list[idx][:-4]
|
||||
next_n = frame_list[idx+1][:-4]
|
||||
flow_f_path = os.path.join(self.flow_root, video_name, f'{current_n}_{next_n}_f.flo')
|
||||
flow_b_path = os.path.join(self.flow_root, video_name, f'{next_n}_{current_n}_b.flo')
|
||||
flow_f = flowread(flow_f_path, quantize=False)
|
||||
flow_b = flowread(flow_b_path, quantize=False)
|
||||
flow_f = resize_flow(flow_f, self.h, self.w)
|
||||
flow_b = resize_flow(flow_b, self.h, self.w)
|
||||
flows_f.append(flow_f)
|
||||
flows_b.append(flow_b)
|
||||
|
||||
# normalizate, to tensors
|
||||
frames_PIL = [np.array(f).astype(np.uint8) for f in frames]
|
||||
frame_tensors = self._to_tensors(frames) * 2.0 - 1.0
|
||||
mask_tensors = self._to_tensors(masks)
|
||||
|
||||
if self.load_flow:
|
||||
flows_f = np.stack(flows_f, axis=-1) # H W 2 T-1
|
||||
flows_b = np.stack(flows_b, axis=-1)
|
||||
flows_f = torch.from_numpy(flows_f).permute(3, 2, 0, 1).contiguous().float()
|
||||
flows_b = torch.from_numpy(flows_b).permute(3, 2, 0, 1).contiguous().float()
|
||||
|
||||
if self.load_flow:
|
||||
return frame_tensors, mask_tensors, flows_f, flows_b, video_name, frames_PIL
|
||||
else:
|
||||
return frame_tensors, mask_tensors, 'None', 'None', video_name
|
||||
@@ -0,0 +1,47 @@
|
||||
import os
|
||||
import torch
|
||||
|
||||
|
||||
def get_world_size():
|
||||
"""Find OMPI world size without calling mpi functions
|
||||
:rtype: int
|
||||
"""
|
||||
if os.environ.get('PMI_SIZE') is not None:
|
||||
return int(os.environ.get('PMI_SIZE') or 1)
|
||||
elif os.environ.get('OMPI_COMM_WORLD_SIZE') is not None:
|
||||
return int(os.environ.get('OMPI_COMM_WORLD_SIZE') or 1)
|
||||
else:
|
||||
return torch.cuda.device_count()
|
||||
|
||||
|
||||
def get_global_rank():
|
||||
"""Find OMPI world rank without calling mpi functions
|
||||
:rtype: int
|
||||
"""
|
||||
if os.environ.get('PMI_RANK') is not None:
|
||||
return int(os.environ.get('PMI_RANK') or 0)
|
||||
elif os.environ.get('OMPI_COMM_WORLD_RANK') is not None:
|
||||
return int(os.environ.get('OMPI_COMM_WORLD_RANK') or 0)
|
||||
else:
|
||||
return 0
|
||||
|
||||
|
||||
def get_local_rank():
|
||||
"""Find OMPI local rank without calling mpi functions
|
||||
:rtype: int
|
||||
"""
|
||||
if os.environ.get('MPI_LOCALRANKID') is not None:
|
||||
return int(os.environ.get('MPI_LOCALRANKID') or 0)
|
||||
elif os.environ.get('OMPI_COMM_WORLD_LOCAL_RANK') is not None:
|
||||
return int(os.environ.get('OMPI_COMM_WORLD_LOCAL_RANK') or 0)
|
||||
else:
|
||||
return 0
|
||||
|
||||
|
||||
def get_master_ip():
|
||||
if os.environ.get('AZ_BATCH_MASTER_NODE') is not None:
|
||||
return os.environ.get('AZ_BATCH_MASTER_NODE').split(':')[0]
|
||||
elif os.environ.get('AZ_BATCHAI_MPI_MASTER_NODE') is not None:
|
||||
return os.environ.get('AZ_BATCHAI_MPI_MASTER_NODE')
|
||||
else:
|
||||
return "127.0.0.1"
|
||||
@@ -0,0 +1,180 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import lpips
|
||||
from model.vgg_arch import VGGFeatureExtractor
|
||||
|
||||
class PerceptualLoss(nn.Module):
|
||||
"""Perceptual loss with commonly used style loss.
|
||||
|
||||
Args:
|
||||
layer_weights (dict): The weight for each layer of vgg feature.
|
||||
Here is an example: {'conv5_4': 1.}, which means the conv5_4
|
||||
feature layer (before relu5_4) will be extracted with weight
|
||||
1.0 in calculting losses.
|
||||
vgg_type (str): The type of vgg network used as feature extractor.
|
||||
Default: 'vgg19'.
|
||||
use_input_norm (bool): If True, normalize the input image in vgg.
|
||||
Default: True.
|
||||
range_norm (bool): If True, norm images with range [-1, 1] to [0, 1].
|
||||
Default: False.
|
||||
perceptual_weight (float): If `perceptual_weight > 0`, the perceptual
|
||||
loss will be calculated and the loss will multiplied by the
|
||||
weight. Default: 1.0.
|
||||
style_weight (float): If `style_weight > 0`, the style loss will be
|
||||
calculated and the loss will multiplied by the weight.
|
||||
Default: 0.
|
||||
criterion (str): Criterion used for perceptual loss. Default: 'l1'.
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
layer_weights,
|
||||
vgg_type='vgg19',
|
||||
use_input_norm=True,
|
||||
range_norm=False,
|
||||
perceptual_weight=1.0,
|
||||
style_weight=0.,
|
||||
criterion='l1'):
|
||||
super(PerceptualLoss, self).__init__()
|
||||
self.perceptual_weight = perceptual_weight
|
||||
self.style_weight = style_weight
|
||||
self.layer_weights = layer_weights
|
||||
self.vgg = VGGFeatureExtractor(
|
||||
layer_name_list=list(layer_weights.keys()),
|
||||
vgg_type=vgg_type,
|
||||
use_input_norm=use_input_norm,
|
||||
range_norm=range_norm)
|
||||
|
||||
self.criterion_type = criterion
|
||||
if self.criterion_type == 'l1':
|
||||
self.criterion = torch.nn.L1Loss()
|
||||
elif self.criterion_type == 'l2':
|
||||
self.criterion = torch.nn.L2loss()
|
||||
elif self.criterion_type == 'mse':
|
||||
self.criterion = torch.nn.MSELoss(reduction='mean')
|
||||
elif self.criterion_type == 'fro':
|
||||
self.criterion = None
|
||||
else:
|
||||
raise NotImplementedError(f'{criterion} criterion has not been supported.')
|
||||
|
||||
def forward(self, x, gt):
|
||||
"""Forward function.
|
||||
|
||||
Args:
|
||||
x (Tensor): Input tensor with shape (n, c, h, w).
|
||||
gt (Tensor): Ground-truth tensor with shape (n, c, h, w).
|
||||
|
||||
Returns:
|
||||
Tensor: Forward results.
|
||||
"""
|
||||
# extract vgg features
|
||||
x_features = self.vgg(x)
|
||||
gt_features = self.vgg(gt.detach())
|
||||
|
||||
# calculate perceptual loss
|
||||
if self.perceptual_weight > 0:
|
||||
percep_loss = 0
|
||||
for k in x_features.keys():
|
||||
if self.criterion_type == 'fro':
|
||||
percep_loss += torch.norm(x_features[k] - gt_features[k], p='fro') * self.layer_weights[k]
|
||||
else:
|
||||
percep_loss += self.criterion(x_features[k], gt_features[k]) * self.layer_weights[k]
|
||||
percep_loss *= self.perceptual_weight
|
||||
else:
|
||||
percep_loss = None
|
||||
|
||||
# calculate style loss
|
||||
if self.style_weight > 0:
|
||||
style_loss = 0
|
||||
for k in x_features.keys():
|
||||
if self.criterion_type == 'fro':
|
||||
style_loss += torch.norm(
|
||||
self._gram_mat(x_features[k]) - self._gram_mat(gt_features[k]), p='fro') * self.layer_weights[k]
|
||||
else:
|
||||
style_loss += self.criterion(self._gram_mat(x_features[k]), self._gram_mat(
|
||||
gt_features[k])) * self.layer_weights[k]
|
||||
style_loss *= self.style_weight
|
||||
else:
|
||||
style_loss = None
|
||||
|
||||
return percep_loss, style_loss
|
||||
|
||||
def _gram_mat(self, x):
|
||||
"""Calculate Gram matrix.
|
||||
|
||||
Args:
|
||||
x (torch.Tensor): Tensor with shape of (n, c, h, w).
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Gram matrix.
|
||||
"""
|
||||
n, c, h, w = x.size()
|
||||
features = x.view(n, c, w * h)
|
||||
features_t = features.transpose(1, 2)
|
||||
gram = features.bmm(features_t) / (c * h * w)
|
||||
return gram
|
||||
|
||||
class LPIPSLoss(nn.Module):
|
||||
def __init__(self,
|
||||
loss_weight=1.0,
|
||||
use_input_norm=True,
|
||||
range_norm=False,):
|
||||
super(LPIPSLoss, self).__init__()
|
||||
self.perceptual = lpips.LPIPS(net="vgg", spatial=False).eval()
|
||||
self.loss_weight = loss_weight
|
||||
self.use_input_norm = use_input_norm
|
||||
self.range_norm = range_norm
|
||||
|
||||
if self.use_input_norm:
|
||||
# the mean is for image with range [0, 1]
|
||||
self.register_buffer('mean', torch.Tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1))
|
||||
# the std is for image with range [0, 1]
|
||||
self.register_buffer('std', torch.Tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1))
|
||||
|
||||
def forward(self, pred, target):
|
||||
if self.range_norm:
|
||||
pred = (pred + 1) / 2
|
||||
target = (target + 1) / 2
|
||||
if self.use_input_norm:
|
||||
pred = (pred - self.mean) / self.std
|
||||
target = (target - self.mean) / self.std
|
||||
lpips_loss = self.perceptual(target.contiguous(), pred.contiguous())
|
||||
return self.loss_weight * lpips_loss.mean(), None
|
||||
|
||||
|
||||
class AdversarialLoss(nn.Module):
|
||||
r"""
|
||||
Adversarial loss
|
||||
https://arxiv.org/abs/1711.10337
|
||||
"""
|
||||
def __init__(self,
|
||||
type='nsgan',
|
||||
target_real_label=1.0,
|
||||
target_fake_label=0.0):
|
||||
r"""
|
||||
type = nsgan | lsgan | hinge
|
||||
"""
|
||||
super(AdversarialLoss, self).__init__()
|
||||
self.type = type
|
||||
self.register_buffer('real_label', torch.tensor(target_real_label))
|
||||
self.register_buffer('fake_label', torch.tensor(target_fake_label))
|
||||
|
||||
if type == 'nsgan':
|
||||
self.criterion = nn.BCELoss()
|
||||
elif type == 'lsgan':
|
||||
self.criterion = nn.MSELoss()
|
||||
elif type == 'hinge':
|
||||
self.criterion = nn.ReLU()
|
||||
|
||||
def __call__(self, outputs, is_real, is_disc=None):
|
||||
if self.type == 'hinge':
|
||||
if is_disc:
|
||||
if is_real:
|
||||
outputs = -outputs
|
||||
return self.criterion(1 + outputs).mean()
|
||||
else:
|
||||
return (-outputs).mean()
|
||||
else:
|
||||
labels = (self.real_label
|
||||
if is_real else self.fake_label).expand_as(outputs)
|
||||
loss = self.criterion(outputs, labels)
|
||||
return loss
|
||||
@@ -0,0 +1,112 @@
|
||||
"""
|
||||
LR scheduler from BasicSR https://github.com/xinntao/BasicSR
|
||||
"""
|
||||
import math
|
||||
from collections import Counter
|
||||
from torch.optim.lr_scheduler import _LRScheduler
|
||||
|
||||
|
||||
class MultiStepRestartLR(_LRScheduler):
|
||||
""" MultiStep with restarts learning rate scheme.
|
||||
Args:
|
||||
optimizer (torch.nn.optimizer): Torch optimizer.
|
||||
milestones (list): Iterations that will decrease learning rate.
|
||||
gamma (float): Decrease ratio. Default: 0.1.
|
||||
restarts (list): Restart iterations. Default: [0].
|
||||
restart_weights (list): Restart weights at each restart iteration.
|
||||
Default: [1].
|
||||
last_epoch (int): Used in _LRScheduler. Default: -1.
|
||||
"""
|
||||
def __init__(self,
|
||||
optimizer,
|
||||
milestones,
|
||||
gamma=0.1,
|
||||
restarts=(0, ),
|
||||
restart_weights=(1, ),
|
||||
last_epoch=-1):
|
||||
self.milestones = Counter(milestones)
|
||||
self.gamma = gamma
|
||||
self.restarts = restarts
|
||||
self.restart_weights = restart_weights
|
||||
assert len(self.restarts) == len(
|
||||
self.restart_weights), 'restarts and their weights do not match.'
|
||||
super(MultiStepRestartLR, self).__init__(optimizer, last_epoch)
|
||||
|
||||
def get_lr(self):
|
||||
if self.last_epoch in self.restarts:
|
||||
weight = self.restart_weights[self.restarts.index(self.last_epoch)]
|
||||
return [
|
||||
group['initial_lr'] * weight
|
||||
for group in self.optimizer.param_groups
|
||||
]
|
||||
if self.last_epoch not in self.milestones:
|
||||
return [group['lr'] for group in self.optimizer.param_groups]
|
||||
return [
|
||||
group['lr'] * self.gamma**self.milestones[self.last_epoch]
|
||||
for group in self.optimizer.param_groups
|
||||
]
|
||||
|
||||
|
||||
def get_position_from_periods(iteration, cumulative_period):
|
||||
"""Get the position from a period list.
|
||||
It will return the index of the right-closest number in the period list.
|
||||
For example, the cumulative_period = [100, 200, 300, 400],
|
||||
if iteration == 50, return 0;
|
||||
if iteration == 210, return 2;
|
||||
if iteration == 300, return 2.
|
||||
Args:
|
||||
iteration (int): Current iteration.
|
||||
cumulative_period (list[int]): Cumulative period list.
|
||||
Returns:
|
||||
int: The position of the right-closest number in the period list.
|
||||
"""
|
||||
for i, period in enumerate(cumulative_period):
|
||||
if iteration <= period:
|
||||
return i
|
||||
|
||||
|
||||
class CosineAnnealingRestartLR(_LRScheduler):
|
||||
""" Cosine annealing with restarts learning rate scheme.
|
||||
An example of config:
|
||||
periods = [10, 10, 10, 10]
|
||||
restart_weights = [1, 0.5, 0.5, 0.5]
|
||||
eta_min=1e-7
|
||||
It has four cycles, each has 10 iterations. At 10th, 20th, 30th, the
|
||||
scheduler will restart with the weights in restart_weights.
|
||||
Args:
|
||||
optimizer (torch.nn.optimizer): Torch optimizer.
|
||||
periods (list): Period for each cosine anneling cycle.
|
||||
restart_weights (list): Restart weights at each restart iteration.
|
||||
Default: [1].
|
||||
eta_min (float): The mimimum lr. Default: 0.
|
||||
last_epoch (int): Used in _LRScheduler. Default: -1.
|
||||
"""
|
||||
def __init__(self,
|
||||
optimizer,
|
||||
periods,
|
||||
restart_weights=(1, ),
|
||||
eta_min=1e-7,
|
||||
last_epoch=-1):
|
||||
self.periods = periods
|
||||
self.restart_weights = restart_weights
|
||||
self.eta_min = eta_min
|
||||
assert (len(self.periods) == len(self.restart_weights)
|
||||
), 'periods and restart_weights should have the same length.'
|
||||
self.cumulative_period = [
|
||||
sum(self.periods[0:i + 1]) for i in range(0, len(self.periods))
|
||||
]
|
||||
super(CosineAnnealingRestartLR, self).__init__(optimizer, last_epoch)
|
||||
|
||||
def get_lr(self):
|
||||
idx = get_position_from_periods(self.last_epoch,
|
||||
self.cumulative_period)
|
||||
current_weight = self.restart_weights[idx]
|
||||
nearest_restart = 0 if idx == 0 else self.cumulative_period[idx - 1]
|
||||
current_period = self.periods[idx]
|
||||
|
||||
return [
|
||||
self.eta_min + current_weight * 0.5 * (base_lr - self.eta_min) *
|
||||
(1 + math.cos(math.pi * (
|
||||
(self.last_epoch - nearest_restart) / current_period)))
|
||||
for base_lr in self.base_lrs
|
||||
]
|
||||
@@ -0,0 +1,569 @@
|
||||
import numpy as np
|
||||
from skimage import measure
|
||||
from scipy import linalg
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from core.utils import to_tensors
|
||||
|
||||
|
||||
def calculate_epe(flow1, flow2):
|
||||
"""Calculate End point errors."""
|
||||
|
||||
epe = torch.sum((flow1 - flow2)**2, dim=1).sqrt()
|
||||
epe = epe.view(-1)
|
||||
return epe.mean().item()
|
||||
|
||||
|
||||
def calculate_psnr(img1, img2):
|
||||
"""Calculate PSNR (Peak Signal-to-Noise Ratio).
|
||||
Ref: https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio
|
||||
Args:
|
||||
img1 (ndarray): Images with range [0, 255].
|
||||
img2 (ndarray): Images with range [0, 255].
|
||||
Returns:
|
||||
float: psnr result.
|
||||
"""
|
||||
|
||||
assert img1.shape == img2.shape, \
|
||||
(f'Image shapes are differnet: {img1.shape}, {img2.shape}.')
|
||||
|
||||
mse = np.mean((img1 - img2)**2)
|
||||
if mse == 0:
|
||||
return float('inf')
|
||||
return 20. * np.log10(255. / np.sqrt(mse))
|
||||
|
||||
|
||||
def calc_psnr_and_ssim(img1, img2):
|
||||
"""Calculate PSNR and SSIM for images.
|
||||
img1: ndarray, range [0, 255]
|
||||
img2: ndarray, range [0, 255]
|
||||
"""
|
||||
img1 = img1.astype(np.float64)
|
||||
img2 = img2.astype(np.float64)
|
||||
|
||||
psnr = calculate_psnr(img1, img2)
|
||||
ssim = measure.compare_ssim(img1,
|
||||
img2,
|
||||
data_range=255,
|
||||
multichannel=True,
|
||||
win_size=65)
|
||||
|
||||
return psnr, ssim
|
||||
|
||||
|
||||
###########################
|
||||
# I3D models
|
||||
###########################
|
||||
|
||||
|
||||
def init_i3d_model(i3d_model_path):
|
||||
print(f"[Loading I3D model from {i3d_model_path} for FID score ..]")
|
||||
i3d_model = InceptionI3d(400, in_channels=3, final_endpoint='Logits')
|
||||
i3d_model.load_state_dict(torch.load(i3d_model_path))
|
||||
i3d_model.to(torch.device('cuda:0'))
|
||||
return i3d_model
|
||||
|
||||
|
||||
def calculate_i3d_activations(video1, video2, i3d_model, device):
|
||||
"""Calculate VFID metric.
|
||||
video1: list[PIL.Image]
|
||||
video2: list[PIL.Image]
|
||||
"""
|
||||
video1 = to_tensors()(video1).unsqueeze(0).to(device)
|
||||
video2 = to_tensors()(video2).unsqueeze(0).to(device)
|
||||
video1_activations = get_i3d_activations(
|
||||
video1, i3d_model).cpu().numpy().flatten()
|
||||
video2_activations = get_i3d_activations(
|
||||
video2, i3d_model).cpu().numpy().flatten()
|
||||
|
||||
return video1_activations, video2_activations
|
||||
|
||||
|
||||
def calculate_vfid(real_activations, fake_activations):
|
||||
"""
|
||||
Given two distribution of features, compute the FID score between them
|
||||
Params:
|
||||
real_activations: list[ndarray]
|
||||
fake_activations: list[ndarray]
|
||||
"""
|
||||
m1 = np.mean(real_activations, axis=0)
|
||||
m2 = np.mean(fake_activations, axis=0)
|
||||
s1 = np.cov(real_activations, rowvar=False)
|
||||
s2 = np.cov(fake_activations, rowvar=False)
|
||||
return calculate_frechet_distance(m1, s1, m2, s2)
|
||||
|
||||
|
||||
def calculate_frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6):
|
||||
"""Numpy implementation of the Frechet Distance.
|
||||
The Frechet distance between two multivariate Gaussians X_1 ~ N(mu_1, C_1)
|
||||
and X_2 ~ N(mu_2, C_2) is
|
||||
d^2 = ||mu_1 - mu_2||^2 + Tr(C_1 + C_2 - 2*sqrt(C_1*C_2)).
|
||||
Stable version by Dougal J. Sutherland.
|
||||
Params:
|
||||
-- mu1 : Numpy array containing the activations of a layer of the
|
||||
inception net (like returned by the function 'get_predictions')
|
||||
for generated samples.
|
||||
-- mu2 : The sample mean over activations, precalculated on an
|
||||
representive data set.
|
||||
-- sigma1: The covariance matrix over activations for generated samples.
|
||||
-- sigma2: The covariance matrix over activations, precalculated on an
|
||||
representive data set.
|
||||
Returns:
|
||||
-- : The Frechet Distance.
|
||||
"""
|
||||
|
||||
mu1 = np.atleast_1d(mu1)
|
||||
mu2 = np.atleast_1d(mu2)
|
||||
|
||||
sigma1 = np.atleast_2d(sigma1)
|
||||
sigma2 = np.atleast_2d(sigma2)
|
||||
|
||||
assert mu1.shape == mu2.shape, \
|
||||
'Training and test mean vectors have different lengths'
|
||||
assert sigma1.shape == sigma2.shape, \
|
||||
'Training and test covariances have different dimensions'
|
||||
|
||||
diff = mu1 - mu2
|
||||
|
||||
# Product might be almost singular
|
||||
covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)
|
||||
if not np.isfinite(covmean).all():
|
||||
msg = ('fid calculation produces singular product; '
|
||||
'adding %s to diagonal of cov estimates') % eps
|
||||
print(msg)
|
||||
offset = np.eye(sigma1.shape[0]) * eps
|
||||
covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))
|
||||
|
||||
# Numerical error might give slight imaginary component
|
||||
if np.iscomplexobj(covmean):
|
||||
if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):
|
||||
m = np.max(np.abs(covmean.imag))
|
||||
raise ValueError('Imaginary component {}'.format(m))
|
||||
covmean = covmean.real
|
||||
|
||||
tr_covmean = np.trace(covmean)
|
||||
|
||||
return (diff.dot(diff) + np.trace(sigma1) + # NOQA
|
||||
np.trace(sigma2) - 2 * tr_covmean)
|
||||
|
||||
|
||||
def get_i3d_activations(batched_video,
|
||||
i3d_model,
|
||||
target_endpoint='Logits',
|
||||
flatten=True,
|
||||
grad_enabled=False):
|
||||
"""
|
||||
Get features from i3d model and flatten them to 1d feature,
|
||||
valid target endpoints are defined in InceptionI3d.VALID_ENDPOINTS
|
||||
VALID_ENDPOINTS = (
|
||||
'Conv3d_1a_7x7',
|
||||
'MaxPool3d_2a_3x3',
|
||||
'Conv3d_2b_1x1',
|
||||
'Conv3d_2c_3x3',
|
||||
'MaxPool3d_3a_3x3',
|
||||
'Mixed_3b',
|
||||
'Mixed_3c',
|
||||
'MaxPool3d_4a_3x3',
|
||||
'Mixed_4b',
|
||||
'Mixed_4c',
|
||||
'Mixed_4d',
|
||||
'Mixed_4e',
|
||||
'Mixed_4f',
|
||||
'MaxPool3d_5a_2x2',
|
||||
'Mixed_5b',
|
||||
'Mixed_5c',
|
||||
'Logits',
|
||||
'Predictions',
|
||||
)
|
||||
"""
|
||||
with torch.set_grad_enabled(grad_enabled):
|
||||
feat = i3d_model.extract_features(batched_video.transpose(1, 2),
|
||||
target_endpoint)
|
||||
if flatten:
|
||||
feat = feat.view(feat.size(0), -1)
|
||||
|
||||
return feat
|
||||
|
||||
|
||||
# This code is from https://github.com/piergiaj/pytorch-i3d/blob/master/pytorch_i3d.py
|
||||
# I only fix flake8 errors and do some cleaning here
|
||||
|
||||
|
||||
class MaxPool3dSamePadding(nn.MaxPool3d):
|
||||
def compute_pad(self, dim, s):
|
||||
if s % self.stride[dim] == 0:
|
||||
return max(self.kernel_size[dim] - self.stride[dim], 0)
|
||||
else:
|
||||
return max(self.kernel_size[dim] - (s % self.stride[dim]), 0)
|
||||
|
||||
def forward(self, x):
|
||||
# compute 'same' padding
|
||||
(batch, channel, t, h, w) = x.size()
|
||||
pad_t = self.compute_pad(0, t)
|
||||
pad_h = self.compute_pad(1, h)
|
||||
pad_w = self.compute_pad(2, w)
|
||||
|
||||
pad_t_f = pad_t // 2
|
||||
pad_t_b = pad_t - pad_t_f
|
||||
pad_h_f = pad_h // 2
|
||||
pad_h_b = pad_h - pad_h_f
|
||||
pad_w_f = pad_w // 2
|
||||
pad_w_b = pad_w - pad_w_f
|
||||
|
||||
pad = (pad_w_f, pad_w_b, pad_h_f, pad_h_b, pad_t_f, pad_t_b)
|
||||
x = F.pad(x, pad)
|
||||
return super(MaxPool3dSamePadding, self).forward(x)
|
||||
|
||||
|
||||
class Unit3D(nn.Module):
|
||||
def __init__(self,
|
||||
in_channels,
|
||||
output_channels,
|
||||
kernel_shape=(1, 1, 1),
|
||||
stride=(1, 1, 1),
|
||||
padding=0,
|
||||
activation_fn=F.relu,
|
||||
use_batch_norm=True,
|
||||
use_bias=False,
|
||||
name='unit_3d'):
|
||||
"""Initializes Unit3D module."""
|
||||
super(Unit3D, self).__init__()
|
||||
|
||||
self._output_channels = output_channels
|
||||
self._kernel_shape = kernel_shape
|
||||
self._stride = stride
|
||||
self._use_batch_norm = use_batch_norm
|
||||
self._activation_fn = activation_fn
|
||||
self._use_bias = use_bias
|
||||
self.name = name
|
||||
self.padding = padding
|
||||
|
||||
self.conv3d = nn.Conv3d(
|
||||
in_channels=in_channels,
|
||||
out_channels=self._output_channels,
|
||||
kernel_size=self._kernel_shape,
|
||||
stride=self._stride,
|
||||
padding=0, # we always want padding to be 0 here. We will
|
||||
# dynamically pad based on input size in forward function
|
||||
bias=self._use_bias)
|
||||
|
||||
if self._use_batch_norm:
|
||||
self.bn = nn.BatchNorm3d(self._output_channels,
|
||||
eps=0.001,
|
||||
momentum=0.01)
|
||||
|
||||
def compute_pad(self, dim, s):
|
||||
if s % self._stride[dim] == 0:
|
||||
return max(self._kernel_shape[dim] - self._stride[dim], 0)
|
||||
else:
|
||||
return max(self._kernel_shape[dim] - (s % self._stride[dim]), 0)
|
||||
|
||||
def forward(self, x):
|
||||
# compute 'same' padding
|
||||
(batch, channel, t, h, w) = x.size()
|
||||
pad_t = self.compute_pad(0, t)
|
||||
pad_h = self.compute_pad(1, h)
|
||||
pad_w = self.compute_pad(2, w)
|
||||
|
||||
pad_t_f = pad_t // 2
|
||||
pad_t_b = pad_t - pad_t_f
|
||||
pad_h_f = pad_h // 2
|
||||
pad_h_b = pad_h - pad_h_f
|
||||
pad_w_f = pad_w // 2
|
||||
pad_w_b = pad_w - pad_w_f
|
||||
|
||||
pad = (pad_w_f, pad_w_b, pad_h_f, pad_h_b, pad_t_f, pad_t_b)
|
||||
x = F.pad(x, pad)
|
||||
|
||||
x = self.conv3d(x)
|
||||
if self._use_batch_norm:
|
||||
x = self.bn(x)
|
||||
if self._activation_fn is not None:
|
||||
x = self._activation_fn(x)
|
||||
return x
|
||||
|
||||
|
||||
class InceptionModule(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, name):
|
||||
super(InceptionModule, self).__init__()
|
||||
|
||||
self.b0 = Unit3D(in_channels=in_channels,
|
||||
output_channels=out_channels[0],
|
||||
kernel_shape=[1, 1, 1],
|
||||
padding=0,
|
||||
name=name + '/Branch_0/Conv3d_0a_1x1')
|
||||
self.b1a = Unit3D(in_channels=in_channels,
|
||||
output_channels=out_channels[1],
|
||||
kernel_shape=[1, 1, 1],
|
||||
padding=0,
|
||||
name=name + '/Branch_1/Conv3d_0a_1x1')
|
||||
self.b1b = Unit3D(in_channels=out_channels[1],
|
||||
output_channels=out_channels[2],
|
||||
kernel_shape=[3, 3, 3],
|
||||
name=name + '/Branch_1/Conv3d_0b_3x3')
|
||||
self.b2a = Unit3D(in_channels=in_channels,
|
||||
output_channels=out_channels[3],
|
||||
kernel_shape=[1, 1, 1],
|
||||
padding=0,
|
||||
name=name + '/Branch_2/Conv3d_0a_1x1')
|
||||
self.b2b = Unit3D(in_channels=out_channels[3],
|
||||
output_channels=out_channels[4],
|
||||
kernel_shape=[3, 3, 3],
|
||||
name=name + '/Branch_2/Conv3d_0b_3x3')
|
||||
self.b3a = MaxPool3dSamePadding(kernel_size=[3, 3, 3],
|
||||
stride=(1, 1, 1),
|
||||
padding=0)
|
||||
self.b3b = Unit3D(in_channels=in_channels,
|
||||
output_channels=out_channels[5],
|
||||
kernel_shape=[1, 1, 1],
|
||||
padding=0,
|
||||
name=name + '/Branch_3/Conv3d_0b_1x1')
|
||||
self.name = name
|
||||
|
||||
def forward(self, x):
|
||||
b0 = self.b0(x)
|
||||
b1 = self.b1b(self.b1a(x))
|
||||
b2 = self.b2b(self.b2a(x))
|
||||
b3 = self.b3b(self.b3a(x))
|
||||
return torch.cat([b0, b1, b2, b3], dim=1)
|
||||
|
||||
|
||||
class InceptionI3d(nn.Module):
|
||||
"""Inception-v1 I3D architecture.
|
||||
The model is introduced in:
|
||||
Quo Vadis, Action Recognition? A New Model and the Kinetics Dataset
|
||||
Joao Carreira, Andrew Zisserman
|
||||
https://arxiv.org/pdf/1705.07750v1.pdf.
|
||||
See also the Inception architecture, introduced in:
|
||||
Going deeper with convolutions
|
||||
Christian Szegedy, Wei Liu, Yangqing Jia, Pierre Sermanet, Scott Reed,
|
||||
Dragomir Anguelov, Dumitru Erhan, Vincent Vanhoucke, Andrew Rabinovich.
|
||||
http://arxiv.org/pdf/1409.4842v1.pdf.
|
||||
"""
|
||||
|
||||
# Endpoints of the model in order. During construction, all the endpoints up
|
||||
# to a designated `final_endpoint` are returned in a dictionary as the
|
||||
# second return value.
|
||||
VALID_ENDPOINTS = (
|
||||
'Conv3d_1a_7x7',
|
||||
'MaxPool3d_2a_3x3',
|
||||
'Conv3d_2b_1x1',
|
||||
'Conv3d_2c_3x3',
|
||||
'MaxPool3d_3a_3x3',
|
||||
'Mixed_3b',
|
||||
'Mixed_3c',
|
||||
'MaxPool3d_4a_3x3',
|
||||
'Mixed_4b',
|
||||
'Mixed_4c',
|
||||
'Mixed_4d',
|
||||
'Mixed_4e',
|
||||
'Mixed_4f',
|
||||
'MaxPool3d_5a_2x2',
|
||||
'Mixed_5b',
|
||||
'Mixed_5c',
|
||||
'Logits',
|
||||
'Predictions',
|
||||
)
|
||||
|
||||
def __init__(self,
|
||||
num_classes=400,
|
||||
spatial_squeeze=True,
|
||||
final_endpoint='Logits',
|
||||
name='inception_i3d',
|
||||
in_channels=3,
|
||||
dropout_keep_prob=0.5):
|
||||
"""Initializes I3D model instance.
|
||||
Args:
|
||||
num_classes: The number of outputs in the logit layer (default 400, which
|
||||
matches the Kinetics dataset).
|
||||
spatial_squeeze: Whether to squeeze the spatial dimensions for the logits
|
||||
before returning (default True).
|
||||
final_endpoint: The model contains many possible endpoints.
|
||||
`final_endpoint` specifies the last endpoint for the model to be built
|
||||
up to. In addition to the output at `final_endpoint`, all the outputs
|
||||
at endpoints up to `final_endpoint` will also be returned, in a
|
||||
dictionary. `final_endpoint` must be one of
|
||||
InceptionI3d.VALID_ENDPOINTS (default 'Logits').
|
||||
name: A string (optional). The name of this module.
|
||||
Raises:
|
||||
ValueError: if `final_endpoint` is not recognized.
|
||||
"""
|
||||
|
||||
if final_endpoint not in self.VALID_ENDPOINTS:
|
||||
raise ValueError('Unknown final endpoint %s' % final_endpoint)
|
||||
|
||||
super(InceptionI3d, self).__init__()
|
||||
self._num_classes = num_classes
|
||||
self._spatial_squeeze = spatial_squeeze
|
||||
self._final_endpoint = final_endpoint
|
||||
self.logits = None
|
||||
|
||||
if self._final_endpoint not in self.VALID_ENDPOINTS:
|
||||
raise ValueError('Unknown final endpoint %s' %
|
||||
self._final_endpoint)
|
||||
|
||||
self.end_points = {}
|
||||
end_point = 'Conv3d_1a_7x7'
|
||||
self.end_points[end_point] = Unit3D(in_channels=in_channels,
|
||||
output_channels=64,
|
||||
kernel_shape=[7, 7, 7],
|
||||
stride=(2, 2, 2),
|
||||
padding=(3, 3, 3),
|
||||
name=name + end_point)
|
||||
if self._final_endpoint == end_point:
|
||||
return
|
||||
|
||||
end_point = 'MaxPool3d_2a_3x3'
|
||||
self.end_points[end_point] = MaxPool3dSamePadding(
|
||||
kernel_size=[1, 3, 3], stride=(1, 2, 2), padding=0)
|
||||
if self._final_endpoint == end_point:
|
||||
return
|
||||
|
||||
end_point = 'Conv3d_2b_1x1'
|
||||
self.end_points[end_point] = Unit3D(in_channels=64,
|
||||
output_channels=64,
|
||||
kernel_shape=[1, 1, 1],
|
||||
padding=0,
|
||||
name=name + end_point)
|
||||
if self._final_endpoint == end_point:
|
||||
return
|
||||
|
||||
end_point = 'Conv3d_2c_3x3'
|
||||
self.end_points[end_point] = Unit3D(in_channels=64,
|
||||
output_channels=192,
|
||||
kernel_shape=[3, 3, 3],
|
||||
padding=1,
|
||||
name=name + end_point)
|
||||
if self._final_endpoint == end_point:
|
||||
return
|
||||
|
||||
end_point = 'MaxPool3d_3a_3x3'
|
||||
self.end_points[end_point] = MaxPool3dSamePadding(
|
||||
kernel_size=[1, 3, 3], stride=(1, 2, 2), padding=0)
|
||||
if self._final_endpoint == end_point:
|
||||
return
|
||||
|
||||
end_point = 'Mixed_3b'
|
||||
self.end_points[end_point] = InceptionModule(192,
|
||||
[64, 96, 128, 16, 32, 32],
|
||||
name + end_point)
|
||||
if self._final_endpoint == end_point:
|
||||
return
|
||||
|
||||
end_point = 'Mixed_3c'
|
||||
self.end_points[end_point] = InceptionModule(
|
||||
256, [128, 128, 192, 32, 96, 64], name + end_point)
|
||||
if self._final_endpoint == end_point:
|
||||
return
|
||||
|
||||
end_point = 'MaxPool3d_4a_3x3'
|
||||
self.end_points[end_point] = MaxPool3dSamePadding(
|
||||
kernel_size=[3, 3, 3], stride=(2, 2, 2), padding=0)
|
||||
if self._final_endpoint == end_point:
|
||||
return
|
||||
|
||||
end_point = 'Mixed_4b'
|
||||
self.end_points[end_point] = InceptionModule(
|
||||
128 + 192 + 96 + 64, [192, 96, 208, 16, 48, 64], name + end_point)
|
||||
if self._final_endpoint == end_point:
|
||||
return
|
||||
|
||||
end_point = 'Mixed_4c'
|
||||
self.end_points[end_point] = InceptionModule(
|
||||
192 + 208 + 48 + 64, [160, 112, 224, 24, 64, 64], name + end_point)
|
||||
if self._final_endpoint == end_point:
|
||||
return
|
||||
|
||||
end_point = 'Mixed_4d'
|
||||
self.end_points[end_point] = InceptionModule(
|
||||
160 + 224 + 64 + 64, [128, 128, 256, 24, 64, 64], name + end_point)
|
||||
if self._final_endpoint == end_point:
|
||||
return
|
||||
|
||||
end_point = 'Mixed_4e'
|
||||
self.end_points[end_point] = InceptionModule(
|
||||
128 + 256 + 64 + 64, [112, 144, 288, 32, 64, 64], name + end_point)
|
||||
if self._final_endpoint == end_point:
|
||||
return
|
||||
|
||||
end_point = 'Mixed_4f'
|
||||
self.end_points[end_point] = InceptionModule(
|
||||
112 + 288 + 64 + 64, [256, 160, 320, 32, 128, 128],
|
||||
name + end_point)
|
||||
if self._final_endpoint == end_point:
|
||||
return
|
||||
|
||||
end_point = 'MaxPool3d_5a_2x2'
|
||||
self.end_points[end_point] = MaxPool3dSamePadding(
|
||||
kernel_size=[2, 2, 2], stride=(2, 2, 2), padding=0)
|
||||
if self._final_endpoint == end_point:
|
||||
return
|
||||
|
||||
end_point = 'Mixed_5b'
|
||||
self.end_points[end_point] = InceptionModule(
|
||||
256 + 320 + 128 + 128, [256, 160, 320, 32, 128, 128],
|
||||
name + end_point)
|
||||
if self._final_endpoint == end_point:
|
||||
return
|
||||
|
||||
end_point = 'Mixed_5c'
|
||||
self.end_points[end_point] = InceptionModule(
|
||||
256 + 320 + 128 + 128, [384, 192, 384, 48, 128, 128],
|
||||
name + end_point)
|
||||
if self._final_endpoint == end_point:
|
||||
return
|
||||
|
||||
end_point = 'Logits'
|
||||
self.avg_pool = nn.AvgPool3d(kernel_size=[2, 7, 7], stride=(1, 1, 1))
|
||||
self.dropout = nn.Dropout(dropout_keep_prob)
|
||||
self.logits = Unit3D(in_channels=384 + 384 + 128 + 128,
|
||||
output_channels=self._num_classes,
|
||||
kernel_shape=[1, 1, 1],
|
||||
padding=0,
|
||||
activation_fn=None,
|
||||
use_batch_norm=False,
|
||||
use_bias=True,
|
||||
name='logits')
|
||||
|
||||
self.build()
|
||||
|
||||
def replace_logits(self, num_classes):
|
||||
self._num_classes = num_classes
|
||||
self.logits = Unit3D(in_channels=384 + 384 + 128 + 128,
|
||||
output_channels=self._num_classes,
|
||||
kernel_shape=[1, 1, 1],
|
||||
padding=0,
|
||||
activation_fn=None,
|
||||
use_batch_norm=False,
|
||||
use_bias=True,
|
||||
name='logits')
|
||||
|
||||
def build(self):
|
||||
for k in self.end_points.keys():
|
||||
self.add_module(k, self.end_points[k])
|
||||
|
||||
def forward(self, x):
|
||||
for end_point in self.VALID_ENDPOINTS:
|
||||
if end_point in self.end_points:
|
||||
x = self._modules[end_point](
|
||||
x) # use _modules to work with dataparallel
|
||||
|
||||
x = self.logits(self.dropout(self.avg_pool(x)))
|
||||
if self._spatial_squeeze:
|
||||
logits = x.squeeze(3).squeeze(3)
|
||||
# logits is batch X time X classes, which is what we want to work with
|
||||
return logits
|
||||
|
||||
def extract_features(self, x, target_endpoint='Logits'):
|
||||
for end_point in self.VALID_ENDPOINTS:
|
||||
if end_point in self.end_points:
|
||||
x = self._modules[end_point](x)
|
||||
if end_point == target_endpoint:
|
||||
break
|
||||
if target_endpoint == 'Logits':
|
||||
return x.mean(4).mean(3).mean(2)
|
||||
else:
|
||||
return x
|
||||
@@ -0,0 +1,125 @@
|
||||
import queue as Queue
|
||||
import threading
|
||||
import torch
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
|
||||
class PrefetchGenerator(threading.Thread):
|
||||
"""A general prefetch generator.
|
||||
|
||||
Ref:
|
||||
https://stackoverflow.com/questions/7323664/python-generator-pre-fetch
|
||||
|
||||
Args:
|
||||
generator: Python generator.
|
||||
num_prefetch_queue (int): Number of prefetch queue.
|
||||
"""
|
||||
|
||||
def __init__(self, generator, num_prefetch_queue):
|
||||
threading.Thread.__init__(self)
|
||||
self.queue = Queue.Queue(num_prefetch_queue)
|
||||
self.generator = generator
|
||||
self.daemon = True
|
||||
self.start()
|
||||
|
||||
def run(self):
|
||||
for item in self.generator:
|
||||
self.queue.put(item)
|
||||
self.queue.put(None)
|
||||
|
||||
def __next__(self):
|
||||
next_item = self.queue.get()
|
||||
if next_item is None:
|
||||
raise StopIteration
|
||||
return next_item
|
||||
|
||||
def __iter__(self):
|
||||
return self
|
||||
|
||||
|
||||
class PrefetchDataLoader(DataLoader):
|
||||
"""Prefetch version of dataloader.
|
||||
|
||||
Ref:
|
||||
https://github.com/IgorSusmelj/pytorch-styleguide/issues/5#
|
||||
|
||||
TODO:
|
||||
Need to test on single gpu and ddp (multi-gpu). There is a known issue in
|
||||
ddp.
|
||||
|
||||
Args:
|
||||
num_prefetch_queue (int): Number of prefetch queue.
|
||||
kwargs (dict): Other arguments for dataloader.
|
||||
"""
|
||||
|
||||
def __init__(self, num_prefetch_queue, **kwargs):
|
||||
self.num_prefetch_queue = num_prefetch_queue
|
||||
super(PrefetchDataLoader, self).__init__(**kwargs)
|
||||
|
||||
def __iter__(self):
|
||||
return PrefetchGenerator(super().__iter__(), self.num_prefetch_queue)
|
||||
|
||||
|
||||
class CPUPrefetcher():
|
||||
"""CPU prefetcher.
|
||||
|
||||
Args:
|
||||
loader: Dataloader.
|
||||
"""
|
||||
|
||||
def __init__(self, loader):
|
||||
self.ori_loader = loader
|
||||
self.loader = iter(loader)
|
||||
|
||||
def next(self):
|
||||
try:
|
||||
return next(self.loader)
|
||||
except StopIteration:
|
||||
return None
|
||||
|
||||
def reset(self):
|
||||
self.loader = iter(self.ori_loader)
|
||||
|
||||
|
||||
class CUDAPrefetcher():
|
||||
"""CUDA prefetcher.
|
||||
|
||||
Ref:
|
||||
https://github.com/NVIDIA/apex/issues/304#
|
||||
|
||||
It may consums more GPU memory.
|
||||
|
||||
Args:
|
||||
loader: Dataloader.
|
||||
opt (dict): Options.
|
||||
"""
|
||||
|
||||
def __init__(self, loader, opt):
|
||||
self.ori_loader = loader
|
||||
self.loader = iter(loader)
|
||||
self.opt = opt
|
||||
self.stream = torch.cuda.Stream()
|
||||
self.device = torch.device('cuda' if opt['num_gpu'] != 0 else 'cpu')
|
||||
self.preload()
|
||||
|
||||
def preload(self):
|
||||
try:
|
||||
self.batch = next(self.loader) # self.batch is a dict
|
||||
except StopIteration:
|
||||
self.batch = None
|
||||
return None
|
||||
# put tensors to gpu
|
||||
with torch.cuda.stream(self.stream):
|
||||
for k, v in self.batch.items():
|
||||
if torch.is_tensor(v):
|
||||
self.batch[k] = self.batch[k].to(device=self.device, non_blocking=True)
|
||||
|
||||
def next(self):
|
||||
torch.cuda.current_stream().wait_stream(self.stream)
|
||||
batch = self.batch
|
||||
self.preload()
|
||||
return batch
|
||||
|
||||
def reset(self):
|
||||
self.loader = iter(self.ori_loader)
|
||||
self.preload()
|
||||
@@ -0,0 +1,509 @@
|
||||
import os
|
||||
import glob
|
||||
import logging
|
||||
import importlib
|
||||
from tqdm import tqdm
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from core.prefetch_dataloader import PrefetchDataLoader, CPUPrefetcher
|
||||
from torch.utils.data.distributed import DistributedSampler
|
||||
from torch.nn.parallel import DistributedDataParallel as DDP
|
||||
import torchvision
|
||||
from torch.utils.tensorboard import SummaryWriter
|
||||
|
||||
from core.lr_scheduler import MultiStepRestartLR, CosineAnnealingRestartLR
|
||||
from core.loss import AdversarialLoss, PerceptualLoss, LPIPSLoss
|
||||
from core.dataset import TrainDataset
|
||||
|
||||
from model.modules.flow_comp_raft import RAFT_bi, FlowLoss, EdgeLoss
|
||||
from model.recurrent_flow_completion import RecurrentFlowCompleteNet
|
||||
|
||||
from RAFT.utils.flow_viz_pt import flow_to_image
|
||||
|
||||
|
||||
class Trainer:
|
||||
def __init__(self, config):
|
||||
self.config = config
|
||||
self.epoch = 0
|
||||
self.iteration = 0
|
||||
self.num_local_frames = config['train_data_loader']['num_local_frames']
|
||||
self.num_ref_frames = config['train_data_loader']['num_ref_frames']
|
||||
|
||||
# setup data set and data loader
|
||||
self.train_dataset = TrainDataset(config['train_data_loader'])
|
||||
|
||||
self.train_sampler = None
|
||||
self.train_args = config['trainer']
|
||||
if config['distributed']:
|
||||
self.train_sampler = DistributedSampler(
|
||||
self.train_dataset,
|
||||
num_replicas=config['world_size'],
|
||||
rank=config['global_rank'])
|
||||
|
||||
dataloader_args = dict(
|
||||
dataset=self.train_dataset,
|
||||
batch_size=self.train_args['batch_size'] // config['world_size'],
|
||||
shuffle=(self.train_sampler is None),
|
||||
num_workers=self.train_args['num_workers'],
|
||||
sampler=self.train_sampler,
|
||||
drop_last=True)
|
||||
|
||||
self.train_loader = PrefetchDataLoader(self.train_args['num_prefetch_queue'], **dataloader_args)
|
||||
self.prefetcher = CPUPrefetcher(self.train_loader)
|
||||
|
||||
# set loss functions
|
||||
self.adversarial_loss = AdversarialLoss(type=self.config['losses']['GAN_LOSS'])
|
||||
self.adversarial_loss = self.adversarial_loss.to(self.config['device'])
|
||||
self.l1_loss = nn.L1Loss()
|
||||
# self.perc_loss = PerceptualLoss(
|
||||
# layer_weights={'conv3_4': 0.25, 'conv4_4': 0.25, 'conv5_4': 0.5},
|
||||
# use_input_norm=True,
|
||||
# range_norm=True,
|
||||
# criterion='l1'
|
||||
# ).to(self.config['device'])
|
||||
|
||||
if self.config['losses']['perceptual_weight'] > 0:
|
||||
self.perc_loss = LPIPSLoss(use_input_norm=True, range_norm=True).to(self.config['device'])
|
||||
|
||||
# self.flow_comp_loss = FlowCompletionLoss().to(self.config['device'])
|
||||
# self.flow_comp_loss = FlowCompletionLoss(self.config['device'])
|
||||
|
||||
# set raft
|
||||
self.fix_raft = RAFT_bi(device = self.config['device'])
|
||||
self.fix_flow_complete = RecurrentFlowCompleteNet('/mnt/lustre/sczhou/VQGANs/CodeMOVI/experiments_model/recurrent_flow_completion_v5_train_flowcomp_v5/gen_760000.pth')
|
||||
for p in self.fix_flow_complete.parameters():
|
||||
p.requires_grad = False
|
||||
self.fix_flow_complete.to(self.config['device'])
|
||||
self.fix_flow_complete.eval()
|
||||
|
||||
# self.flow_loss = FlowLoss()
|
||||
|
||||
# setup models including generator and discriminator
|
||||
net = importlib.import_module('model.' + config['model']['net'])
|
||||
self.netG = net.InpaintGenerator()
|
||||
# print(self.netG)
|
||||
self.netG = self.netG.to(self.config['device'])
|
||||
if not self.config['model'].get('no_dis', False):
|
||||
if self.config['model'].get('dis_2d', False):
|
||||
self.netD = net.Discriminator_2D(
|
||||
in_channels=3,
|
||||
use_sigmoid=config['losses']['GAN_LOSS'] != 'hinge')
|
||||
else:
|
||||
self.netD = net.Discriminator(
|
||||
in_channels=3,
|
||||
use_sigmoid=config['losses']['GAN_LOSS'] != 'hinge')
|
||||
self.netD = self.netD.to(self.config['device'])
|
||||
|
||||
self.interp_mode = self.config['model']['interp_mode']
|
||||
# setup optimizers and schedulers
|
||||
self.setup_optimizers()
|
||||
self.setup_schedulers()
|
||||
self.load()
|
||||
|
||||
if config['distributed']:
|
||||
self.netG = DDP(self.netG,
|
||||
device_ids=[self.config['local_rank']],
|
||||
output_device=self.config['local_rank'],
|
||||
broadcast_buffers=True,
|
||||
find_unused_parameters=True)
|
||||
if not self.config['model']['no_dis']:
|
||||
self.netD = DDP(self.netD,
|
||||
device_ids=[self.config['local_rank']],
|
||||
output_device=self.config['local_rank'],
|
||||
broadcast_buffers=True,
|
||||
find_unused_parameters=False)
|
||||
|
||||
# set summary writer
|
||||
self.dis_writer = None
|
||||
self.gen_writer = None
|
||||
self.summary = {}
|
||||
if self.config['global_rank'] == 0 or (not config['distributed']):
|
||||
if not self.config['model']['no_dis']:
|
||||
self.dis_writer = SummaryWriter(
|
||||
os.path.join(config['save_dir'], 'dis'))
|
||||
self.gen_writer = SummaryWriter(
|
||||
os.path.join(config['save_dir'], 'gen'))
|
||||
|
||||
def setup_optimizers(self):
|
||||
"""Set up optimizers."""
|
||||
backbone_params = []
|
||||
for name, param in self.netG.named_parameters():
|
||||
if param.requires_grad:
|
||||
backbone_params.append(param)
|
||||
else:
|
||||
print(f'Params {name} will not be optimized.')
|
||||
|
||||
optim_params = [
|
||||
{
|
||||
'params': backbone_params,
|
||||
'lr': self.config['trainer']['lr']
|
||||
},
|
||||
]
|
||||
|
||||
self.optimG = torch.optim.Adam(optim_params,
|
||||
betas=(self.config['trainer']['beta1'],
|
||||
self.config['trainer']['beta2']))
|
||||
|
||||
if not self.config['model']['no_dis']:
|
||||
self.optimD = torch.optim.Adam(
|
||||
self.netD.parameters(),
|
||||
lr=self.config['trainer']['lr'],
|
||||
betas=(self.config['trainer']['beta1'],
|
||||
self.config['trainer']['beta2']))
|
||||
|
||||
def setup_schedulers(self):
|
||||
"""Set up schedulers."""
|
||||
scheduler_opt = self.config['trainer']['scheduler']
|
||||
scheduler_type = scheduler_opt.pop('type')
|
||||
|
||||
if scheduler_type in ['MultiStepLR', 'MultiStepRestartLR']:
|
||||
self.scheG = MultiStepRestartLR(
|
||||
self.optimG,
|
||||
milestones=scheduler_opt['milestones'],
|
||||
gamma=scheduler_opt['gamma'])
|
||||
if not self.config['model']['no_dis']:
|
||||
self.scheD = MultiStepRestartLR(
|
||||
self.optimD,
|
||||
milestones=scheduler_opt['milestones'],
|
||||
gamma=scheduler_opt['gamma'])
|
||||
elif scheduler_type == 'CosineAnnealingRestartLR':
|
||||
self.scheG = CosineAnnealingRestartLR(
|
||||
self.optimG,
|
||||
periods=scheduler_opt['periods'],
|
||||
restart_weights=scheduler_opt['restart_weights'],
|
||||
eta_min=scheduler_opt['eta_min'])
|
||||
if not self.config['model']['no_dis']:
|
||||
self.scheD = CosineAnnealingRestartLR(
|
||||
self.optimD,
|
||||
periods=scheduler_opt['periods'],
|
||||
restart_weights=scheduler_opt['restart_weights'],
|
||||
eta_min=scheduler_opt['eta_min'])
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f'Scheduler {scheduler_type} is not implemented yet.')
|
||||
|
||||
def update_learning_rate(self):
|
||||
"""Update learning rate."""
|
||||
self.scheG.step()
|
||||
if not self.config['model']['no_dis']:
|
||||
self.scheD.step()
|
||||
|
||||
def get_lr(self):
|
||||
"""Get current learning rate."""
|
||||
return self.optimG.param_groups[0]['lr']
|
||||
|
||||
def add_summary(self, writer, name, val):
|
||||
"""Add tensorboard summary."""
|
||||
if name not in self.summary:
|
||||
self.summary[name] = 0
|
||||
self.summary[name] += val
|
||||
n = self.train_args['log_freq']
|
||||
if writer is not None and self.iteration % n == 0:
|
||||
writer.add_scalar(name, self.summary[name] / n, self.iteration)
|
||||
self.summary[name] = 0
|
||||
|
||||
def load(self):
|
||||
"""Load netG (and netD)."""
|
||||
# get the latest checkpoint
|
||||
model_path = self.config['save_dir']
|
||||
# TODO: add resume name
|
||||
if os.path.isfile(os.path.join(model_path, 'latest.ckpt')):
|
||||
latest_epoch = open(os.path.join(model_path, 'latest.ckpt'),
|
||||
'r').read().splitlines()[-1]
|
||||
else:
|
||||
ckpts = [
|
||||
os.path.basename(i).split('.pth')[0]
|
||||
for i in glob.glob(os.path.join(model_path, '*.pth'))
|
||||
]
|
||||
ckpts.sort()
|
||||
latest_epoch = ckpts[-1][4:] if len(ckpts) > 0 else None
|
||||
|
||||
if latest_epoch is not None:
|
||||
gen_path = os.path.join(model_path,
|
||||
f'gen_{int(latest_epoch):06d}.pth')
|
||||
dis_path = os.path.join(model_path,
|
||||
f'dis_{int(latest_epoch):06d}.pth')
|
||||
opt_path = os.path.join(model_path,
|
||||
f'opt_{int(latest_epoch):06d}.pth')
|
||||
|
||||
if self.config['global_rank'] == 0:
|
||||
print(f'Loading model from {gen_path}...')
|
||||
dataG = torch.load(gen_path, map_location=self.config['device'])
|
||||
self.netG.load_state_dict(dataG)
|
||||
if not self.config['model']['no_dis'] and self.config['model']['load_d']:
|
||||
dataD = torch.load(dis_path, map_location=self.config['device'])
|
||||
self.netD.load_state_dict(dataD)
|
||||
|
||||
data_opt = torch.load(opt_path, map_location=self.config['device'])
|
||||
self.optimG.load_state_dict(data_opt['optimG'])
|
||||
# self.scheG.load_state_dict(data_opt['scheG'])
|
||||
if not self.config['model']['no_dis'] and self.config['model']['load_d']:
|
||||
self.optimD.load_state_dict(data_opt['optimD'])
|
||||
# self.scheD.load_state_dict(data_opt['scheD'])
|
||||
self.epoch = data_opt['epoch']
|
||||
self.iteration = data_opt['iteration']
|
||||
else:
|
||||
gen_path = self.config['trainer'].get('gen_path', None)
|
||||
dis_path = self.config['trainer'].get('dis_path', None)
|
||||
opt_path = self.config['trainer'].get('opt_path', None)
|
||||
if gen_path is not None:
|
||||
if self.config['global_rank'] == 0:
|
||||
print(f'Loading Gen-Net from {gen_path}...')
|
||||
dataG = torch.load(gen_path, map_location=self.config['device'])
|
||||
self.netG.load_state_dict(dataG)
|
||||
|
||||
if dis_path is not None and not self.config['model']['no_dis'] and self.config['model']['load_d']:
|
||||
if self.config['global_rank'] == 0:
|
||||
print(f'Loading Dis-Net from {dis_path}...')
|
||||
dataD = torch.load(dis_path, map_location=self.config['device'])
|
||||
self.netD.load_state_dict(dataD)
|
||||
if opt_path is not None:
|
||||
data_opt = torch.load(opt_path, map_location=self.config['device'])
|
||||
self.optimG.load_state_dict(data_opt['optimG'])
|
||||
self.scheG.load_state_dict(data_opt['scheG'])
|
||||
if not self.config['model']['no_dis'] and self.config['model']['load_d']:
|
||||
self.optimD.load_state_dict(data_opt['optimD'])
|
||||
self.scheD.load_state_dict(data_opt['scheD'])
|
||||
else:
|
||||
if self.config['global_rank'] == 0:
|
||||
print('Warnning: There is no trained model found.'
|
||||
'An initialized model will be used.')
|
||||
|
||||
def save(self, it):
|
||||
"""Save parameters every eval_epoch"""
|
||||
if self.config['global_rank'] == 0:
|
||||
# configure path
|
||||
gen_path = os.path.join(self.config['save_dir'],
|
||||
f'gen_{it:06d}.pth')
|
||||
dis_path = os.path.join(self.config['save_dir'],
|
||||
f'dis_{it:06d}.pth')
|
||||
opt_path = os.path.join(self.config['save_dir'],
|
||||
f'opt_{it:06d}.pth')
|
||||
print(f'\nsaving model to {gen_path} ...')
|
||||
|
||||
# remove .module for saving
|
||||
if isinstance(self.netG, torch.nn.DataParallel) or isinstance(self.netG, DDP):
|
||||
netG = self.netG.module
|
||||
if not self.config['model']['no_dis']:
|
||||
netD = self.netD.module
|
||||
else:
|
||||
netG = self.netG
|
||||
if not self.config['model']['no_dis']:
|
||||
netD = self.netD
|
||||
|
||||
# save checkpoints
|
||||
torch.save(netG.state_dict(), gen_path)
|
||||
if not self.config['model']['no_dis']:
|
||||
torch.save(netD.state_dict(), dis_path)
|
||||
torch.save(
|
||||
{
|
||||
'epoch': self.epoch,
|
||||
'iteration': self.iteration,
|
||||
'optimG': self.optimG.state_dict(),
|
||||
'optimD': self.optimD.state_dict(),
|
||||
'scheG': self.scheG.state_dict(),
|
||||
'scheD': self.scheD.state_dict()
|
||||
}, opt_path)
|
||||
else:
|
||||
torch.save(
|
||||
{
|
||||
'epoch': self.epoch,
|
||||
'iteration': self.iteration,
|
||||
'optimG': self.optimG.state_dict(),
|
||||
'scheG': self.scheG.state_dict()
|
||||
}, opt_path)
|
||||
|
||||
latest_path = os.path.join(self.config['save_dir'], 'latest.ckpt')
|
||||
os.system(f"echo {it:06d} > {latest_path}")
|
||||
|
||||
def train(self):
|
||||
"""training entry"""
|
||||
pbar = range(int(self.train_args['iterations']))
|
||||
if self.config['global_rank'] == 0:
|
||||
pbar = tqdm(pbar,
|
||||
initial=self.iteration,
|
||||
dynamic_ncols=True,
|
||||
smoothing=0.01)
|
||||
|
||||
os.makedirs('logs', exist_ok=True)
|
||||
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s %(filename)s[line:%(lineno)d]"
|
||||
"%(levelname)s %(message)s",
|
||||
datefmt="%a, %d %b %Y %H:%M:%S",
|
||||
filename=f"logs/{self.config['save_dir'].split('/')[-1]}.log",
|
||||
filemode='w')
|
||||
|
||||
while True:
|
||||
self.epoch += 1
|
||||
self.prefetcher.reset()
|
||||
if self.config['distributed']:
|
||||
self.train_sampler.set_epoch(self.epoch)
|
||||
self._train_epoch(pbar)
|
||||
if self.iteration > self.train_args['iterations']:
|
||||
break
|
||||
print('\nEnd training....')
|
||||
|
||||
def _train_epoch(self, pbar):
|
||||
"""Process input and calculate loss every training epoch"""
|
||||
device = self.config['device']
|
||||
train_data = self.prefetcher.next()
|
||||
while train_data is not None:
|
||||
self.iteration += 1
|
||||
frames, masks, flows_f, flows_b, _ = train_data
|
||||
frames, masks = frames.to(device), masks.to(device).float()
|
||||
l_t = self.num_local_frames
|
||||
b, t, c, h, w = frames.size()
|
||||
gt_local_frames = frames[:, :l_t, ...]
|
||||
local_masks = masks[:, :l_t, ...].contiguous()
|
||||
|
||||
masked_frames = frames * (1 - masks)
|
||||
masked_local_frames = masked_frames[:, :l_t, ...]
|
||||
# get gt optical flow
|
||||
if flows_f[0] == 'None' or flows_b[0] == 'None':
|
||||
gt_flows_bi = self.fix_raft(gt_local_frames)
|
||||
else:
|
||||
gt_flows_bi = (flows_f.to(device), flows_b.to(device))
|
||||
|
||||
# ---- complete flow ----
|
||||
pred_flows_bi, _ = self.fix_flow_complete.forward_bidirect_flow(gt_flows_bi, local_masks)
|
||||
pred_flows_bi = self.fix_flow_complete.combine_flow(gt_flows_bi, pred_flows_bi, local_masks)
|
||||
# pred_flows_bi = gt_flows_bi
|
||||
|
||||
# ---- image propagation ----
|
||||
prop_imgs, updated_local_masks = self.netG.module.img_propagation(masked_local_frames, pred_flows_bi, local_masks, interpolation=self.interp_mode)
|
||||
updated_masks = masks.clone()
|
||||
updated_masks[:, :l_t, ...] = updated_local_masks.view(b, l_t, 1, h, w)
|
||||
updated_frames = masked_frames.clone()
|
||||
prop_local_frames = gt_local_frames * (1-local_masks) + prop_imgs.view(b, l_t, 3, h, w) * local_masks # merge
|
||||
updated_frames[:, :l_t, ...] = prop_local_frames
|
||||
|
||||
# ---- feature propagation + Transformer ----
|
||||
pred_imgs = self.netG(updated_frames, pred_flows_bi, masks, updated_masks, l_t)
|
||||
pred_imgs = pred_imgs.view(b, -1, c, h, w)
|
||||
|
||||
# get the local frames
|
||||
pred_local_frames = pred_imgs[:, :l_t, ...]
|
||||
comp_local_frames = gt_local_frames * (1. - local_masks) + pred_local_frames * local_masks
|
||||
comp_imgs = frames * (1. - masks) + pred_imgs * masks
|
||||
|
||||
gen_loss = 0
|
||||
dis_loss = 0
|
||||
# optimize net_g
|
||||
if not self.config['model']['no_dis']:
|
||||
for p in self.netD.parameters():
|
||||
p.requires_grad = False
|
||||
|
||||
self.optimG.zero_grad()
|
||||
|
||||
# generator l1 loss
|
||||
hole_loss = self.l1_loss(pred_imgs * masks, frames * masks)
|
||||
hole_loss = hole_loss / torch.mean(masks) * self.config['losses']['hole_weight']
|
||||
gen_loss += hole_loss
|
||||
self.add_summary(self.gen_writer, 'loss/hole_loss', hole_loss.item())
|
||||
|
||||
valid_loss = self.l1_loss(pred_imgs * (1 - masks), frames * (1 - masks))
|
||||
valid_loss = valid_loss / torch.mean(1-masks) * self.config['losses']['valid_weight']
|
||||
gen_loss += valid_loss
|
||||
self.add_summary(self.gen_writer, 'loss/valid_loss', valid_loss.item())
|
||||
|
||||
# perceptual loss
|
||||
if self.config['losses']['perceptual_weight'] > 0:
|
||||
perc_loss = self.perc_loss(pred_imgs.view(-1,3,h,w), frames.view(-1,3,h,w))[0] * self.config['losses']['perceptual_weight']
|
||||
gen_loss += perc_loss
|
||||
self.add_summary(self.gen_writer, 'loss/perc_loss', perc_loss.item())
|
||||
|
||||
# gan loss
|
||||
if not self.config['model']['no_dis']:
|
||||
# generator adversarial loss
|
||||
gen_clip = self.netD(comp_imgs)
|
||||
gan_loss = self.adversarial_loss(gen_clip, True, False)
|
||||
gan_loss = gan_loss * self.config['losses']['adversarial_weight']
|
||||
gen_loss += gan_loss
|
||||
self.add_summary(self.gen_writer, 'loss/gan_loss', gan_loss.item())
|
||||
gen_loss.backward()
|
||||
self.optimG.step()
|
||||
|
||||
if not self.config['model']['no_dis']:
|
||||
# optimize net_d
|
||||
for p in self.netD.parameters():
|
||||
p.requires_grad = True
|
||||
self.optimD.zero_grad()
|
||||
|
||||
# discriminator adversarial loss
|
||||
real_clip = self.netD(frames)
|
||||
fake_clip = self.netD(comp_imgs.detach())
|
||||
dis_real_loss = self.adversarial_loss(real_clip, True, True)
|
||||
dis_fake_loss = self.adversarial_loss(fake_clip, False, True)
|
||||
dis_loss += (dis_real_loss + dis_fake_loss) / 2
|
||||
self.add_summary(self.dis_writer, 'loss/dis_vid_real', dis_real_loss.item())
|
||||
self.add_summary(self.dis_writer, 'loss/dis_vid_fake', dis_fake_loss.item())
|
||||
dis_loss.backward()
|
||||
self.optimD.step()
|
||||
|
||||
self.update_learning_rate()
|
||||
|
||||
# write image to tensorboard
|
||||
if self.iteration % 200 == 0:
|
||||
# img to cpu
|
||||
t = 0
|
||||
gt_local_frames_cpu = ((gt_local_frames.view(b,-1,3,h,w) + 1)/2.0).cpu()
|
||||
masked_local_frames = ((masked_local_frames.view(b,-1,3,h,w) + 1)/2.0).cpu()
|
||||
prop_local_frames_cpu = ((prop_local_frames.view(b,-1,3,h,w) + 1)/2.0).cpu()
|
||||
pred_local_frames_cpu = ((pred_local_frames.view(b,-1,3,h,w) + 1)/2.0).cpu()
|
||||
img_results = torch.cat([masked_local_frames[0][t], gt_local_frames_cpu[0][t],
|
||||
prop_local_frames_cpu[0][t], pred_local_frames_cpu[0][t]], 1)
|
||||
img_results = torchvision.utils.make_grid(img_results, nrow=1, normalize=True)
|
||||
if self.gen_writer is not None:
|
||||
self.gen_writer.add_image(f'img/img:inp-gt-res-{t}', img_results, self.iteration)
|
||||
|
||||
t = 5
|
||||
if masked_local_frames.shape[1] > 5:
|
||||
img_results = torch.cat([masked_local_frames[0][t], gt_local_frames_cpu[0][t],
|
||||
prop_local_frames_cpu[0][t], pred_local_frames_cpu[0][t]], 1)
|
||||
img_results = torchvision.utils.make_grid(img_results, nrow=1, normalize=True)
|
||||
if self.gen_writer is not None:
|
||||
self.gen_writer.add_image(f'img/img:inp-gt-res-{t}', img_results, self.iteration)
|
||||
|
||||
# flow to cpu
|
||||
gt_flows_forward_cpu = flow_to_image(gt_flows_bi[0][0]).cpu()
|
||||
masked_flows_forward_cpu = (gt_flows_forward_cpu[0] * (1-local_masks[0][0].cpu())).to(gt_flows_forward_cpu)
|
||||
pred_flows_forward_cpu = flow_to_image(pred_flows_bi[0][0]).cpu()
|
||||
|
||||
flow_results = torch.cat([gt_flows_forward_cpu[0], masked_flows_forward_cpu, pred_flows_forward_cpu[0]], 1)
|
||||
if self.gen_writer is not None:
|
||||
self.gen_writer.add_image('img/flow:gt-pred', flow_results, self.iteration)
|
||||
|
||||
# console logs
|
||||
if self.config['global_rank'] == 0:
|
||||
pbar.update(1)
|
||||
if not self.config['model']['no_dis']:
|
||||
pbar.set_description((f"d: {dis_loss.item():.3f}; "
|
||||
f"hole: {hole_loss.item():.3f}; "
|
||||
f"valid: {valid_loss.item():.3f}"))
|
||||
else:
|
||||
pbar.set_description((f"hole: {hole_loss.item():.3f}; "
|
||||
f"valid: {valid_loss.item():.3f}"))
|
||||
|
||||
if self.iteration % self.train_args['log_freq'] == 0:
|
||||
if not self.config['model']['no_dis']:
|
||||
logging.info(f"[Iter {self.iteration}] "
|
||||
f"d: {dis_loss.item():.4f}; "
|
||||
f"hole: {hole_loss.item():.4f}; "
|
||||
f"valid: {valid_loss.item():.4f}")
|
||||
else:
|
||||
logging.info(f"[Iter {self.iteration}] "
|
||||
f"hole: {hole_loss.item():.4f}; "
|
||||
f"valid: {valid_loss.item():.4f}")
|
||||
|
||||
# saving models
|
||||
if self.iteration % self.train_args['save_freq'] == 0:
|
||||
self.save(int(self.iteration))
|
||||
|
||||
if self.iteration > self.train_args['iterations']:
|
||||
break
|
||||
|
||||
train_data = self.prefetcher.next()
|
||||
@@ -0,0 +1,380 @@
|
||||
import os
|
||||
import glob
|
||||
import logging
|
||||
import importlib
|
||||
from tqdm import tqdm
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from core.prefetch_dataloader import PrefetchDataLoader, CPUPrefetcher
|
||||
from torch.utils.data.distributed import DistributedSampler
|
||||
from torch.nn.parallel import DistributedDataParallel as DDP
|
||||
|
||||
from torch.utils.tensorboard import SummaryWriter
|
||||
|
||||
from core.lr_scheduler import MultiStepRestartLR, CosineAnnealingRestartLR
|
||||
from core.dataset import TrainDataset
|
||||
|
||||
from model.modules.flow_comp_raft import RAFT_bi, FlowLoss, EdgeLoss
|
||||
|
||||
# from skimage.feature import canny
|
||||
from model.canny.canny_filter import Canny
|
||||
from RAFT.utils.flow_viz_pt import flow_to_image
|
||||
|
||||
|
||||
class Trainer:
|
||||
def __init__(self, config):
|
||||
self.config = config
|
||||
self.epoch = 0
|
||||
self.iteration = 0
|
||||
self.num_local_frames = config['train_data_loader']['num_local_frames']
|
||||
self.num_ref_frames = config['train_data_loader']['num_ref_frames']
|
||||
|
||||
# setup data set and data loader
|
||||
self.train_dataset = TrainDataset(config['train_data_loader'])
|
||||
|
||||
self.train_sampler = None
|
||||
self.train_args = config['trainer']
|
||||
if config['distributed']:
|
||||
self.train_sampler = DistributedSampler(
|
||||
self.train_dataset,
|
||||
num_replicas=config['world_size'],
|
||||
rank=config['global_rank'])
|
||||
|
||||
dataloader_args = dict(
|
||||
dataset=self.train_dataset,
|
||||
batch_size=self.train_args['batch_size'] // config['world_size'],
|
||||
shuffle=(self.train_sampler is None),
|
||||
num_workers=self.train_args['num_workers'],
|
||||
sampler=self.train_sampler,
|
||||
drop_last=True)
|
||||
|
||||
self.train_loader = PrefetchDataLoader(self.train_args['num_prefetch_queue'], **dataloader_args)
|
||||
self.prefetcher = CPUPrefetcher(self.train_loader)
|
||||
|
||||
# set raft
|
||||
self.fix_raft = RAFT_bi(device = self.config['device'])
|
||||
self.flow_loss = FlowLoss()
|
||||
self.edge_loss = EdgeLoss()
|
||||
self.canny = Canny(sigma=(2,2), low_threshold=0.1, high_threshold=0.2)
|
||||
|
||||
# setup models including generator and discriminator
|
||||
net = importlib.import_module('model.' + config['model']['net'])
|
||||
self.netG = net.RecurrentFlowCompleteNet()
|
||||
# print(self.netG)
|
||||
self.netG = self.netG.to(self.config['device'])
|
||||
|
||||
# setup optimizers and schedulers
|
||||
self.setup_optimizers()
|
||||
self.setup_schedulers()
|
||||
self.load()
|
||||
|
||||
if config['distributed']:
|
||||
self.netG = DDP(self.netG,
|
||||
device_ids=[self.config['local_rank']],
|
||||
output_device=self.config['local_rank'],
|
||||
broadcast_buffers=True,
|
||||
find_unused_parameters=True)
|
||||
|
||||
# set summary writer
|
||||
self.dis_writer = None
|
||||
self.gen_writer = None
|
||||
self.summary = {}
|
||||
if self.config['global_rank'] == 0 or (not config['distributed']):
|
||||
self.gen_writer = SummaryWriter(
|
||||
os.path.join(config['save_dir'], 'gen'))
|
||||
|
||||
def setup_optimizers(self):
|
||||
"""Set up optimizers."""
|
||||
backbone_params = []
|
||||
for name, param in self.netG.named_parameters():
|
||||
if param.requires_grad:
|
||||
backbone_params.append(param)
|
||||
else:
|
||||
print(f'Params {name} will not be optimized.')
|
||||
|
||||
optim_params = [
|
||||
{
|
||||
'params': backbone_params,
|
||||
'lr': self.config['trainer']['lr']
|
||||
},
|
||||
]
|
||||
|
||||
self.optimG = torch.optim.Adam(optim_params,
|
||||
betas=(self.config['trainer']['beta1'],
|
||||
self.config['trainer']['beta2']))
|
||||
|
||||
|
||||
def setup_schedulers(self):
|
||||
"""Set up schedulers."""
|
||||
scheduler_opt = self.config['trainer']['scheduler']
|
||||
scheduler_type = scheduler_opt.pop('type')
|
||||
|
||||
if scheduler_type in ['MultiStepLR', 'MultiStepRestartLR']:
|
||||
self.scheG = MultiStepRestartLR(
|
||||
self.optimG,
|
||||
milestones=scheduler_opt['milestones'],
|
||||
gamma=scheduler_opt['gamma'])
|
||||
elif scheduler_type == 'CosineAnnealingRestartLR':
|
||||
self.scheG = CosineAnnealingRestartLR(
|
||||
self.optimG,
|
||||
periods=scheduler_opt['periods'],
|
||||
restart_weights=scheduler_opt['restart_weights'])
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f'Scheduler {scheduler_type} is not implemented yet.')
|
||||
|
||||
def update_learning_rate(self):
|
||||
"""Update learning rate."""
|
||||
self.scheG.step()
|
||||
|
||||
def get_lr(self):
|
||||
"""Get current learning rate."""
|
||||
return self.optimG.param_groups[0]['lr']
|
||||
|
||||
def add_summary(self, writer, name, val):
|
||||
"""Add tensorboard summary."""
|
||||
if name not in self.summary:
|
||||
self.summary[name] = 0
|
||||
self.summary[name] += val
|
||||
n = self.train_args['log_freq']
|
||||
if writer is not None and self.iteration % n == 0:
|
||||
writer.add_scalar(name, self.summary[name] / n, self.iteration)
|
||||
self.summary[name] = 0
|
||||
|
||||
def load(self):
|
||||
"""Load netG."""
|
||||
# get the latest checkpoint
|
||||
model_path = self.config['save_dir']
|
||||
if os.path.isfile(os.path.join(model_path, 'latest.ckpt')):
|
||||
latest_epoch = open(os.path.join(model_path, 'latest.ckpt'),
|
||||
'r').read().splitlines()[-1]
|
||||
else:
|
||||
ckpts = [
|
||||
os.path.basename(i).split('.pth')[0]
|
||||
for i in glob.glob(os.path.join(model_path, '*.pth'))
|
||||
]
|
||||
ckpts.sort()
|
||||
latest_epoch = ckpts[-1][4:] if len(ckpts) > 0 else None
|
||||
|
||||
if latest_epoch is not None:
|
||||
gen_path = os.path.join(model_path, f'gen_{int(latest_epoch):06d}.pth')
|
||||
opt_path = os.path.join(model_path,f'opt_{int(latest_epoch):06d}.pth')
|
||||
|
||||
if self.config['global_rank'] == 0:
|
||||
print(f'Loading model from {gen_path}...')
|
||||
dataG = torch.load(gen_path, map_location=self.config['device'])
|
||||
self.netG.load_state_dict(dataG)
|
||||
|
||||
|
||||
data_opt = torch.load(opt_path, map_location=self.config['device'])
|
||||
self.optimG.load_state_dict(data_opt['optimG'])
|
||||
self.scheG.load_state_dict(data_opt['scheG'])
|
||||
|
||||
self.epoch = data_opt['epoch']
|
||||
self.iteration = data_opt['iteration']
|
||||
|
||||
else:
|
||||
if self.config['global_rank'] == 0:
|
||||
print('Warnning: There is no trained model found.'
|
||||
'An initialized model will be used.')
|
||||
|
||||
def save(self, it):
|
||||
"""Save parameters every eval_epoch"""
|
||||
if self.config['global_rank'] == 0:
|
||||
# configure path
|
||||
gen_path = os.path.join(self.config['save_dir'],
|
||||
f'gen_{it:06d}.pth')
|
||||
opt_path = os.path.join(self.config['save_dir'],
|
||||
f'opt_{it:06d}.pth')
|
||||
print(f'\nsaving model to {gen_path} ...')
|
||||
|
||||
# remove .module for saving
|
||||
if isinstance(self.netG, torch.nn.DataParallel) or isinstance(self.netG, DDP):
|
||||
netG = self.netG.module
|
||||
else:
|
||||
netG = self.netG
|
||||
|
||||
# save checkpoints
|
||||
torch.save(netG.state_dict(), gen_path)
|
||||
torch.save(
|
||||
{
|
||||
'epoch': self.epoch,
|
||||
'iteration': self.iteration,
|
||||
'optimG': self.optimG.state_dict(),
|
||||
'scheG': self.scheG.state_dict()
|
||||
}, opt_path)
|
||||
|
||||
latest_path = os.path.join(self.config['save_dir'], 'latest.ckpt')
|
||||
os.system(f"echo {it:06d} > {latest_path}")
|
||||
|
||||
def train(self):
|
||||
"""training entry"""
|
||||
pbar = range(int(self.train_args['iterations']))
|
||||
if self.config['global_rank'] == 0:
|
||||
pbar = tqdm(pbar,
|
||||
initial=self.iteration,
|
||||
dynamic_ncols=True,
|
||||
smoothing=0.01)
|
||||
|
||||
os.makedirs('logs', exist_ok=True)
|
||||
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s %(filename)s[line:%(lineno)d]"
|
||||
"%(levelname)s %(message)s",
|
||||
datefmt="%a, %d %b %Y %H:%M:%S",
|
||||
filename=f"logs/{self.config['save_dir'].split('/')[-1]}.log",
|
||||
filemode='w')
|
||||
|
||||
while True:
|
||||
self.epoch += 1
|
||||
self.prefetcher.reset()
|
||||
if self.config['distributed']:
|
||||
self.train_sampler.set_epoch(self.epoch)
|
||||
self._train_epoch(pbar)
|
||||
if self.iteration > self.train_args['iterations']:
|
||||
break
|
||||
print('\nEnd training....')
|
||||
|
||||
# def get_edges(self, flows): # fgvc
|
||||
# # (b, t, 2, H, W)
|
||||
# b, t, _, h, w = flows.shape
|
||||
# flows = flows.view(-1, 2, h, w)
|
||||
# flows_list = flows.permute(0, 2, 3, 1).cpu().numpy()
|
||||
# edges = []
|
||||
# for f in list(flows_list):
|
||||
# flows_gray = (f[:, :, 0] ** 2 + f[:, :, 1] ** 2) ** 0.5
|
||||
# if flows_gray.max() < 1:
|
||||
# flows_gray = flows_gray*0
|
||||
# else:
|
||||
# flows_gray = flows_gray / flows_gray.max()
|
||||
|
||||
# edge = canny(flows_gray, sigma=2, low_threshold=0.1, high_threshold=0.2) # fgvc
|
||||
# edge = torch.from_numpy(edge).view(1, 1, h, w).float()
|
||||
# edges.append(edge)
|
||||
# edges = torch.stack(edges, dim=0).to(self.config['device'])
|
||||
# edges = edges.view(b, t, 1, h, w)
|
||||
# return edges
|
||||
|
||||
def get_edges(self, flows):
|
||||
# (b, t, 2, H, W)
|
||||
b, t, _, h, w = flows.shape
|
||||
flows = flows.view(-1, 2, h, w)
|
||||
flows_gray = (flows[:, 0, None] ** 2 + flows[:, 1, None] ** 2) ** 0.5
|
||||
if flows_gray.max() < 1:
|
||||
flows_gray = flows_gray*0
|
||||
else:
|
||||
flows_gray = flows_gray / flows_gray.max()
|
||||
|
||||
magnitude, edges = self.canny(flows_gray.float())
|
||||
edges = edges.view(b, t, 1, h, w)
|
||||
return edges
|
||||
|
||||
def _train_epoch(self, pbar):
|
||||
"""Process input and calculate loss every training epoch"""
|
||||
device = self.config['device']
|
||||
train_data = self.prefetcher.next()
|
||||
while train_data is not None:
|
||||
self.iteration += 1
|
||||
frames, masks, flows_f, flows_b, _ = train_data
|
||||
frames, masks = frames.to(device), masks.to(device)
|
||||
masks = masks.float()
|
||||
|
||||
l_t = self.num_local_frames
|
||||
b, t, c, h, w = frames.size()
|
||||
gt_local_frames = frames[:, :l_t, ...]
|
||||
local_masks = masks[:, :l_t, ...].contiguous()
|
||||
|
||||
# get gt optical flow
|
||||
if flows_f[0] == 'None' or flows_b[0] == 'None':
|
||||
gt_flows_bi = self.fix_raft(gt_local_frames)
|
||||
else:
|
||||
gt_flows_bi = (flows_f.to(device), flows_b.to(device))
|
||||
|
||||
# get gt edge
|
||||
gt_edges_forward = self.get_edges(gt_flows_bi[0])
|
||||
gt_edges_backward = self.get_edges(gt_flows_bi[1])
|
||||
gt_edges_bi = [gt_edges_forward, gt_edges_backward]
|
||||
|
||||
# complete flow
|
||||
pred_flows_bi, pred_edges_bi = self.netG.module.forward_bidirect_flow(gt_flows_bi, local_masks)
|
||||
|
||||
# optimize net_g
|
||||
self.optimG.zero_grad()
|
||||
|
||||
# compulte flow_loss
|
||||
flow_loss, warp_loss = self.flow_loss(pred_flows_bi, gt_flows_bi, local_masks, gt_local_frames)
|
||||
flow_loss = flow_loss * self.config['losses']['flow_weight']
|
||||
warp_loss = warp_loss * 0.01
|
||||
self.add_summary(self.gen_writer, 'loss/flow_loss', flow_loss.item())
|
||||
self.add_summary(self.gen_writer, 'loss/warp_loss', warp_loss.item())
|
||||
|
||||
# compute edge loss
|
||||
edge_loss = self.edge_loss(pred_edges_bi, gt_edges_bi, local_masks)
|
||||
edge_loss = edge_loss*1.0
|
||||
self.add_summary(self.gen_writer, 'loss/edge_loss', edge_loss.item())
|
||||
|
||||
loss = flow_loss + warp_loss + edge_loss
|
||||
loss.backward()
|
||||
self.optimG.step()
|
||||
self.update_learning_rate()
|
||||
|
||||
# write image to tensorboard
|
||||
# if self.iteration % 200 == 0:
|
||||
if self.iteration % 200 == 0 and self.gen_writer is not None:
|
||||
t = 5
|
||||
# forward to cpu
|
||||
gt_flows_forward_cpu = flow_to_image(gt_flows_bi[0][0]).cpu()
|
||||
masked_flows_forward_cpu = (gt_flows_forward_cpu[t] * (1-local_masks[0][t].cpu())).to(gt_flows_forward_cpu)
|
||||
pred_flows_forward_cpu = flow_to_image(pred_flows_bi[0][0]).cpu()
|
||||
|
||||
flow_results = torch.cat([gt_flows_forward_cpu[t], masked_flows_forward_cpu, pred_flows_forward_cpu[t]], 1)
|
||||
self.gen_writer.add_image('img/flow-f:gt-pred', flow_results, self.iteration)
|
||||
|
||||
# backward to cpu
|
||||
gt_flows_backward_cpu = flow_to_image(gt_flows_bi[1][0]).cpu()
|
||||
masked_flows_backward_cpu = (gt_flows_backward_cpu[t] * (1-local_masks[0][t+1].cpu())).to(gt_flows_backward_cpu)
|
||||
pred_flows_backward_cpu = flow_to_image(pred_flows_bi[1][0]).cpu()
|
||||
|
||||
flow_results = torch.cat([gt_flows_backward_cpu[t], masked_flows_backward_cpu, pred_flows_backward_cpu[t]], 1)
|
||||
self.gen_writer.add_image('img/flow-b:gt-pred', flow_results, self.iteration)
|
||||
|
||||
# TODO: show edge
|
||||
# forward
|
||||
gt_edges_forward_cpu = gt_edges_bi[0][0].cpu()
|
||||
masked_edges_forward_cpu = (gt_edges_forward_cpu[t] * (1-local_masks[0][t].cpu())).to(gt_edges_forward_cpu)
|
||||
pred_edges_forward_cpu = pred_edges_bi[0][0].cpu()
|
||||
|
||||
edge_results = torch.cat([gt_edges_forward_cpu[t], masked_edges_forward_cpu, pred_edges_forward_cpu[t]], 1)
|
||||
self.gen_writer.add_image('img/edge-f:gt-pred', edge_results, self.iteration)
|
||||
# backward
|
||||
gt_edges_backward_cpu = gt_edges_bi[1][0].cpu()
|
||||
masked_edges_backward_cpu = (gt_edges_backward_cpu[t] * (1-local_masks[0][t+1].cpu())).to(gt_edges_backward_cpu)
|
||||
pred_edges_backward_cpu = pred_edges_bi[1][0].cpu()
|
||||
|
||||
edge_results = torch.cat([gt_edges_backward_cpu[t], masked_edges_backward_cpu, pred_edges_backward_cpu[t]], 1)
|
||||
self.gen_writer.add_image('img/edge-b:gt-pred', edge_results, self.iteration)
|
||||
|
||||
# console logs
|
||||
if self.config['global_rank'] == 0:
|
||||
pbar.update(1)
|
||||
pbar.set_description((f"flow: {flow_loss.item():.3f}; "
|
||||
f"warp: {warp_loss.item():.3f}; "
|
||||
f"edge: {edge_loss.item():.3f}; "
|
||||
f"lr: {self.get_lr()}"))
|
||||
|
||||
if self.iteration % self.train_args['log_freq'] == 0:
|
||||
logging.info(f"[Iter {self.iteration}] "
|
||||
f"flow: {flow_loss.item():.4f}; "
|
||||
f"warp: {warp_loss.item():.4f}")
|
||||
|
||||
# saving models
|
||||
if self.iteration % self.train_args['save_freq'] == 0:
|
||||
self.save(int(self.iteration))
|
||||
|
||||
if self.iteration > self.train_args['iterations']:
|
||||
break
|
||||
|
||||
train_data = self.prefetcher.next()
|
||||
@@ -0,0 +1,371 @@
|
||||
import os
|
||||
import io
|
||||
import cv2
|
||||
import random
|
||||
import numpy as np
|
||||
from PIL import Image, ImageOps
|
||||
import zipfile
|
||||
import math
|
||||
|
||||
import torch
|
||||
import matplotlib
|
||||
import matplotlib.patches as patches
|
||||
from matplotlib.path import Path
|
||||
from matplotlib import pyplot as plt
|
||||
from torchvision import transforms
|
||||
|
||||
# matplotlib.use('agg')
|
||||
|
||||
# ###########################################################################
|
||||
# Directory IO
|
||||
# ###########################################################################
|
||||
|
||||
|
||||
def read_dirnames_under_root(root_dir):
|
||||
dirnames = [
|
||||
name for i, name in enumerate(sorted(os.listdir(root_dir)))
|
||||
if os.path.isdir(os.path.join(root_dir, name))
|
||||
]
|
||||
print(f'Reading directories under {root_dir}, num: {len(dirnames)}')
|
||||
return dirnames
|
||||
|
||||
|
||||
class TrainZipReader(object):
|
||||
file_dict = dict()
|
||||
|
||||
def __init__(self):
|
||||
super(TrainZipReader, self).__init__()
|
||||
|
||||
@staticmethod
|
||||
def build_file_dict(path):
|
||||
file_dict = TrainZipReader.file_dict
|
||||
if path in file_dict:
|
||||
return file_dict[path]
|
||||
else:
|
||||
file_handle = zipfile.ZipFile(path, 'r')
|
||||
file_dict[path] = file_handle
|
||||
return file_dict[path]
|
||||
|
||||
@staticmethod
|
||||
def imread(path, idx):
|
||||
zfile = TrainZipReader.build_file_dict(path)
|
||||
filelist = zfile.namelist()
|
||||
filelist.sort()
|
||||
data = zfile.read(filelist[idx])
|
||||
#
|
||||
im = Image.open(io.BytesIO(data))
|
||||
return im
|
||||
|
||||
|
||||
class TestZipReader(object):
|
||||
file_dict = dict()
|
||||
|
||||
def __init__(self):
|
||||
super(TestZipReader, self).__init__()
|
||||
|
||||
@staticmethod
|
||||
def build_file_dict(path):
|
||||
file_dict = TestZipReader.file_dict
|
||||
if path in file_dict:
|
||||
return file_dict[path]
|
||||
else:
|
||||
file_handle = zipfile.ZipFile(path, 'r')
|
||||
file_dict[path] = file_handle
|
||||
return file_dict[path]
|
||||
|
||||
@staticmethod
|
||||
def imread(path, idx):
|
||||
zfile = TestZipReader.build_file_dict(path)
|
||||
filelist = zfile.namelist()
|
||||
filelist.sort()
|
||||
data = zfile.read(filelist[idx])
|
||||
file_bytes = np.asarray(bytearray(data), dtype=np.uint8)
|
||||
im = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR)
|
||||
im = Image.fromarray(cv2.cvtColor(im, cv2.COLOR_BGR2RGB))
|
||||
# im = Image.open(io.BytesIO(data))
|
||||
return im
|
||||
|
||||
|
||||
# ###########################################################################
|
||||
# Data augmentation
|
||||
# ###########################################################################
|
||||
|
||||
|
||||
def to_tensors():
|
||||
return transforms.Compose([Stack(), ToTorchFormatTensor()])
|
||||
|
||||
|
||||
class GroupRandomHorizontalFlowFlip(object):
|
||||
"""Randomly horizontally flips the given PIL.Image with a probability of 0.5
|
||||
"""
|
||||
def __call__(self, img_group, flowF_group, flowB_group):
|
||||
v = random.random()
|
||||
if v < 0.5:
|
||||
ret_img = [
|
||||
img.transpose(Image.FLIP_LEFT_RIGHT) for img in img_group
|
||||
]
|
||||
ret_flowF = [ff[:, ::-1] * [-1.0, 1.0] for ff in flowF_group]
|
||||
ret_flowB = [fb[:, ::-1] * [-1.0, 1.0] for fb in flowB_group]
|
||||
return ret_img, ret_flowF, ret_flowB
|
||||
else:
|
||||
return img_group, flowF_group, flowB_group
|
||||
|
||||
|
||||
class GroupRandomHorizontalFlip(object):
|
||||
"""Randomly horizontally flips the given PIL.Image with a probability of 0.5
|
||||
"""
|
||||
def __call__(self, img_group, is_flow=False):
|
||||
v = random.random()
|
||||
if v < 0.5:
|
||||
ret = [img.transpose(Image.FLIP_LEFT_RIGHT) for img in img_group]
|
||||
if is_flow:
|
||||
for i in range(0, len(ret), 2):
|
||||
# invert flow pixel values when flipping
|
||||
ret[i] = ImageOps.invert(ret[i])
|
||||
return ret
|
||||
else:
|
||||
return img_group
|
||||
|
||||
|
||||
class Stack(object):
|
||||
def __init__(self, roll=False):
|
||||
self.roll = roll
|
||||
|
||||
def __call__(self, img_group):
|
||||
mode = img_group[0].mode
|
||||
if mode == '1':
|
||||
img_group = [img.convert('L') for img in img_group]
|
||||
mode = 'L'
|
||||
if mode == 'L':
|
||||
return np.stack([np.expand_dims(x, 2) for x in img_group], axis=2)
|
||||
elif mode == 'RGB':
|
||||
if self.roll:
|
||||
return np.stack([np.array(x)[:, :, ::-1] for x in img_group],
|
||||
axis=2)
|
||||
else:
|
||||
return np.stack(img_group, axis=2)
|
||||
else:
|
||||
raise NotImplementedError(f"Image mode {mode}")
|
||||
|
||||
|
||||
class ToTorchFormatTensor(object):
|
||||
""" Converts a PIL.Image (RGB) or numpy.ndarray (H x W x C) in the range [0, 255]
|
||||
to a torch.FloatTensor of shape (C x H x W) in the range [0.0, 1.0] """
|
||||
def __init__(self, div=True):
|
||||
self.div = div
|
||||
|
||||
def __call__(self, pic):
|
||||
if isinstance(pic, np.ndarray):
|
||||
# numpy img: [L, C, H, W]
|
||||
img = torch.from_numpy(pic).permute(2, 3, 0, 1).contiguous()
|
||||
else:
|
||||
# handle PIL Image
|
||||
img = torch.ByteTensor(torch.ByteStorage.from_buffer(
|
||||
pic.tobytes()))
|
||||
img = img.view(pic.size[1], pic.size[0], len(pic.mode))
|
||||
# put it from HWC to CHW format
|
||||
# yikes, this transpose takes 80% of the loading time/CPU
|
||||
img = img.transpose(0, 1).transpose(0, 2).contiguous()
|
||||
img = img.float().div(255) if self.div else img.float()
|
||||
return img
|
||||
|
||||
|
||||
# ###########################################################################
|
||||
# Create masks with random shape
|
||||
# ###########################################################################
|
||||
|
||||
|
||||
def create_random_shape_with_random_motion(video_length,
|
||||
imageHeight=240,
|
||||
imageWidth=432):
|
||||
# get a random shape
|
||||
height = random.randint(imageHeight // 3, imageHeight - 1)
|
||||
width = random.randint(imageWidth // 3, imageWidth - 1)
|
||||
edge_num = random.randint(6, 8)
|
||||
ratio = random.randint(6, 8) / 10
|
||||
|
||||
region = get_random_shape(edge_num=edge_num,
|
||||
ratio=ratio,
|
||||
height=height,
|
||||
width=width)
|
||||
region_width, region_height = region.size
|
||||
# get random position
|
||||
x, y = random.randint(0, imageHeight - region_height), random.randint(
|
||||
0, imageWidth - region_width)
|
||||
velocity = get_random_velocity(max_speed=3)
|
||||
m = Image.fromarray(np.zeros((imageHeight, imageWidth)).astype(np.uint8))
|
||||
m.paste(region, (y, x, y + region.size[0], x + region.size[1]))
|
||||
masks = [m.convert('L')]
|
||||
# return fixed masks
|
||||
if random.uniform(0, 1) > 0.5:
|
||||
return masks * video_length
|
||||
# return moving masks
|
||||
for _ in range(video_length - 1):
|
||||
x, y, velocity = random_move_control_points(x,
|
||||
y,
|
||||
imageHeight,
|
||||
imageWidth,
|
||||
velocity,
|
||||
region.size,
|
||||
maxLineAcceleration=(3,
|
||||
0.5),
|
||||
maxInitSpeed=3)
|
||||
m = Image.fromarray(
|
||||
np.zeros((imageHeight, imageWidth)).astype(np.uint8))
|
||||
m.paste(region, (y, x, y + region.size[0], x + region.size[1]))
|
||||
masks.append(m.convert('L'))
|
||||
return masks
|
||||
|
||||
|
||||
def create_random_shape_with_random_motion_zoom_rotation(video_length, zoomin=0.9, zoomout=1.1, rotmin=1, rotmax=10, imageHeight=240, imageWidth=432):
|
||||
# get a random shape
|
||||
assert zoomin < 1, "Zoom-in parameter must be smaller than 1"
|
||||
assert zoomout > 1, "Zoom-out parameter must be larger than 1"
|
||||
assert rotmin < rotmax, "Minimum value of rotation must be smaller than maximun value !"
|
||||
height = random.randint(imageHeight//3, imageHeight-1)
|
||||
width = random.randint(imageWidth//3, imageWidth-1)
|
||||
edge_num = random.randint(6, 8)
|
||||
ratio = random.randint(6, 8)/10
|
||||
region = get_random_shape(
|
||||
edge_num=edge_num, ratio=ratio, height=height, width=width)
|
||||
region_width, region_height = region.size
|
||||
# get random position
|
||||
x, y = random.randint(
|
||||
0, imageHeight-region_height), random.randint(0, imageWidth-region_width)
|
||||
velocity = get_random_velocity(max_speed=3)
|
||||
m = Image.fromarray(np.zeros((imageHeight, imageWidth)).astype(np.uint8))
|
||||
m.paste(region, (y, x, y+region.size[0], x+region.size[1]))
|
||||
masks = [m.convert('L')]
|
||||
# return fixed masks
|
||||
if random.uniform(0, 1) > 0.5:
|
||||
return masks*video_length # -> directly copy all the base masks
|
||||
# return moving masks
|
||||
for _ in range(video_length-1):
|
||||
x, y, velocity = random_move_control_points(
|
||||
x, y, imageHeight, imageWidth, velocity, region.size, maxLineAcceleration=(3, 0.5), maxInitSpeed=3)
|
||||
m = Image.fromarray(
|
||||
np.zeros((imageHeight, imageWidth)).astype(np.uint8))
|
||||
### add by kaidong, to simulate zoon-in, zoom-out and rotation
|
||||
extra_transform = random.uniform(0, 1)
|
||||
# zoom in and zoom out
|
||||
if extra_transform > 0.75:
|
||||
resize_coefficient = random.uniform(zoomin, zoomout)
|
||||
region = region.resize((math.ceil(region_width * resize_coefficient), math.ceil(region_height * resize_coefficient)), Image.NEAREST)
|
||||
m.paste(region, (y, x, y + region.size[0], x + region.size[1]))
|
||||
region_width, region_height = region.size
|
||||
# rotation
|
||||
elif extra_transform > 0.5:
|
||||
m.paste(region, (y, x, y + region.size[0], x + region.size[1]))
|
||||
m = m.rotate(random.randint(rotmin, rotmax))
|
||||
# region_width, region_height = region.size
|
||||
### end
|
||||
else:
|
||||
m.paste(region, (y, x, y+region.size[0], x+region.size[1]))
|
||||
masks.append(m.convert('L'))
|
||||
return masks
|
||||
|
||||
|
||||
def get_random_shape(edge_num=9, ratio=0.7, width=432, height=240):
|
||||
'''
|
||||
There is the initial point and 3 points per cubic bezier curve.
|
||||
Thus, the curve will only pass though n points, which will be the sharp edges.
|
||||
The other 2 modify the shape of the bezier curve.
|
||||
edge_num, Number of possibly sharp edges
|
||||
points_num, number of points in the Path
|
||||
ratio, (0, 1) magnitude of the perturbation from the unit circle,
|
||||
'''
|
||||
points_num = edge_num*3 + 1
|
||||
angles = np.linspace(0, 2*np.pi, points_num)
|
||||
codes = np.full(points_num, Path.CURVE4)
|
||||
codes[0] = Path.MOVETO
|
||||
# Using this instead of Path.CLOSEPOLY avoids an innecessary straight line
|
||||
verts = np.stack((np.cos(angles), np.sin(angles))).T * \
|
||||
(2*ratio*np.random.random(points_num)+1-ratio)[:, None]
|
||||
verts[-1, :] = verts[0, :]
|
||||
path = Path(verts, codes)
|
||||
# draw paths into images
|
||||
fig = plt.figure()
|
||||
ax = fig.add_subplot(111)
|
||||
patch = patches.PathPatch(path, facecolor='black', lw=2)
|
||||
ax.add_patch(patch)
|
||||
ax.set_xlim(np.min(verts)*1.1, np.max(verts)*1.1)
|
||||
ax.set_ylim(np.min(verts)*1.1, np.max(verts)*1.1)
|
||||
ax.axis('off') # removes the axis to leave only the shape
|
||||
fig.canvas.draw()
|
||||
# convert plt images into numpy images
|
||||
data = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
|
||||
data = data.reshape((fig.canvas.get_width_height()[::-1] + (3,)))
|
||||
plt.close(fig)
|
||||
# postprocess
|
||||
data = cv2.resize(data, (width, height))[:, :, 0]
|
||||
data = (1 - np.array(data > 0).astype(np.uint8))*255
|
||||
corrdinates = np.where(data > 0)
|
||||
xmin, xmax, ymin, ymax = np.min(corrdinates[0]), np.max(
|
||||
corrdinates[0]), np.min(corrdinates[1]), np.max(corrdinates[1])
|
||||
region = Image.fromarray(data).crop((ymin, xmin, ymax, xmax))
|
||||
return region
|
||||
|
||||
|
||||
def random_accelerate(velocity, maxAcceleration, dist='uniform'):
|
||||
speed, angle = velocity
|
||||
d_speed, d_angle = maxAcceleration
|
||||
if dist == 'uniform':
|
||||
speed += np.random.uniform(-d_speed, d_speed)
|
||||
angle += np.random.uniform(-d_angle, d_angle)
|
||||
elif dist == 'guassian':
|
||||
speed += np.random.normal(0, d_speed / 2)
|
||||
angle += np.random.normal(0, d_angle / 2)
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f'Distribution type {dist} is not supported.')
|
||||
return (speed, angle)
|
||||
|
||||
|
||||
def get_random_velocity(max_speed=3, dist='uniform'):
|
||||
if dist == 'uniform':
|
||||
speed = np.random.uniform(max_speed)
|
||||
elif dist == 'guassian':
|
||||
speed = np.abs(np.random.normal(0, max_speed / 2))
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f'Distribution type {dist} is not supported.')
|
||||
angle = np.random.uniform(0, 2 * np.pi)
|
||||
return (speed, angle)
|
||||
|
||||
|
||||
def random_move_control_points(X,
|
||||
Y,
|
||||
imageHeight,
|
||||
imageWidth,
|
||||
lineVelocity,
|
||||
region_size,
|
||||
maxLineAcceleration=(3, 0.5),
|
||||
maxInitSpeed=3):
|
||||
region_width, region_height = region_size
|
||||
speed, angle = lineVelocity
|
||||
X += int(speed * np.cos(angle))
|
||||
Y += int(speed * np.sin(angle))
|
||||
lineVelocity = random_accelerate(lineVelocity,
|
||||
maxLineAcceleration,
|
||||
dist='guassian')
|
||||
if ((X > imageHeight - region_height) or (X < 0)
|
||||
or (Y > imageWidth - region_width) or (Y < 0)):
|
||||
lineVelocity = get_random_velocity(maxInitSpeed, dist='guassian')
|
||||
new_X = np.clip(X, 0, imageHeight - region_height)
|
||||
new_Y = np.clip(Y, 0, imageWidth - region_width)
|
||||
return new_X, new_Y, lineVelocity
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
trials = 10
|
||||
for _ in range(trials):
|
||||
video_length = 10
|
||||
# The returned masks are either stationary (50%) or moving (50%)
|
||||
masks = create_random_shape_with_random_motion(video_length,
|
||||
imageHeight=240,
|
||||
imageWidth=432)
|
||||
|
||||
for m in masks:
|
||||
cv2.imshow('mask', np.array(m))
|
||||
cv2.waitKey(500)
|
||||
@@ -0,0 +1 @@
|
||||
|
||||
@@ -0,0 +1,256 @@
|
||||
import math
|
||||
from typing import Tuple
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from .gaussian import gaussian_blur2d
|
||||
from .kernels import get_canny_nms_kernel, get_hysteresis_kernel
|
||||
from .sobel import spatial_gradient
|
||||
|
||||
def rgb_to_grayscale(image, rgb_weights = None):
|
||||
if len(image.shape) < 3 or image.shape[-3] != 3:
|
||||
raise ValueError(f"Input size must have a shape of (*, 3, H, W). Got {image.shape}")
|
||||
|
||||
if rgb_weights is None:
|
||||
# 8 bit images
|
||||
if image.dtype == torch.uint8:
|
||||
rgb_weights = torch.tensor([76, 150, 29], device=image.device, dtype=torch.uint8)
|
||||
# floating point images
|
||||
elif image.dtype in (torch.float16, torch.float32, torch.float64):
|
||||
rgb_weights = torch.tensor([0.299, 0.587, 0.114], device=image.device, dtype=image.dtype)
|
||||
else:
|
||||
raise TypeError(f"Unknown data type: {image.dtype}")
|
||||
else:
|
||||
# is tensor that we make sure is in the same device/dtype
|
||||
rgb_weights = rgb_weights.to(image)
|
||||
|
||||
# unpack the color image channels with RGB order
|
||||
r = image[..., 0:1, :, :]
|
||||
g = image[..., 1:2, :, :]
|
||||
b = image[..., 2:3, :, :]
|
||||
|
||||
w_r, w_g, w_b = rgb_weights.unbind()
|
||||
return w_r * r + w_g * g + w_b * b
|
||||
|
||||
|
||||
def canny(
|
||||
input: torch.Tensor,
|
||||
low_threshold: float = 0.1,
|
||||
high_threshold: float = 0.2,
|
||||
kernel_size: Tuple[int, int] = (5, 5),
|
||||
sigma: Tuple[float, float] = (1, 1),
|
||||
hysteresis: bool = True,
|
||||
eps: float = 1e-6,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
r"""Find edges of the input image and filters them using the Canny algorithm.
|
||||
|
||||
.. image:: _static/img/canny.png
|
||||
|
||||
Args:
|
||||
input: input image tensor with shape :math:`(B,C,H,W)`.
|
||||
low_threshold: lower threshold for the hysteresis procedure.
|
||||
high_threshold: upper threshold for the hysteresis procedure.
|
||||
kernel_size: the size of the kernel for the gaussian blur.
|
||||
sigma: the standard deviation of the kernel for the gaussian blur.
|
||||
hysteresis: if True, applies the hysteresis edge tracking.
|
||||
Otherwise, the edges are divided between weak (0.5) and strong (1) edges.
|
||||
eps: regularization number to avoid NaN during backprop.
|
||||
|
||||
Returns:
|
||||
- the canny edge magnitudes map, shape of :math:`(B,1,H,W)`.
|
||||
- the canny edge detection filtered by thresholds and hysteresis, shape of :math:`(B,1,H,W)`.
|
||||
|
||||
.. note::
|
||||
See a working example `here <https://kornia-tutorials.readthedocs.io/en/latest/
|
||||
canny.html>`__.
|
||||
|
||||
Example:
|
||||
>>> input = torch.rand(5, 3, 4, 4)
|
||||
>>> magnitude, edges = canny(input) # 5x3x4x4
|
||||
>>> magnitude.shape
|
||||
torch.Size([5, 1, 4, 4])
|
||||
>>> edges.shape
|
||||
torch.Size([5, 1, 4, 4])
|
||||
"""
|
||||
if not isinstance(input, torch.Tensor):
|
||||
raise TypeError(f"Input type is not a torch.Tensor. Got {type(input)}")
|
||||
|
||||
if not len(input.shape) == 4:
|
||||
raise ValueError(f"Invalid input shape, we expect BxCxHxW. Got: {input.shape}")
|
||||
|
||||
if low_threshold > high_threshold:
|
||||
raise ValueError(
|
||||
"Invalid input thresholds. low_threshold should be smaller than the high_threshold. Got: {}>{}".format(
|
||||
low_threshold, high_threshold
|
||||
)
|
||||
)
|
||||
|
||||
if low_threshold < 0 and low_threshold > 1:
|
||||
raise ValueError(f"Invalid input threshold. low_threshold should be in range (0,1). Got: {low_threshold}")
|
||||
|
||||
if high_threshold < 0 and high_threshold > 1:
|
||||
raise ValueError(f"Invalid input threshold. high_threshold should be in range (0,1). Got: {high_threshold}")
|
||||
|
||||
device: torch.device = input.device
|
||||
dtype: torch.dtype = input.dtype
|
||||
|
||||
# To Grayscale
|
||||
if input.shape[1] == 3:
|
||||
input = rgb_to_grayscale(input)
|
||||
|
||||
# Gaussian filter
|
||||
blurred: torch.Tensor = gaussian_blur2d(input, kernel_size, sigma)
|
||||
|
||||
# Compute the gradients
|
||||
gradients: torch.Tensor = spatial_gradient(blurred, normalized=False)
|
||||
|
||||
# Unpack the edges
|
||||
gx: torch.Tensor = gradients[:, :, 0]
|
||||
gy: torch.Tensor = gradients[:, :, 1]
|
||||
|
||||
# Compute gradient magnitude and angle
|
||||
magnitude: torch.Tensor = torch.sqrt(gx * gx + gy * gy + eps)
|
||||
angle: torch.Tensor = torch.atan2(gy, gx)
|
||||
|
||||
# Radians to Degrees
|
||||
angle = 180.0 * angle / math.pi
|
||||
|
||||
# Round angle to the nearest 45 degree
|
||||
angle = torch.round(angle / 45) * 45
|
||||
|
||||
# Non-maximal suppression
|
||||
nms_kernels: torch.Tensor = get_canny_nms_kernel(device, dtype)
|
||||
nms_magnitude: torch.Tensor = F.conv2d(magnitude, nms_kernels, padding=nms_kernels.shape[-1] // 2)
|
||||
|
||||
# Get the indices for both directions
|
||||
positive_idx: torch.Tensor = (angle / 45) % 8
|
||||
positive_idx = positive_idx.long()
|
||||
|
||||
negative_idx: torch.Tensor = ((angle / 45) + 4) % 8
|
||||
negative_idx = negative_idx.long()
|
||||
|
||||
# Apply the non-maximum suppression to the different directions
|
||||
channel_select_filtered_positive: torch.Tensor = torch.gather(nms_magnitude, 1, positive_idx)
|
||||
channel_select_filtered_negative: torch.Tensor = torch.gather(nms_magnitude, 1, negative_idx)
|
||||
|
||||
channel_select_filtered: torch.Tensor = torch.stack(
|
||||
[channel_select_filtered_positive, channel_select_filtered_negative], 1
|
||||
)
|
||||
|
||||
is_max: torch.Tensor = channel_select_filtered.min(dim=1)[0] > 0.0
|
||||
|
||||
magnitude = magnitude * is_max
|
||||
|
||||
# Threshold
|
||||
edges: torch.Tensor = F.threshold(magnitude, low_threshold, 0.0)
|
||||
|
||||
low: torch.Tensor = magnitude > low_threshold
|
||||
high: torch.Tensor = magnitude > high_threshold
|
||||
|
||||
edges = low * 0.5 + high * 0.5
|
||||
edges = edges.to(dtype)
|
||||
|
||||
# Hysteresis
|
||||
if hysteresis:
|
||||
edges_old: torch.Tensor = -torch.ones(edges.shape, device=edges.device, dtype=dtype)
|
||||
hysteresis_kernels: torch.Tensor = get_hysteresis_kernel(device, dtype)
|
||||
|
||||
while ((edges_old - edges).abs() != 0).any():
|
||||
weak: torch.Tensor = (edges == 0.5).float()
|
||||
strong: torch.Tensor = (edges == 1).float()
|
||||
|
||||
hysteresis_magnitude: torch.Tensor = F.conv2d(
|
||||
edges, hysteresis_kernels, padding=hysteresis_kernels.shape[-1] // 2
|
||||
)
|
||||
hysteresis_magnitude = (hysteresis_magnitude == 1).any(1, keepdim=True).to(dtype)
|
||||
hysteresis_magnitude = hysteresis_magnitude * weak + strong
|
||||
|
||||
edges_old = edges.clone()
|
||||
edges = hysteresis_magnitude + (hysteresis_magnitude == 0) * weak * 0.5
|
||||
|
||||
edges = hysteresis_magnitude
|
||||
|
||||
return magnitude, edges
|
||||
|
||||
|
||||
class Canny(nn.Module):
|
||||
r"""Module that finds edges of the input image and filters them using the Canny algorithm.
|
||||
|
||||
Args:
|
||||
input: input image tensor with shape :math:`(B,C,H,W)`.
|
||||
low_threshold: lower threshold for the hysteresis procedure.
|
||||
high_threshold: upper threshold for the hysteresis procedure.
|
||||
kernel_size: the size of the kernel for the gaussian blur.
|
||||
sigma: the standard deviation of the kernel for the gaussian blur.
|
||||
hysteresis: if True, applies the hysteresis edge tracking.
|
||||
Otherwise, the edges are divided between weak (0.5) and strong (1) edges.
|
||||
eps: regularization number to avoid NaN during backprop.
|
||||
|
||||
Returns:
|
||||
- the canny edge magnitudes map, shape of :math:`(B,1,H,W)`.
|
||||
- the canny edge detection filtered by thresholds and hysteresis, shape of :math:`(B,1,H,W)`.
|
||||
|
||||
Example:
|
||||
>>> input = torch.rand(5, 3, 4, 4)
|
||||
>>> magnitude, edges = Canny()(input) # 5x3x4x4
|
||||
>>> magnitude.shape
|
||||
torch.Size([5, 1, 4, 4])
|
||||
>>> edges.shape
|
||||
torch.Size([5, 1, 4, 4])
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
low_threshold: float = 0.1,
|
||||
high_threshold: float = 0.2,
|
||||
kernel_size: Tuple[int, int] = (5, 5),
|
||||
sigma: Tuple[float, float] = (1, 1),
|
||||
hysteresis: bool = True,
|
||||
eps: float = 1e-6,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
if low_threshold > high_threshold:
|
||||
raise ValueError(
|
||||
"Invalid input thresholds. low_threshold should be\
|
||||
smaller than the high_threshold. Got: {}>{}".format(
|
||||
low_threshold, high_threshold
|
||||
)
|
||||
)
|
||||
|
||||
if low_threshold < 0 or low_threshold > 1:
|
||||
raise ValueError(f"Invalid input threshold. low_threshold should be in range (0,1). Got: {low_threshold}")
|
||||
|
||||
if high_threshold < 0 or high_threshold > 1:
|
||||
raise ValueError(f"Invalid input threshold. high_threshold should be in range (0,1). Got: {high_threshold}")
|
||||
|
||||
# Gaussian blur parameters
|
||||
self.kernel_size = kernel_size
|
||||
self.sigma = sigma
|
||||
|
||||
# Double threshold
|
||||
self.low_threshold = low_threshold
|
||||
self.high_threshold = high_threshold
|
||||
|
||||
# Hysteresis
|
||||
self.hysteresis = hysteresis
|
||||
|
||||
self.eps: float = eps
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return ''.join(
|
||||
(
|
||||
f'{type(self).__name__}(',
|
||||
', '.join(
|
||||
f'{name}={getattr(self, name)}' for name in sorted(self.__dict__) if not name.startswith('_')
|
||||
),
|
||||
')',
|
||||
)
|
||||
)
|
||||
|
||||
def forward(self, input: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
return canny(
|
||||
input, self.low_threshold, self.high_threshold, self.kernel_size, self.sigma, self.hysteresis, self.eps
|
||||
)
|
||||
@@ -0,0 +1,288 @@
|
||||
from typing import List
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from .kernels import normalize_kernel2d
|
||||
|
||||
|
||||
def _compute_padding(kernel_size: List[int]) -> List[int]:
|
||||
"""Compute padding tuple."""
|
||||
# 4 or 6 ints: (padding_left, padding_right,padding_top,padding_bottom)
|
||||
# https://pytorch.org/docs/stable/nn.html#torch.nn.functional.pad
|
||||
if len(kernel_size) < 2:
|
||||
raise AssertionError(kernel_size)
|
||||
computed = [k - 1 for k in kernel_size]
|
||||
|
||||
# for even kernels we need to do asymmetric padding :(
|
||||
out_padding = 2 * len(kernel_size) * [0]
|
||||
|
||||
for i in range(len(kernel_size)):
|
||||
computed_tmp = computed[-(i + 1)]
|
||||
|
||||
pad_front = computed_tmp // 2
|
||||
pad_rear = computed_tmp - pad_front
|
||||
|
||||
out_padding[2 * i + 0] = pad_front
|
||||
out_padding[2 * i + 1] = pad_rear
|
||||
|
||||
return out_padding
|
||||
|
||||
|
||||
def filter2d(
|
||||
input: torch.Tensor,
|
||||
kernel: torch.Tensor,
|
||||
border_type: str = 'reflect',
|
||||
normalized: bool = False,
|
||||
padding: str = 'same',
|
||||
) -> torch.Tensor:
|
||||
r"""Convolve a tensor with a 2d kernel.
|
||||
|
||||
The function applies a given kernel to a tensor. The kernel is applied
|
||||
independently at each depth channel of the tensor. Before applying the
|
||||
kernel, the function applies padding according to the specified mode so
|
||||
that the output remains in the same shape.
|
||||
|
||||
Args:
|
||||
input: the input tensor with shape of
|
||||
:math:`(B, C, H, W)`.
|
||||
kernel: the kernel to be convolved with the input
|
||||
tensor. The kernel shape must be :math:`(1, kH, kW)` or :math:`(B, kH, kW)`.
|
||||
border_type: the padding mode to be applied before convolving.
|
||||
The expected modes are: ``'constant'``, ``'reflect'``,
|
||||
``'replicate'`` or ``'circular'``.
|
||||
normalized: If True, kernel will be L1 normalized.
|
||||
padding: This defines the type of padding.
|
||||
2 modes available ``'same'`` or ``'valid'``.
|
||||
|
||||
Return:
|
||||
torch.Tensor: the convolved tensor of same size and numbers of channels
|
||||
as the input with shape :math:`(B, C, H, W)`.
|
||||
|
||||
Example:
|
||||
>>> input = torch.tensor([[[
|
||||
... [0., 0., 0., 0., 0.],
|
||||
... [0., 0., 0., 0., 0.],
|
||||
... [0., 0., 5., 0., 0.],
|
||||
... [0., 0., 0., 0., 0.],
|
||||
... [0., 0., 0., 0., 0.],]]])
|
||||
>>> kernel = torch.ones(1, 3, 3)
|
||||
>>> filter2d(input, kernel, padding='same')
|
||||
tensor([[[[0., 0., 0., 0., 0.],
|
||||
[0., 5., 5., 5., 0.],
|
||||
[0., 5., 5., 5., 0.],
|
||||
[0., 5., 5., 5., 0.],
|
||||
[0., 0., 0., 0., 0.]]]])
|
||||
"""
|
||||
if not isinstance(input, torch.Tensor):
|
||||
raise TypeError(f"Input input is not torch.Tensor. Got {type(input)}")
|
||||
|
||||
if not isinstance(kernel, torch.Tensor):
|
||||
raise TypeError(f"Input kernel is not torch.Tensor. Got {type(kernel)}")
|
||||
|
||||
if not isinstance(border_type, str):
|
||||
raise TypeError(f"Input border_type is not string. Got {type(border_type)}")
|
||||
|
||||
if border_type not in ['constant', 'reflect', 'replicate', 'circular']:
|
||||
raise ValueError(
|
||||
f"Invalid border type, we expect 'constant', \
|
||||
'reflect', 'replicate', 'circular'. Got:{border_type}"
|
||||
)
|
||||
|
||||
if not isinstance(padding, str):
|
||||
raise TypeError(f"Input padding is not string. Got {type(padding)}")
|
||||
|
||||
if padding not in ['valid', 'same']:
|
||||
raise ValueError(f"Invalid padding mode, we expect 'valid' or 'same'. Got: {padding}")
|
||||
|
||||
if not len(input.shape) == 4:
|
||||
raise ValueError(f"Invalid input shape, we expect BxCxHxW. Got: {input.shape}")
|
||||
|
||||
if (not len(kernel.shape) == 3) and not ((kernel.shape[0] == 0) or (kernel.shape[0] == input.shape[0])):
|
||||
raise ValueError(f"Invalid kernel shape, we expect 1xHxW or BxHxW. Got: {kernel.shape}")
|
||||
|
||||
# prepare kernel
|
||||
b, c, h, w = input.shape
|
||||
tmp_kernel: torch.Tensor = kernel.unsqueeze(1).to(input)
|
||||
|
||||
if normalized:
|
||||
tmp_kernel = normalize_kernel2d(tmp_kernel)
|
||||
|
||||
tmp_kernel = tmp_kernel.expand(-1, c, -1, -1)
|
||||
|
||||
height, width = tmp_kernel.shape[-2:]
|
||||
|
||||
# pad the input tensor
|
||||
if padding == 'same':
|
||||
padding_shape: List[int] = _compute_padding([height, width])
|
||||
input = F.pad(input, padding_shape, mode=border_type)
|
||||
|
||||
# kernel and input tensor reshape to align element-wise or batch-wise params
|
||||
tmp_kernel = tmp_kernel.reshape(-1, 1, height, width)
|
||||
input = input.view(-1, tmp_kernel.size(0), input.size(-2), input.size(-1))
|
||||
|
||||
# convolve the tensor with the kernel.
|
||||
output = F.conv2d(input, tmp_kernel, groups=tmp_kernel.size(0), padding=0, stride=1)
|
||||
|
||||
if padding == 'same':
|
||||
out = output.view(b, c, h, w)
|
||||
else:
|
||||
out = output.view(b, c, h - height + 1, w - width + 1)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
def filter2d_separable(
|
||||
input: torch.Tensor,
|
||||
kernel_x: torch.Tensor,
|
||||
kernel_y: torch.Tensor,
|
||||
border_type: str = 'reflect',
|
||||
normalized: bool = False,
|
||||
padding: str = 'same',
|
||||
) -> torch.Tensor:
|
||||
r"""Convolve a tensor with two 1d kernels, in x and y directions.
|
||||
|
||||
The function applies a given kernel to a tensor. The kernel is applied
|
||||
independently at each depth channel of the tensor. Before applying the
|
||||
kernel, the function applies padding according to the specified mode so
|
||||
that the output remains in the same shape.
|
||||
|
||||
Args:
|
||||
input: the input tensor with shape of
|
||||
:math:`(B, C, H, W)`.
|
||||
kernel_x: the kernel to be convolved with the input
|
||||
tensor. The kernel shape must be :math:`(1, kW)` or :math:`(B, kW)`.
|
||||
kernel_y: the kernel to be convolved with the input
|
||||
tensor. The kernel shape must be :math:`(1, kH)` or :math:`(B, kH)`.
|
||||
border_type: the padding mode to be applied before convolving.
|
||||
The expected modes are: ``'constant'``, ``'reflect'``,
|
||||
``'replicate'`` or ``'circular'``.
|
||||
normalized: If True, kernel will be L1 normalized.
|
||||
padding: This defines the type of padding.
|
||||
2 modes available ``'same'`` or ``'valid'``.
|
||||
|
||||
Return:
|
||||
torch.Tensor: the convolved tensor of same size and numbers of channels
|
||||
as the input with shape :math:`(B, C, H, W)`.
|
||||
|
||||
Example:
|
||||
>>> input = torch.tensor([[[
|
||||
... [0., 0., 0., 0., 0.],
|
||||
... [0., 0., 0., 0., 0.],
|
||||
... [0., 0., 5., 0., 0.],
|
||||
... [0., 0., 0., 0., 0.],
|
||||
... [0., 0., 0., 0., 0.],]]])
|
||||
>>> kernel = torch.ones(1, 3)
|
||||
|
||||
>>> filter2d_separable(input, kernel, kernel, padding='same')
|
||||
tensor([[[[0., 0., 0., 0., 0.],
|
||||
[0., 5., 5., 5., 0.],
|
||||
[0., 5., 5., 5., 0.],
|
||||
[0., 5., 5., 5., 0.],
|
||||
[0., 0., 0., 0., 0.]]]])
|
||||
"""
|
||||
out_x = filter2d(input, kernel_x.unsqueeze(0), border_type, normalized, padding)
|
||||
out = filter2d(out_x, kernel_y.unsqueeze(-1), border_type, normalized, padding)
|
||||
return out
|
||||
|
||||
|
||||
def filter3d(
|
||||
input: torch.Tensor, kernel: torch.Tensor, border_type: str = 'replicate', normalized: bool = False
|
||||
) -> torch.Tensor:
|
||||
r"""Convolve a tensor with a 3d kernel.
|
||||
|
||||
The function applies a given kernel to a tensor. The kernel is applied
|
||||
independently at each depth channel of the tensor. Before applying the
|
||||
kernel, the function applies padding according to the specified mode so
|
||||
that the output remains in the same shape.
|
||||
|
||||
Args:
|
||||
input: the input tensor with shape of
|
||||
:math:`(B, C, D, H, W)`.
|
||||
kernel: the kernel to be convolved with the input
|
||||
tensor. The kernel shape must be :math:`(1, kD, kH, kW)` or :math:`(B, kD, kH, kW)`.
|
||||
border_type: the padding mode to be applied before convolving.
|
||||
The expected modes are: ``'constant'``,
|
||||
``'replicate'`` or ``'circular'``.
|
||||
normalized: If True, kernel will be L1 normalized.
|
||||
|
||||
Return:
|
||||
the convolved tensor of same size and numbers of channels
|
||||
as the input with shape :math:`(B, C, D, H, W)`.
|
||||
|
||||
Example:
|
||||
>>> input = torch.tensor([[[
|
||||
... [[0., 0., 0., 0., 0.],
|
||||
... [0., 0., 0., 0., 0.],
|
||||
... [0., 0., 0., 0., 0.],
|
||||
... [0., 0., 0., 0., 0.],
|
||||
... [0., 0., 0., 0., 0.]],
|
||||
... [[0., 0., 0., 0., 0.],
|
||||
... [0., 0., 0., 0., 0.],
|
||||
... [0., 0., 5., 0., 0.],
|
||||
... [0., 0., 0., 0., 0.],
|
||||
... [0., 0., 0., 0., 0.]],
|
||||
... [[0., 0., 0., 0., 0.],
|
||||
... [0., 0., 0., 0., 0.],
|
||||
... [0., 0., 0., 0., 0.],
|
||||
... [0., 0., 0., 0., 0.],
|
||||
... [0., 0., 0., 0., 0.]]
|
||||
... ]]])
|
||||
>>> kernel = torch.ones(1, 3, 3, 3)
|
||||
>>> filter3d(input, kernel)
|
||||
tensor([[[[[0., 0., 0., 0., 0.],
|
||||
[0., 5., 5., 5., 0.],
|
||||
[0., 5., 5., 5., 0.],
|
||||
[0., 5., 5., 5., 0.],
|
||||
[0., 0., 0., 0., 0.]],
|
||||
<BLANKLINE>
|
||||
[[0., 0., 0., 0., 0.],
|
||||
[0., 5., 5., 5., 0.],
|
||||
[0., 5., 5., 5., 0.],
|
||||
[0., 5., 5., 5., 0.],
|
||||
[0., 0., 0., 0., 0.]],
|
||||
<BLANKLINE>
|
||||
[[0., 0., 0., 0., 0.],
|
||||
[0., 5., 5., 5., 0.],
|
||||
[0., 5., 5., 5., 0.],
|
||||
[0., 5., 5., 5., 0.],
|
||||
[0., 0., 0., 0., 0.]]]]])
|
||||
"""
|
||||
if not isinstance(input, torch.Tensor):
|
||||
raise TypeError(f"Input border_type is not torch.Tensor. Got {type(input)}")
|
||||
|
||||
if not isinstance(kernel, torch.Tensor):
|
||||
raise TypeError(f"Input border_type is not torch.Tensor. Got {type(kernel)}")
|
||||
|
||||
if not isinstance(border_type, str):
|
||||
raise TypeError(f"Input border_type is not string. Got {type(kernel)}")
|
||||
|
||||
if not len(input.shape) == 5:
|
||||
raise ValueError(f"Invalid input shape, we expect BxCxDxHxW. Got: {input.shape}")
|
||||
|
||||
if not len(kernel.shape) == 4 and kernel.shape[0] != 1:
|
||||
raise ValueError(f"Invalid kernel shape, we expect 1xDxHxW. Got: {kernel.shape}")
|
||||
|
||||
# prepare kernel
|
||||
b, c, d, h, w = input.shape
|
||||
tmp_kernel: torch.Tensor = kernel.unsqueeze(1).to(input)
|
||||
|
||||
if normalized:
|
||||
bk, dk, hk, wk = kernel.shape
|
||||
tmp_kernel = normalize_kernel2d(tmp_kernel.view(bk, dk, hk * wk)).view_as(tmp_kernel)
|
||||
|
||||
tmp_kernel = tmp_kernel.expand(-1, c, -1, -1, -1)
|
||||
|
||||
# pad the input tensor
|
||||
depth, height, width = tmp_kernel.shape[-3:]
|
||||
padding_shape: List[int] = _compute_padding([depth, height, width])
|
||||
input_pad: torch.Tensor = F.pad(input, padding_shape, mode=border_type)
|
||||
|
||||
# kernel and input tensor reshape to align element-wise or batch-wise params
|
||||
tmp_kernel = tmp_kernel.reshape(-1, 1, depth, height, width)
|
||||
input_pad = input_pad.view(-1, tmp_kernel.size(0), input_pad.size(-3), input_pad.size(-2), input_pad.size(-1))
|
||||
|
||||
# convolve the tensor with the kernel.
|
||||
output = F.conv3d(input_pad, tmp_kernel, groups=tmp_kernel.size(0), padding=0, stride=1)
|
||||
|
||||
return output.view(b, c, d, h, w)
|
||||
@@ -0,0 +1,116 @@
|
||||
from typing import Tuple
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from .filter import filter2d, filter2d_separable
|
||||
from .kernels import get_gaussian_kernel1d, get_gaussian_kernel2d
|
||||
|
||||
|
||||
def gaussian_blur2d(
|
||||
input: torch.Tensor,
|
||||
kernel_size: Tuple[int, int],
|
||||
sigma: Tuple[float, float],
|
||||
border_type: str = 'reflect',
|
||||
separable: bool = True,
|
||||
) -> torch.Tensor:
|
||||
r"""Create an operator that blurs a tensor using a Gaussian filter.
|
||||
|
||||
.. image:: _static/img/gaussian_blur2d.png
|
||||
|
||||
The operator smooths the given tensor with a gaussian kernel by convolving
|
||||
it to each channel. It supports batched operation.
|
||||
|
||||
Arguments:
|
||||
input: the input tensor with shape :math:`(B,C,H,W)`.
|
||||
kernel_size: the size of the kernel.
|
||||
sigma: the standard deviation of the kernel.
|
||||
border_type: the padding mode to be applied before convolving.
|
||||
The expected modes are: ``'constant'``, ``'reflect'``,
|
||||
``'replicate'`` or ``'circular'``. Default: ``'reflect'``.
|
||||
separable: run as composition of two 1d-convolutions.
|
||||
|
||||
Returns:
|
||||
the blurred tensor with shape :math:`(B, C, H, W)`.
|
||||
|
||||
.. note::
|
||||
See a working example `here <https://kornia-tutorials.readthedocs.io/en/latest/
|
||||
gaussian_blur.html>`__.
|
||||
|
||||
Examples:
|
||||
>>> input = torch.rand(2, 4, 5, 5)
|
||||
>>> output = gaussian_blur2d(input, (3, 3), (1.5, 1.5))
|
||||
>>> output.shape
|
||||
torch.Size([2, 4, 5, 5])
|
||||
"""
|
||||
if separable:
|
||||
kernel_x: torch.Tensor = get_gaussian_kernel1d(kernel_size[1], sigma[1])
|
||||
kernel_y: torch.Tensor = get_gaussian_kernel1d(kernel_size[0], sigma[0])
|
||||
out = filter2d_separable(input, kernel_x[None], kernel_y[None], border_type)
|
||||
else:
|
||||
kernel: torch.Tensor = get_gaussian_kernel2d(kernel_size, sigma)
|
||||
out = filter2d(input, kernel[None], border_type)
|
||||
return out
|
||||
|
||||
|
||||
class GaussianBlur2d(nn.Module):
|
||||
r"""Create an operator that blurs a tensor using a Gaussian filter.
|
||||
|
||||
The operator smooths the given tensor with a gaussian kernel by convolving
|
||||
it to each channel. It supports batched operation.
|
||||
|
||||
Arguments:
|
||||
kernel_size: the size of the kernel.
|
||||
sigma: the standard deviation of the kernel.
|
||||
border_type: the padding mode to be applied before convolving.
|
||||
The expected modes are: ``'constant'``, ``'reflect'``,
|
||||
``'replicate'`` or ``'circular'``. Default: ``'reflect'``.
|
||||
separable: run as composition of two 1d-convolutions.
|
||||
|
||||
Returns:
|
||||
the blurred tensor.
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(B, C, H, W)`
|
||||
- Output: :math:`(B, C, H, W)`
|
||||
|
||||
Examples::
|
||||
|
||||
>>> input = torch.rand(2, 4, 5, 5)
|
||||
>>> gauss = GaussianBlur2d((3, 3), (1.5, 1.5))
|
||||
>>> output = gauss(input) # 2x4x5x5
|
||||
>>> output.shape
|
||||
torch.Size([2, 4, 5, 5])
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
kernel_size: Tuple[int, int],
|
||||
sigma: Tuple[float, float],
|
||||
border_type: str = 'reflect',
|
||||
separable: bool = True,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.kernel_size: Tuple[int, int] = kernel_size
|
||||
self.sigma: Tuple[float, float] = sigma
|
||||
self.border_type = border_type
|
||||
self.separable = separable
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
self.__class__.__name__
|
||||
+ '(kernel_size='
|
||||
+ str(self.kernel_size)
|
||||
+ ', '
|
||||
+ 'sigma='
|
||||
+ str(self.sigma)
|
||||
+ ', '
|
||||
+ 'border_type='
|
||||
+ self.border_type
|
||||
+ 'separable='
|
||||
+ str(self.separable)
|
||||
+ ')'
|
||||
)
|
||||
|
||||
def forward(self, input: torch.Tensor) -> torch.Tensor:
|
||||
return gaussian_blur2d(input, self.kernel_size, self.sigma, self.border_type, self.separable)
|
||||
@@ -0,0 +1,690 @@
|
||||
import math
|
||||
from math import sqrt
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def normalize_kernel2d(input: torch.Tensor) -> torch.Tensor:
|
||||
r"""Normalize both derivative and smoothing kernel."""
|
||||
if len(input.size()) < 2:
|
||||
raise TypeError(f"input should be at least 2D tensor. Got {input.size()}")
|
||||
norm: torch.Tensor = input.abs().sum(dim=-1).sum(dim=-1)
|
||||
return input / (norm.unsqueeze(-1).unsqueeze(-1))
|
||||
|
||||
|
||||
def gaussian(window_size: int, sigma: float) -> torch.Tensor:
|
||||
device, dtype = None, None
|
||||
if isinstance(sigma, torch.Tensor):
|
||||
device, dtype = sigma.device, sigma.dtype
|
||||
x = torch.arange(window_size, device=device, dtype=dtype) - window_size // 2
|
||||
if window_size % 2 == 0:
|
||||
x = x + 0.5
|
||||
|
||||
gauss = torch.exp((-x.pow(2.0) / (2 * sigma**2)).float())
|
||||
return gauss / gauss.sum()
|
||||
|
||||
|
||||
def gaussian_discrete_erf(window_size: int, sigma) -> torch.Tensor:
|
||||
r"""Discrete Gaussian by interpolating the error function.
|
||||
|
||||
Adapted from:
|
||||
https://github.com/Project-MONAI/MONAI/blob/master/monai/networks/layers/convutils.py
|
||||
"""
|
||||
device = sigma.device if isinstance(sigma, torch.Tensor) else None
|
||||
sigma = torch.as_tensor(sigma, dtype=torch.float, device=device)
|
||||
x = torch.arange(window_size).float() - window_size // 2
|
||||
t = 0.70710678 / torch.abs(sigma)
|
||||
gauss = 0.5 * ((t * (x + 0.5)).erf() - (t * (x - 0.5)).erf())
|
||||
gauss = gauss.clamp(min=0)
|
||||
return gauss / gauss.sum()
|
||||
|
||||
|
||||
def _modified_bessel_0(x: torch.Tensor) -> torch.Tensor:
|
||||
r"""Adapted from:
|
||||
|
||||
https://github.com/Project-MONAI/MONAI/blob/master/monai/networks/layers/convutils.py
|
||||
"""
|
||||
if torch.abs(x) < 3.75:
|
||||
y = (x / 3.75) * (x / 3.75)
|
||||
return 1.0 + y * (
|
||||
3.5156229 + y * (3.0899424 + y * (1.2067492 + y * (0.2659732 + y * (0.360768e-1 + y * 0.45813e-2))))
|
||||
)
|
||||
ax = torch.abs(x)
|
||||
y = 3.75 / ax
|
||||
ans = 0.916281e-2 + y * (-0.2057706e-1 + y * (0.2635537e-1 + y * (-0.1647633e-1 + y * 0.392377e-2)))
|
||||
coef = 0.39894228 + y * (0.1328592e-1 + y * (0.225319e-2 + y * (-0.157565e-2 + y * ans)))
|
||||
return (torch.exp(ax) / torch.sqrt(ax)) * coef
|
||||
|
||||
|
||||
def _modified_bessel_1(x: torch.Tensor) -> torch.Tensor:
|
||||
r"""adapted from:
|
||||
|
||||
https://github.com/Project-MONAI/MONAI/blob/master/monai/networks/layers/convutils.py
|
||||
"""
|
||||
if torch.abs(x) < 3.75:
|
||||
y = (x / 3.75) * (x / 3.75)
|
||||
ans = 0.51498869 + y * (0.15084934 + y * (0.2658733e-1 + y * (0.301532e-2 + y * 0.32411e-3)))
|
||||
return torch.abs(x) * (0.5 + y * (0.87890594 + y * ans))
|
||||
ax = torch.abs(x)
|
||||
y = 3.75 / ax
|
||||
ans = 0.2282967e-1 + y * (-0.2895312e-1 + y * (0.1787654e-1 - y * 0.420059e-2))
|
||||
ans = 0.39894228 + y * (-0.3988024e-1 + y * (-0.362018e-2 + y * (0.163801e-2 + y * (-0.1031555e-1 + y * ans))))
|
||||
ans = ans * torch.exp(ax) / torch.sqrt(ax)
|
||||
return -ans if x < 0.0 else ans
|
||||
|
||||
|
||||
def _modified_bessel_i(n: int, x: torch.Tensor) -> torch.Tensor:
|
||||
r"""adapted from:
|
||||
|
||||
https://github.com/Project-MONAI/MONAI/blob/master/monai/networks/layers/convutils.py
|
||||
"""
|
||||
if n < 2:
|
||||
raise ValueError("n must be greater than 1.")
|
||||
if x == 0.0:
|
||||
return x
|
||||
device = x.device
|
||||
tox = 2.0 / torch.abs(x)
|
||||
ans = torch.tensor(0.0, device=device)
|
||||
bip = torch.tensor(0.0, device=device)
|
||||
bi = torch.tensor(1.0, device=device)
|
||||
m = int(2 * (n + int(sqrt(40.0 * n))))
|
||||
for j in range(m, 0, -1):
|
||||
bim = bip + float(j) * tox * bi
|
||||
bip = bi
|
||||
bi = bim
|
||||
if abs(bi) > 1.0e10:
|
||||
ans = ans * 1.0e-10
|
||||
bi = bi * 1.0e-10
|
||||
bip = bip * 1.0e-10
|
||||
if j == n:
|
||||
ans = bip
|
||||
ans = ans * _modified_bessel_0(x) / bi
|
||||
return -ans if x < 0.0 and (n % 2) == 1 else ans
|
||||
|
||||
|
||||
def gaussian_discrete(window_size, sigma) -> torch.Tensor:
|
||||
r"""Discrete Gaussian kernel based on the modified Bessel functions.
|
||||
|
||||
Adapted from:
|
||||
https://github.com/Project-MONAI/MONAI/blob/master/monai/networks/layers/convutils.py
|
||||
"""
|
||||
device = sigma.device if isinstance(sigma, torch.Tensor) else None
|
||||
sigma = torch.as_tensor(sigma, dtype=torch.float, device=device)
|
||||
sigma2 = sigma * sigma
|
||||
tail = int(window_size // 2)
|
||||
out_pos: List[Optional[torch.Tensor]] = [None] * (tail + 1)
|
||||
out_pos[0] = _modified_bessel_0(sigma2)
|
||||
out_pos[1] = _modified_bessel_1(sigma2)
|
||||
for k in range(2, len(out_pos)):
|
||||
out_pos[k] = _modified_bessel_i(k, sigma2)
|
||||
out = out_pos[:0:-1]
|
||||
out.extend(out_pos)
|
||||
out = torch.stack(out) * torch.exp(sigma2) # type: ignore
|
||||
return out / out.sum() # type: ignore
|
||||
|
||||
|
||||
def laplacian_1d(window_size) -> torch.Tensor:
|
||||
r"""One could also use the Laplacian of Gaussian formula to design the filter."""
|
||||
|
||||
filter_1d = torch.ones(window_size)
|
||||
filter_1d[window_size // 2] = 1 - window_size
|
||||
laplacian_1d: torch.Tensor = filter_1d
|
||||
return laplacian_1d
|
||||
|
||||
|
||||
def get_box_kernel2d(kernel_size: Tuple[int, int]) -> torch.Tensor:
|
||||
r"""Utility function that returns a box filter."""
|
||||
kx: float = float(kernel_size[0])
|
||||
ky: float = float(kernel_size[1])
|
||||
scale: torch.Tensor = torch.tensor(1.0) / torch.tensor([kx * ky])
|
||||
tmp_kernel: torch.Tensor = torch.ones(1, kernel_size[0], kernel_size[1])
|
||||
return scale.to(tmp_kernel.dtype) * tmp_kernel
|
||||
|
||||
|
||||
def get_binary_kernel2d(window_size: Tuple[int, int]) -> torch.Tensor:
|
||||
r"""Create a binary kernel to extract the patches.
|
||||
|
||||
If the window size is HxW will create a (H*W)xHxW kernel.
|
||||
"""
|
||||
window_range: int = window_size[0] * window_size[1]
|
||||
kernel: torch.Tensor = torch.zeros(window_range, window_range)
|
||||
for i in range(window_range):
|
||||
kernel[i, i] += 1.0
|
||||
return kernel.view(window_range, 1, window_size[0], window_size[1])
|
||||
|
||||
|
||||
def get_sobel_kernel_3x3() -> torch.Tensor:
|
||||
"""Utility function that returns a sobel kernel of 3x3."""
|
||||
return torch.tensor([[-1.0, 0.0, 1.0], [-2.0, 0.0, 2.0], [-1.0, 0.0, 1.0]])
|
||||
|
||||
|
||||
def get_sobel_kernel_5x5_2nd_order() -> torch.Tensor:
|
||||
"""Utility function that returns a 2nd order sobel kernel of 5x5."""
|
||||
return torch.tensor(
|
||||
[
|
||||
[-1.0, 0.0, 2.0, 0.0, -1.0],
|
||||
[-4.0, 0.0, 8.0, 0.0, -4.0],
|
||||
[-6.0, 0.0, 12.0, 0.0, -6.0],
|
||||
[-4.0, 0.0, 8.0, 0.0, -4.0],
|
||||
[-1.0, 0.0, 2.0, 0.0, -1.0],
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def _get_sobel_kernel_5x5_2nd_order_xy() -> torch.Tensor:
|
||||
"""Utility function that returns a 2nd order sobel kernel of 5x5."""
|
||||
return torch.tensor(
|
||||
[
|
||||
[-1.0, -2.0, 0.0, 2.0, 1.0],
|
||||
[-2.0, -4.0, 0.0, 4.0, 2.0],
|
||||
[0.0, 0.0, 0.0, 0.0, 0.0],
|
||||
[2.0, 4.0, 0.0, -4.0, -2.0],
|
||||
[1.0, 2.0, 0.0, -2.0, -1.0],
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def get_diff_kernel_3x3() -> torch.Tensor:
|
||||
"""Utility function that returns a first order derivative kernel of 3x3."""
|
||||
return torch.tensor([[-0.0, 0.0, 0.0], [-1.0, 0.0, 1.0], [-0.0, 0.0, 0.0]])
|
||||
|
||||
|
||||
def get_diff_kernel3d(device=torch.device('cpu'), dtype=torch.float) -> torch.Tensor:
|
||||
"""Utility function that returns a first order derivative kernel of 3x3x3."""
|
||||
kernel: torch.Tensor = torch.tensor(
|
||||
[
|
||||
[
|
||||
[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]],
|
||||
[[0.0, 0.0, 0.0], [-0.5, 0.0, 0.5], [0.0, 0.0, 0.0]],
|
||||
[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]],
|
||||
],
|
||||
[
|
||||
[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]],
|
||||
[[0.0, -0.5, 0.0], [0.0, 0.0, 0.0], [0.0, 0.5, 0.0]],
|
||||
[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]],
|
||||
],
|
||||
[
|
||||
[[0.0, 0.0, 0.0], [0.0, -0.5, 0.0], [0.0, 0.0, 0.0]],
|
||||
[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]],
|
||||
[[0.0, 0.0, 0.0], [0.0, 0.5, 0.0], [0.0, 0.0, 0.0]],
|
||||
],
|
||||
],
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
return kernel.unsqueeze(1)
|
||||
|
||||
|
||||
def get_diff_kernel3d_2nd_order(device=torch.device('cpu'), dtype=torch.float) -> torch.Tensor:
|
||||
"""Utility function that returns a first order derivative kernel of 3x3x3."""
|
||||
kernel: torch.Tensor = torch.tensor(
|
||||
[
|
||||
[
|
||||
[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]],
|
||||
[[0.0, 0.0, 0.0], [1.0, -2.0, 1.0], [0.0, 0.0, 0.0]],
|
||||
[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]],
|
||||
],
|
||||
[
|
||||
[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]],
|
||||
[[0.0, 1.0, 0.0], [0.0, -2.0, 0.0], [0.0, 1.0, 0.0]],
|
||||
[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]],
|
||||
],
|
||||
[
|
||||
[[0.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0]],
|
||||
[[0.0, 0.0, 0.0], [0.0, -2.0, 0.0], [0.0, 0.0, 0.0]],
|
||||
[[0.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0]],
|
||||
],
|
||||
[
|
||||
[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]],
|
||||
[[1.0, 0.0, -1.0], [0.0, 0.0, 0.0], [-1.0, 0.0, 1.0]],
|
||||
[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]],
|
||||
],
|
||||
[
|
||||
[[0.0, 1.0, 0.0], [0.0, 0.0, 0.0], [0.0, -1.0, 0.0]],
|
||||
[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]],
|
||||
[[0.0, -1.0, 0.0], [0.0, 0.0, 0.0], [0.0, 1.0, 0.0]],
|
||||
],
|
||||
[
|
||||
[[0.0, 0.0, 0.0], [1.0, 0.0, -1.0], [0.0, 0.0, 0.0]],
|
||||
[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]],
|
||||
[[0.0, 0.0, 0.0], [-1.0, 0.0, 1.0], [0.0, 0.0, 0.0]],
|
||||
],
|
||||
],
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
return kernel.unsqueeze(1)
|
||||
|
||||
|
||||
def get_sobel_kernel2d() -> torch.Tensor:
|
||||
kernel_x: torch.Tensor = get_sobel_kernel_3x3()
|
||||
kernel_y: torch.Tensor = kernel_x.transpose(0, 1)
|
||||
return torch.stack([kernel_x, kernel_y])
|
||||
|
||||
|
||||
def get_diff_kernel2d() -> torch.Tensor:
|
||||
kernel_x: torch.Tensor = get_diff_kernel_3x3()
|
||||
kernel_y: torch.Tensor = kernel_x.transpose(0, 1)
|
||||
return torch.stack([kernel_x, kernel_y])
|
||||
|
||||
|
||||
def get_sobel_kernel2d_2nd_order() -> torch.Tensor:
|
||||
gxx: torch.Tensor = get_sobel_kernel_5x5_2nd_order()
|
||||
gyy: torch.Tensor = gxx.transpose(0, 1)
|
||||
gxy: torch.Tensor = _get_sobel_kernel_5x5_2nd_order_xy()
|
||||
return torch.stack([gxx, gxy, gyy])
|
||||
|
||||
|
||||
def get_diff_kernel2d_2nd_order() -> torch.Tensor:
|
||||
gxx: torch.Tensor = torch.tensor([[0.0, 0.0, 0.0], [1.0, -2.0, 1.0], [0.0, 0.0, 0.0]])
|
||||
gyy: torch.Tensor = gxx.transpose(0, 1)
|
||||
gxy: torch.Tensor = torch.tensor([[-1.0, 0.0, 1.0], [0.0, 0.0, 0.0], [1.0, 0.0, -1.0]])
|
||||
return torch.stack([gxx, gxy, gyy])
|
||||
|
||||
|
||||
def get_spatial_gradient_kernel2d(mode: str, order: int) -> torch.Tensor:
|
||||
r"""Function that returns kernel for 1st or 2nd order image gradients, using one of the following operators:
|
||||
|
||||
sobel, diff.
|
||||
"""
|
||||
if mode not in ['sobel', 'diff']:
|
||||
raise TypeError(
|
||||
"mode should be either sobel\
|
||||
or diff. Got {}".format(
|
||||
mode
|
||||
)
|
||||
)
|
||||
if order not in [1, 2]:
|
||||
raise TypeError(
|
||||
"order should be either 1 or 2\
|
||||
Got {}".format(
|
||||
order
|
||||
)
|
||||
)
|
||||
if mode == 'sobel' and order == 1:
|
||||
kernel: torch.Tensor = get_sobel_kernel2d()
|
||||
elif mode == 'sobel' and order == 2:
|
||||
kernel = get_sobel_kernel2d_2nd_order()
|
||||
elif mode == 'diff' and order == 1:
|
||||
kernel = get_diff_kernel2d()
|
||||
elif mode == 'diff' and order == 2:
|
||||
kernel = get_diff_kernel2d_2nd_order()
|
||||
else:
|
||||
raise NotImplementedError("")
|
||||
return kernel
|
||||
|
||||
|
||||
def get_spatial_gradient_kernel3d(mode: str, order: int, device=torch.device('cpu'), dtype=torch.float) -> torch.Tensor:
|
||||
r"""Function that returns kernel for 1st or 2nd order scale pyramid gradients, using one of the following
|
||||
operators: sobel, diff."""
|
||||
if mode not in ['sobel', 'diff']:
|
||||
raise TypeError(
|
||||
"mode should be either sobel\
|
||||
or diff. Got {}".format(
|
||||
mode
|
||||
)
|
||||
)
|
||||
if order not in [1, 2]:
|
||||
raise TypeError(
|
||||
"order should be either 1 or 2\
|
||||
Got {}".format(
|
||||
order
|
||||
)
|
||||
)
|
||||
if mode == 'sobel':
|
||||
raise NotImplementedError("Sobel kernel for 3d gradient is not implemented yet")
|
||||
if mode == 'diff' and order == 1:
|
||||
kernel = get_diff_kernel3d(device, dtype)
|
||||
elif mode == 'diff' and order == 2:
|
||||
kernel = get_diff_kernel3d_2nd_order(device, dtype)
|
||||
else:
|
||||
raise NotImplementedError("")
|
||||
return kernel
|
||||
|
||||
|
||||
def get_gaussian_kernel1d(kernel_size: int, sigma: float, force_even: bool = False) -> torch.Tensor:
|
||||
r"""Function that returns Gaussian filter coefficients.
|
||||
|
||||
Args:
|
||||
kernel_size: filter size. It should be odd and positive.
|
||||
sigma: gaussian standard deviation.
|
||||
force_even: overrides requirement for odd kernel size.
|
||||
|
||||
Returns:
|
||||
1D tensor with gaussian filter coefficients.
|
||||
|
||||
Shape:
|
||||
- Output: :math:`(\text{kernel_size})`
|
||||
|
||||
Examples:
|
||||
|
||||
>>> get_gaussian_kernel1d(3, 2.5)
|
||||
tensor([0.3243, 0.3513, 0.3243])
|
||||
|
||||
>>> get_gaussian_kernel1d(5, 1.5)
|
||||
tensor([0.1201, 0.2339, 0.2921, 0.2339, 0.1201])
|
||||
"""
|
||||
if not isinstance(kernel_size, int) or ((kernel_size % 2 == 0) and not force_even) or (kernel_size <= 0):
|
||||
raise TypeError("kernel_size must be an odd positive integer. " "Got {}".format(kernel_size))
|
||||
window_1d: torch.Tensor = gaussian(kernel_size, sigma)
|
||||
return window_1d
|
||||
|
||||
|
||||
def get_gaussian_discrete_kernel1d(kernel_size: int, sigma: float, force_even: bool = False) -> torch.Tensor:
|
||||
r"""Function that returns Gaussian filter coefficients based on the modified Bessel functions. Adapted from:
|
||||
https://github.com/Project-MONAI/MONAI/blob/master/monai/networks/layers/convutils.py.
|
||||
|
||||
Args:
|
||||
kernel_size: filter size. It should be odd and positive.
|
||||
sigma: gaussian standard deviation.
|
||||
force_even: overrides requirement for odd kernel size.
|
||||
|
||||
Returns:
|
||||
1D tensor with gaussian filter coefficients.
|
||||
|
||||
Shape:
|
||||
- Output: :math:`(\text{kernel_size})`
|
||||
|
||||
Examples:
|
||||
|
||||
>>> get_gaussian_discrete_kernel1d(3, 2.5)
|
||||
tensor([0.3235, 0.3531, 0.3235])
|
||||
|
||||
>>> get_gaussian_discrete_kernel1d(5, 1.5)
|
||||
tensor([0.1096, 0.2323, 0.3161, 0.2323, 0.1096])
|
||||
"""
|
||||
if not isinstance(kernel_size, int) or ((kernel_size % 2 == 0) and not force_even) or (kernel_size <= 0):
|
||||
raise TypeError("kernel_size must be an odd positive integer. " "Got {}".format(kernel_size))
|
||||
window_1d = gaussian_discrete(kernel_size, sigma)
|
||||
return window_1d
|
||||
|
||||
|
||||
def get_gaussian_erf_kernel1d(kernel_size: int, sigma: float, force_even: bool = False) -> torch.Tensor:
|
||||
r"""Function that returns Gaussian filter coefficients by interpolating the error function, adapted from:
|
||||
https://github.com/Project-MONAI/MONAI/blob/master/monai/networks/layers/convutils.py.
|
||||
|
||||
Args:
|
||||
kernel_size: filter size. It should be odd and positive.
|
||||
sigma: gaussian standard deviation.
|
||||
force_even: overrides requirement for odd kernel size.
|
||||
|
||||
Returns:
|
||||
1D tensor with gaussian filter coefficients.
|
||||
|
||||
Shape:
|
||||
- Output: :math:`(\text{kernel_size})`
|
||||
|
||||
Examples:
|
||||
|
||||
>>> get_gaussian_erf_kernel1d(3, 2.5)
|
||||
tensor([0.3245, 0.3511, 0.3245])
|
||||
|
||||
>>> get_gaussian_erf_kernel1d(5, 1.5)
|
||||
tensor([0.1226, 0.2331, 0.2887, 0.2331, 0.1226])
|
||||
"""
|
||||
if not isinstance(kernel_size, int) or ((kernel_size % 2 == 0) and not force_even) or (kernel_size <= 0):
|
||||
raise TypeError("kernel_size must be an odd positive integer. " "Got {}".format(kernel_size))
|
||||
window_1d = gaussian_discrete_erf(kernel_size, sigma)
|
||||
return window_1d
|
||||
|
||||
|
||||
def get_gaussian_kernel2d(
|
||||
kernel_size: Tuple[int, int], sigma: Tuple[float, float], force_even: bool = False
|
||||
) -> torch.Tensor:
|
||||
r"""Function that returns Gaussian filter matrix coefficients.
|
||||
|
||||
Args:
|
||||
kernel_size: filter sizes in the x and y direction.
|
||||
Sizes should be odd and positive.
|
||||
sigma: gaussian standard deviation in the x and y
|
||||
direction.
|
||||
force_even: overrides requirement for odd kernel size.
|
||||
|
||||
Returns:
|
||||
2D tensor with gaussian filter matrix coefficients.
|
||||
|
||||
Shape:
|
||||
- Output: :math:`(\text{kernel_size}_x, \text{kernel_size}_y)`
|
||||
|
||||
Examples:
|
||||
>>> get_gaussian_kernel2d((3, 3), (1.5, 1.5))
|
||||
tensor([[0.0947, 0.1183, 0.0947],
|
||||
[0.1183, 0.1478, 0.1183],
|
||||
[0.0947, 0.1183, 0.0947]])
|
||||
>>> get_gaussian_kernel2d((3, 5), (1.5, 1.5))
|
||||
tensor([[0.0370, 0.0720, 0.0899, 0.0720, 0.0370],
|
||||
[0.0462, 0.0899, 0.1123, 0.0899, 0.0462],
|
||||
[0.0370, 0.0720, 0.0899, 0.0720, 0.0370]])
|
||||
"""
|
||||
if not isinstance(kernel_size, tuple) or len(kernel_size) != 2:
|
||||
raise TypeError(f"kernel_size must be a tuple of length two. Got {kernel_size}")
|
||||
if not isinstance(sigma, tuple) or len(sigma) != 2:
|
||||
raise TypeError(f"sigma must be a tuple of length two. Got {sigma}")
|
||||
ksize_x, ksize_y = kernel_size
|
||||
sigma_x, sigma_y = sigma
|
||||
kernel_x: torch.Tensor = get_gaussian_kernel1d(ksize_x, sigma_x, force_even)
|
||||
kernel_y: torch.Tensor = get_gaussian_kernel1d(ksize_y, sigma_y, force_even)
|
||||
kernel_2d: torch.Tensor = torch.matmul(kernel_x.unsqueeze(-1), kernel_y.unsqueeze(-1).t())
|
||||
return kernel_2d
|
||||
|
||||
|
||||
def get_laplacian_kernel1d(kernel_size: int) -> torch.Tensor:
|
||||
r"""Function that returns the coefficients of a 1D Laplacian filter.
|
||||
|
||||
Args:
|
||||
kernel_size: filter size. It should be odd and positive.
|
||||
|
||||
Returns:
|
||||
1D tensor with laplacian filter coefficients.
|
||||
|
||||
Shape:
|
||||
- Output: math:`(\text{kernel_size})`
|
||||
|
||||
Examples:
|
||||
>>> get_laplacian_kernel1d(3)
|
||||
tensor([ 1., -2., 1.])
|
||||
>>> get_laplacian_kernel1d(5)
|
||||
tensor([ 1., 1., -4., 1., 1.])
|
||||
"""
|
||||
if not isinstance(kernel_size, int) or kernel_size % 2 == 0 or kernel_size <= 0:
|
||||
raise TypeError(f"ksize must be an odd positive integer. Got {kernel_size}")
|
||||
window_1d: torch.Tensor = laplacian_1d(kernel_size)
|
||||
return window_1d
|
||||
|
||||
|
||||
def get_laplacian_kernel2d(kernel_size: int) -> torch.Tensor:
|
||||
r"""Function that returns Gaussian filter matrix coefficients.
|
||||
|
||||
Args:
|
||||
kernel_size: filter size should be odd.
|
||||
|
||||
Returns:
|
||||
2D tensor with laplacian filter matrix coefficients.
|
||||
|
||||
Shape:
|
||||
- Output: :math:`(\text{kernel_size}_x, \text{kernel_size}_y)`
|
||||
|
||||
Examples:
|
||||
>>> get_laplacian_kernel2d(3)
|
||||
tensor([[ 1., 1., 1.],
|
||||
[ 1., -8., 1.],
|
||||
[ 1., 1., 1.]])
|
||||
>>> get_laplacian_kernel2d(5)
|
||||
tensor([[ 1., 1., 1., 1., 1.],
|
||||
[ 1., 1., 1., 1., 1.],
|
||||
[ 1., 1., -24., 1., 1.],
|
||||
[ 1., 1., 1., 1., 1.],
|
||||
[ 1., 1., 1., 1., 1.]])
|
||||
"""
|
||||
if not isinstance(kernel_size, int) or kernel_size % 2 == 0 or kernel_size <= 0:
|
||||
raise TypeError(f"ksize must be an odd positive integer. Got {kernel_size}")
|
||||
|
||||
kernel = torch.ones((kernel_size, kernel_size))
|
||||
mid = kernel_size // 2
|
||||
kernel[mid, mid] = 1 - kernel_size**2
|
||||
kernel_2d: torch.Tensor = kernel
|
||||
return kernel_2d
|
||||
|
||||
|
||||
def get_pascal_kernel_2d(kernel_size: int, norm: bool = True) -> torch.Tensor:
|
||||
"""Generate pascal filter kernel by kernel size.
|
||||
|
||||
Args:
|
||||
kernel_size: height and width of the kernel.
|
||||
norm: if to normalize the kernel or not. Default: True.
|
||||
|
||||
Returns:
|
||||
kernel shaped as :math:`(kernel_size, kernel_size)`
|
||||
|
||||
Examples:
|
||||
>>> get_pascal_kernel_2d(1)
|
||||
tensor([[1.]])
|
||||
>>> get_pascal_kernel_2d(4)
|
||||
tensor([[0.0156, 0.0469, 0.0469, 0.0156],
|
||||
[0.0469, 0.1406, 0.1406, 0.0469],
|
||||
[0.0469, 0.1406, 0.1406, 0.0469],
|
||||
[0.0156, 0.0469, 0.0469, 0.0156]])
|
||||
>>> get_pascal_kernel_2d(4, norm=False)
|
||||
tensor([[1., 3., 3., 1.],
|
||||
[3., 9., 9., 3.],
|
||||
[3., 9., 9., 3.],
|
||||
[1., 3., 3., 1.]])
|
||||
"""
|
||||
a = get_pascal_kernel_1d(kernel_size)
|
||||
|
||||
filt = a[:, None] * a[None, :]
|
||||
if norm:
|
||||
filt = filt / torch.sum(filt)
|
||||
return filt
|
||||
|
||||
|
||||
def get_pascal_kernel_1d(kernel_size: int, norm: bool = False) -> torch.Tensor:
|
||||
"""Generate Yang Hui triangle (Pascal's triangle) by a given number.
|
||||
|
||||
Args:
|
||||
kernel_size: height and width of the kernel.
|
||||
norm: if to normalize the kernel or not. Default: False.
|
||||
|
||||
Returns:
|
||||
kernel shaped as :math:`(kernel_size,)`
|
||||
|
||||
Examples:
|
||||
>>> get_pascal_kernel_1d(1)
|
||||
tensor([1.])
|
||||
>>> get_pascal_kernel_1d(2)
|
||||
tensor([1., 1.])
|
||||
>>> get_pascal_kernel_1d(3)
|
||||
tensor([1., 2., 1.])
|
||||
>>> get_pascal_kernel_1d(4)
|
||||
tensor([1., 3., 3., 1.])
|
||||
>>> get_pascal_kernel_1d(5)
|
||||
tensor([1., 4., 6., 4., 1.])
|
||||
>>> get_pascal_kernel_1d(6)
|
||||
tensor([ 1., 5., 10., 10., 5., 1.])
|
||||
"""
|
||||
pre: List[float] = []
|
||||
cur: List[float] = []
|
||||
for i in range(kernel_size):
|
||||
cur = [1.0] * (i + 1)
|
||||
|
||||
for j in range(1, i // 2 + 1):
|
||||
value = pre[j - 1] + pre[j]
|
||||
cur[j] = value
|
||||
if i != 2 * j:
|
||||
cur[-j - 1] = value
|
||||
pre = cur
|
||||
|
||||
out = torch.as_tensor(cur)
|
||||
if norm:
|
||||
out = out / torch.sum(out)
|
||||
return out
|
||||
|
||||
|
||||
def get_canny_nms_kernel(device=torch.device('cpu'), dtype=torch.float) -> torch.Tensor:
|
||||
"""Utility function that returns 3x3 kernels for the Canny Non-maximal suppression."""
|
||||
kernel: torch.Tensor = torch.tensor(
|
||||
[
|
||||
[[0.0, 0.0, 0.0], [0.0, 1.0, -1.0], [0.0, 0.0, 0.0]],
|
||||
[[0.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, -1.0]],
|
||||
[[0.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, -1.0, 0.0]],
|
||||
[[0.0, 0.0, 0.0], [0.0, 1.0, 0.0], [-1.0, 0.0, 0.0]],
|
||||
[[0.0, 0.0, 0.0], [-1.0, 1.0, 0.0], [0.0, 0.0, 0.0]],
|
||||
[[-1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0]],
|
||||
[[0.0, -1.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0]],
|
||||
[[0.0, 0.0, -1.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0]],
|
||||
],
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
return kernel.unsqueeze(1)
|
||||
|
||||
|
||||
def get_hysteresis_kernel(device=torch.device('cpu'), dtype=torch.float) -> torch.Tensor:
|
||||
"""Utility function that returns the 3x3 kernels for the Canny hysteresis."""
|
||||
kernel: torch.Tensor = torch.tensor(
|
||||
[
|
||||
[[0.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 0.0]],
|
||||
[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 1.0]],
|
||||
[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 1.0, 0.0]],
|
||||
[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [1.0, 0.0, 0.0]],
|
||||
[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 0.0]],
|
||||
[[1.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]],
|
||||
[[0.0, 1.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]],
|
||||
[[0.0, 0.0, 1.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]],
|
||||
],
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
return kernel.unsqueeze(1)
|
||||
|
||||
|
||||
def get_hanning_kernel1d(kernel_size: int, device=torch.device('cpu'), dtype=torch.float) -> torch.Tensor:
|
||||
r"""Returns Hanning (also known as Hann) kernel, used in signal processing and KCF tracker.
|
||||
|
||||
.. math:: w(n) = 0.5 - 0.5cos\\left(\\frac{2\\pi{n}}{M-1}\\right)
|
||||
\\qquad 0 \\leq n \\leq M-1
|
||||
|
||||
See further in numpy docs https://numpy.org/doc/stable/reference/generated/numpy.hanning.html
|
||||
|
||||
Args:
|
||||
kernel_size: The size the of the kernel. It should be positive.
|
||||
|
||||
Returns:
|
||||
1D tensor with Hanning filter coefficients.
|
||||
.. math:: w(n) = 0.5 - 0.5cos\\left(\\frac{2\\pi{n}}{M-1}\\right)
|
||||
|
||||
Shape:
|
||||
- Output: math:`(\text{kernel_size})`
|
||||
|
||||
Examples:
|
||||
>>> get_hanning_kernel1d(4)
|
||||
tensor([0.0000, 0.7500, 0.7500, 0.0000])
|
||||
"""
|
||||
if not isinstance(kernel_size, int) or kernel_size <= 2:
|
||||
raise TypeError(f"ksize must be an positive integer > 2. Got {kernel_size}")
|
||||
|
||||
x: torch.Tensor = torch.arange(kernel_size, device=device, dtype=dtype)
|
||||
x = 0.5 - 0.5 * torch.cos(2.0 * math.pi * x / float(kernel_size - 1))
|
||||
return x
|
||||
|
||||
|
||||
def get_hanning_kernel2d(kernel_size: Tuple[int, int], device=torch.device('cpu'), dtype=torch.float) -> torch.Tensor:
|
||||
r"""Returns 2d Hanning kernel, used in signal processing and KCF tracker.
|
||||
|
||||
Args:
|
||||
kernel_size: The size of the kernel for the filter. It should be positive.
|
||||
|
||||
Returns:
|
||||
2D tensor with Hanning filter coefficients.
|
||||
.. math:: w(n) = 0.5 - 0.5cos\\left(\\frac{2\\pi{n}}{M-1}\\right)
|
||||
|
||||
Shape:
|
||||
- Output: math:`(\text{kernel_size[0], kernel_size[1]})`
|
||||
"""
|
||||
if kernel_size[0] <= 2 or kernel_size[1] <= 2:
|
||||
raise TypeError(f"ksize must be an tuple of positive integers > 2. Got {kernel_size}")
|
||||
ky: torch.Tensor = get_hanning_kernel1d(kernel_size[0], device, dtype)[None].T
|
||||
kx: torch.Tensor = get_hanning_kernel1d(kernel_size[1], device, dtype)[None]
|
||||
kernel2d = ky @ kx
|
||||
return kernel2d
|
||||
@@ -0,0 +1,263 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from .kernels import get_spatial_gradient_kernel2d, get_spatial_gradient_kernel3d, normalize_kernel2d
|
||||
|
||||
|
||||
def spatial_gradient(input: torch.Tensor, mode: str = 'sobel', order: int = 1, normalized: bool = True) -> torch.Tensor:
|
||||
r"""Compute the first order image derivative in both x and y using a Sobel operator.
|
||||
|
||||
.. image:: _static/img/spatial_gradient.png
|
||||
|
||||
Args:
|
||||
input: input image tensor with shape :math:`(B, C, H, W)`.
|
||||
mode: derivatives modality, can be: `sobel` or `diff`.
|
||||
order: the order of the derivatives.
|
||||
normalized: whether the output is normalized.
|
||||
|
||||
Return:
|
||||
the derivatives of the input feature map. with shape :math:`(B, C, 2, H, W)`.
|
||||
|
||||
.. note::
|
||||
See a working example `here <https://kornia-tutorials.readthedocs.io/en/latest/
|
||||
filtering_edges.html>`__.
|
||||
|
||||
Examples:
|
||||
>>> input = torch.rand(1, 3, 4, 4)
|
||||
>>> output = spatial_gradient(input) # 1x3x2x4x4
|
||||
>>> output.shape
|
||||
torch.Size([1, 3, 2, 4, 4])
|
||||
"""
|
||||
if not isinstance(input, torch.Tensor):
|
||||
raise TypeError(f"Input type is not a torch.Tensor. Got {type(input)}")
|
||||
|
||||
if not len(input.shape) == 4:
|
||||
raise ValueError(f"Invalid input shape, we expect BxCxHxW. Got: {input.shape}")
|
||||
# allocate kernel
|
||||
kernel: torch.Tensor = get_spatial_gradient_kernel2d(mode, order)
|
||||
if normalized:
|
||||
kernel = normalize_kernel2d(kernel)
|
||||
|
||||
# prepare kernel
|
||||
b, c, h, w = input.shape
|
||||
tmp_kernel: torch.Tensor = kernel.to(input).detach()
|
||||
tmp_kernel = tmp_kernel.unsqueeze(1).unsqueeze(1)
|
||||
|
||||
# convolve input tensor with sobel kernel
|
||||
kernel_flip: torch.Tensor = tmp_kernel.flip(-3)
|
||||
|
||||
# Pad with "replicate for spatial dims, but with zeros for channel
|
||||
spatial_pad = [kernel.size(1) // 2, kernel.size(1) // 2, kernel.size(2) // 2, kernel.size(2) // 2]
|
||||
out_channels: int = 3 if order == 2 else 2
|
||||
padded_inp: torch.Tensor = F.pad(input.reshape(b * c, 1, h, w), spatial_pad, 'replicate')[:, :, None]
|
||||
|
||||
return F.conv3d(padded_inp, kernel_flip, padding=0).view(b, c, out_channels, h, w)
|
||||
|
||||
|
||||
def spatial_gradient3d(input: torch.Tensor, mode: str = 'diff', order: int = 1) -> torch.Tensor:
|
||||
r"""Compute the first and second order volume derivative in x, y and d using a diff operator.
|
||||
|
||||
Args:
|
||||
input: input features tensor with shape :math:`(B, C, D, H, W)`.
|
||||
mode: derivatives modality, can be: `sobel` or `diff`.
|
||||
order: the order of the derivatives.
|
||||
|
||||
Return:
|
||||
the spatial gradients of the input feature map with shape math:`(B, C, 3, D, H, W)`
|
||||
or :math:`(B, C, 6, D, H, W)`.
|
||||
|
||||
Examples:
|
||||
>>> input = torch.rand(1, 4, 2, 4, 4)
|
||||
>>> output = spatial_gradient3d(input)
|
||||
>>> output.shape
|
||||
torch.Size([1, 4, 3, 2, 4, 4])
|
||||
"""
|
||||
if not isinstance(input, torch.Tensor):
|
||||
raise TypeError(f"Input type is not a torch.Tensor. Got {type(input)}")
|
||||
|
||||
if not len(input.shape) == 5:
|
||||
raise ValueError(f"Invalid input shape, we expect BxCxDxHxW. Got: {input.shape}")
|
||||
b, c, d, h, w = input.shape
|
||||
dev = input.device
|
||||
dtype = input.dtype
|
||||
if (mode == 'diff') and (order == 1):
|
||||
# we go for the special case implementation due to conv3d bad speed
|
||||
x: torch.Tensor = F.pad(input, 6 * [1], 'replicate')
|
||||
center = slice(1, -1)
|
||||
left = slice(0, -2)
|
||||
right = slice(2, None)
|
||||
out = torch.empty(b, c, 3, d, h, w, device=dev, dtype=dtype)
|
||||
out[..., 0, :, :, :] = x[..., center, center, right] - x[..., center, center, left]
|
||||
out[..., 1, :, :, :] = x[..., center, right, center] - x[..., center, left, center]
|
||||
out[..., 2, :, :, :] = x[..., right, center, center] - x[..., left, center, center]
|
||||
out = 0.5 * out
|
||||
else:
|
||||
# prepare kernel
|
||||
# allocate kernel
|
||||
kernel: torch.Tensor = get_spatial_gradient_kernel3d(mode, order)
|
||||
|
||||
tmp_kernel: torch.Tensor = kernel.to(input).detach()
|
||||
tmp_kernel = tmp_kernel.repeat(c, 1, 1, 1, 1)
|
||||
|
||||
# convolve input tensor with grad kernel
|
||||
kernel_flip: torch.Tensor = tmp_kernel.flip(-3)
|
||||
|
||||
# Pad with "replicate for spatial dims, but with zeros for channel
|
||||
spatial_pad = [
|
||||
kernel.size(2) // 2,
|
||||
kernel.size(2) // 2,
|
||||
kernel.size(3) // 2,
|
||||
kernel.size(3) // 2,
|
||||
kernel.size(4) // 2,
|
||||
kernel.size(4) // 2,
|
||||
]
|
||||
out_ch: int = 6 if order == 2 else 3
|
||||
out = F.conv3d(F.pad(input, spatial_pad, 'replicate'), kernel_flip, padding=0, groups=c).view(
|
||||
b, c, out_ch, d, h, w
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
def sobel(input: torch.Tensor, normalized: bool = True, eps: float = 1e-6) -> torch.Tensor:
|
||||
r"""Compute the Sobel operator and returns the magnitude per channel.
|
||||
|
||||
.. image:: _static/img/sobel.png
|
||||
|
||||
Args:
|
||||
input: the input image with shape :math:`(B,C,H,W)`.
|
||||
normalized: if True, L1 norm of the kernel is set to 1.
|
||||
eps: regularization number to avoid NaN during backprop.
|
||||
|
||||
Return:
|
||||
the sobel edge gradient magnitudes map with shape :math:`(B,C,H,W)`.
|
||||
|
||||
.. note::
|
||||
See a working example `here <https://kornia-tutorials.readthedocs.io/en/latest/
|
||||
filtering_edges.html>`__.
|
||||
|
||||
Example:
|
||||
>>> input = torch.rand(1, 3, 4, 4)
|
||||
>>> output = sobel(input) # 1x3x4x4
|
||||
>>> output.shape
|
||||
torch.Size([1, 3, 4, 4])
|
||||
"""
|
||||
if not isinstance(input, torch.Tensor):
|
||||
raise TypeError(f"Input type is not a torch.Tensor. Got {type(input)}")
|
||||
|
||||
if not len(input.shape) == 4:
|
||||
raise ValueError(f"Invalid input shape, we expect BxCxHxW. Got: {input.shape}")
|
||||
|
||||
# comput the x/y gradients
|
||||
edges: torch.Tensor = spatial_gradient(input, normalized=normalized)
|
||||
|
||||
# unpack the edges
|
||||
gx: torch.Tensor = edges[:, :, 0]
|
||||
gy: torch.Tensor = edges[:, :, 1]
|
||||
|
||||
# compute gradient maginitude
|
||||
magnitude: torch.Tensor = torch.sqrt(gx * gx + gy * gy + eps)
|
||||
|
||||
return magnitude
|
||||
|
||||
|
||||
class SpatialGradient(nn.Module):
|
||||
r"""Compute the first order image derivative in both x and y using a Sobel operator.
|
||||
|
||||
Args:
|
||||
mode: derivatives modality, can be: `sobel` or `diff`.
|
||||
order: the order of the derivatives.
|
||||
normalized: whether the output is normalized.
|
||||
|
||||
Return:
|
||||
the sobel edges of the input feature map.
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(B, C, H, W)`
|
||||
- Output: :math:`(B, C, 2, H, W)`
|
||||
|
||||
Examples:
|
||||
>>> input = torch.rand(1, 3, 4, 4)
|
||||
>>> output = SpatialGradient()(input) # 1x3x2x4x4
|
||||
"""
|
||||
|
||||
def __init__(self, mode: str = 'sobel', order: int = 1, normalized: bool = True) -> None:
|
||||
super().__init__()
|
||||
self.normalized: bool = normalized
|
||||
self.order: int = order
|
||||
self.mode: str = mode
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
self.__class__.__name__ + '('
|
||||
'order=' + str(self.order) + ', ' + 'normalized=' + str(self.normalized) + ', ' + 'mode=' + self.mode + ')'
|
||||
)
|
||||
|
||||
def forward(self, input: torch.Tensor) -> torch.Tensor:
|
||||
return spatial_gradient(input, self.mode, self.order, self.normalized)
|
||||
|
||||
|
||||
class SpatialGradient3d(nn.Module):
|
||||
r"""Compute the first and second order volume derivative in x, y and d using a diff operator.
|
||||
|
||||
Args:
|
||||
mode: derivatives modality, can be: `sobel` or `diff`.
|
||||
order: the order of the derivatives.
|
||||
|
||||
Return:
|
||||
the spatial gradients of the input feature map.
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(B, C, D, H, W)`. D, H, W are spatial dimensions, gradient is calculated w.r.t to them.
|
||||
- Output: :math:`(B, C, 3, D, H, W)` or :math:`(B, C, 6, D, H, W)`
|
||||
|
||||
Examples:
|
||||
>>> input = torch.rand(1, 4, 2, 4, 4)
|
||||
>>> output = SpatialGradient3d()(input)
|
||||
>>> output.shape
|
||||
torch.Size([1, 4, 3, 2, 4, 4])
|
||||
"""
|
||||
|
||||
def __init__(self, mode: str = 'diff', order: int = 1) -> None:
|
||||
super().__init__()
|
||||
self.order: int = order
|
||||
self.mode: str = mode
|
||||
self.kernel = get_spatial_gradient_kernel3d(mode, order)
|
||||
return
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return self.__class__.__name__ + '(' 'order=' + str(self.order) + ', ' + 'mode=' + self.mode + ')'
|
||||
|
||||
def forward(self, input: torch.Tensor) -> torch.Tensor: # type: ignore
|
||||
return spatial_gradient3d(input, self.mode, self.order)
|
||||
|
||||
|
||||
class Sobel(nn.Module):
|
||||
r"""Compute the Sobel operator and returns the magnitude per channel.
|
||||
|
||||
Args:
|
||||
normalized: if True, L1 norm of the kernel is set to 1.
|
||||
eps: regularization number to avoid NaN during backprop.
|
||||
|
||||
Return:
|
||||
the sobel edge gradient magnitudes map.
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(B, C, H, W)`
|
||||
- Output: :math:`(B, C, H, W)`
|
||||
|
||||
Examples:
|
||||
>>> input = torch.rand(1, 3, 4, 4)
|
||||
>>> output = Sobel()(input) # 1x3x4x4
|
||||
"""
|
||||
|
||||
def __init__(self, normalized: bool = True, eps: float = 1e-6) -> None:
|
||||
super().__init__()
|
||||
self.normalized: bool = normalized
|
||||
self.eps: float = eps
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return self.__class__.__name__ + '(' 'normalized=' + str(self.normalized) + ')'
|
||||
|
||||
def forward(self, input: torch.Tensor) -> torch.Tensor:
|
||||
return sobel(input, self.normalized, self.eps)
|
||||
@@ -0,0 +1,133 @@
|
||||
import os
|
||||
import re
|
||||
import random
|
||||
import time
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import logging
|
||||
import numpy as np
|
||||
from os import path as osp
|
||||
|
||||
def constant_init(module, val, bias=0):
|
||||
if hasattr(module, 'weight') and module.weight is not None:
|
||||
nn.init.constant_(module.weight, val)
|
||||
if hasattr(module, 'bias') and module.bias is not None:
|
||||
nn.init.constant_(module.bias, bias)
|
||||
|
||||
initialized_logger = {}
|
||||
def get_root_logger(logger_name='basicsr', log_level=logging.INFO, log_file=None):
|
||||
"""Get the root logger.
|
||||
The logger will be initialized if it has not been initialized. By default a
|
||||
StreamHandler will be added. If `log_file` is specified, a FileHandler will
|
||||
also be added.
|
||||
Args:
|
||||
logger_name (str): root logger name. Default: 'basicsr'.
|
||||
log_file (str | None): The log filename. If specified, a FileHandler
|
||||
will be added to the root logger.
|
||||
log_level (int): The root logger level. Note that only the process of
|
||||
rank 0 is affected, while other processes will set the level to
|
||||
"Error" and be silent most of the time.
|
||||
Returns:
|
||||
logging.Logger: The root logger.
|
||||
"""
|
||||
logger = logging.getLogger(logger_name)
|
||||
# if the logger has been initialized, just return it
|
||||
if logger_name in initialized_logger:
|
||||
return logger
|
||||
|
||||
format_str = '%(asctime)s %(levelname)s: %(message)s'
|
||||
stream_handler = logging.StreamHandler()
|
||||
stream_handler.setFormatter(logging.Formatter(format_str))
|
||||
logger.addHandler(stream_handler)
|
||||
logger.propagate = False
|
||||
|
||||
if log_file is not None:
|
||||
logger.setLevel(log_level)
|
||||
# add file handler
|
||||
# file_handler = logging.FileHandler(log_file, 'w')
|
||||
file_handler = logging.FileHandler(log_file, 'a') #Shangchen: keep the previous log
|
||||
file_handler.setFormatter(logging.Formatter(format_str))
|
||||
file_handler.setLevel(log_level)
|
||||
logger.addHandler(file_handler)
|
||||
initialized_logger[logger_name] = True
|
||||
return logger
|
||||
|
||||
|
||||
IS_HIGH_VERSION = [int(m) for m in list(re.findall(r"^([0-9]+)\.([0-9]+)\.([0-9]+)([^0-9][a-zA-Z0-9]*)?(\+git.*)?$",\
|
||||
torch.__version__)[0][:3])] >= [1, 12, 0]
|
||||
|
||||
|
||||
def gpu_is_available():
|
||||
if IS_HIGH_VERSION:
|
||||
if torch.backends.mps.is_available():
|
||||
return True
|
||||
return True if torch.cuda.is_available() and torch.backends.cudnn.is_available() else False
|
||||
|
||||
|
||||
def get_device(gpu_id=None):
|
||||
if gpu_id is None:
|
||||
gpu_str = ''
|
||||
elif isinstance(gpu_id, int):
|
||||
gpu_str = f':{gpu_id}'
|
||||
else:
|
||||
raise TypeError('Input should be int value.')
|
||||
|
||||
if IS_HIGH_VERSION:
|
||||
if torch.backends.mps.is_available():
|
||||
return torch.device('mps'+gpu_str)
|
||||
return torch.device('cuda'+gpu_str if torch.cuda.is_available() and torch.backends.cudnn.is_available() else 'cpu')
|
||||
|
||||
|
||||
def set_random_seed(seed):
|
||||
"""Set random seeds."""
|
||||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
torch.manual_seed(seed)
|
||||
torch.cuda.manual_seed(seed)
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
|
||||
|
||||
def get_time_str():
|
||||
return time.strftime('%Y%m%d_%H%M%S', time.localtime())
|
||||
|
||||
|
||||
def scandir(dir_path, suffix=None, recursive=False, full_path=False):
|
||||
"""Scan a directory to find the interested files.
|
||||
|
||||
Args:
|
||||
dir_path (str): Path of the directory.
|
||||
suffix (str | tuple(str), optional): File suffix that we are
|
||||
interested in. Default: None.
|
||||
recursive (bool, optional): If set to True, recursively scan the
|
||||
directory. Default: False.
|
||||
full_path (bool, optional): If set to True, include the dir_path.
|
||||
Default: False.
|
||||
|
||||
Returns:
|
||||
A generator for all the interested files with relative pathes.
|
||||
"""
|
||||
|
||||
if (suffix is not None) and not isinstance(suffix, (str, tuple)):
|
||||
raise TypeError('"suffix" must be a string or tuple of strings')
|
||||
|
||||
root = dir_path
|
||||
|
||||
def _scandir(dir_path, suffix, recursive):
|
||||
for entry in os.scandir(dir_path):
|
||||
if not entry.name.startswith('.') and entry.is_file():
|
||||
if full_path:
|
||||
return_path = entry.path
|
||||
else:
|
||||
return_path = osp.relpath(entry.path, root)
|
||||
|
||||
if suffix is None:
|
||||
yield return_path
|
||||
elif return_path.endswith(suffix):
|
||||
yield return_path
|
||||
else:
|
||||
if recursive:
|
||||
yield from _scandir(entry.path, suffix=suffix, recursive=recursive)
|
||||
else:
|
||||
continue
|
||||
|
||||
return _scandir(dir_path, suffix=suffix, recursive=recursive)
|
||||
@@ -0,0 +1,131 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from functools import reduce
|
||||
|
||||
class BaseNetwork(nn.Module):
|
||||
def __init__(self):
|
||||
super(BaseNetwork, self).__init__()
|
||||
|
||||
def print_network(self):
|
||||
if isinstance(self, list):
|
||||
self = self[0]
|
||||
num_params = 0
|
||||
for param in self.parameters():
|
||||
num_params += param.numel()
|
||||
print(
|
||||
'Network [%s] was created. Total number of parameters: %.1f million. '
|
||||
'To see the architecture, do print(network).' %
|
||||
(type(self).__name__, num_params / 1000000))
|
||||
|
||||
def init_weights(self, init_type='normal', gain=0.02):
|
||||
'''
|
||||
initialize network's weights
|
||||
init_type: normal | xavier | kaiming | orthogonal
|
||||
https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/9451e70673400885567d08a9e97ade2524c700d0/models/networks.py#L39
|
||||
'''
|
||||
def init_func(m):
|
||||
classname = m.__class__.__name__
|
||||
if classname.find('InstanceNorm2d') != -1:
|
||||
if hasattr(m, 'weight') and m.weight is not None:
|
||||
nn.init.constant_(m.weight.data, 1.0)
|
||||
if hasattr(m, 'bias') and m.bias is not None:
|
||||
nn.init.constant_(m.bias.data, 0.0)
|
||||
elif hasattr(m, 'weight') and (classname.find('Conv') != -1
|
||||
or classname.find('Linear') != -1):
|
||||
if init_type == 'normal':
|
||||
nn.init.normal_(m.weight.data, 0.0, gain)
|
||||
elif init_type == 'xavier':
|
||||
nn.init.xavier_normal_(m.weight.data, gain=gain)
|
||||
elif init_type == 'xavier_uniform':
|
||||
nn.init.xavier_uniform_(m.weight.data, gain=1.0)
|
||||
elif init_type == 'kaiming':
|
||||
nn.init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
|
||||
elif init_type == 'orthogonal':
|
||||
nn.init.orthogonal_(m.weight.data, gain=gain)
|
||||
elif init_type == 'none': # uses pytorch's default init method
|
||||
m.reset_parameters()
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
'initialization method [%s] is not implemented' %
|
||||
init_type)
|
||||
if hasattr(m, 'bias') and m.bias is not None:
|
||||
nn.init.constant_(m.bias.data, 0.0)
|
||||
|
||||
self.apply(init_func)
|
||||
|
||||
# propagate to children
|
||||
for m in self.children():
|
||||
if hasattr(m, 'init_weights'):
|
||||
m.init_weights(init_type, gain)
|
||||
|
||||
|
||||
class Vec2Feat(nn.Module):
|
||||
def __init__(self, channel, hidden, kernel_size, stride, padding):
|
||||
super(Vec2Feat, self).__init__()
|
||||
self.relu = nn.LeakyReLU(0.2, inplace=True)
|
||||
c_out = reduce((lambda x, y: x * y), kernel_size) * channel
|
||||
self.embedding = nn.Linear(hidden, c_out)
|
||||
self.kernel_size = kernel_size
|
||||
self.stride = stride
|
||||
self.padding = padding
|
||||
self.bias_conv = nn.Conv2d(channel,
|
||||
channel,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1)
|
||||
|
||||
def forward(self, x, t, output_size):
|
||||
b_, _, _, _, c_ = x.shape
|
||||
x = x.view(b_, -1, c_)
|
||||
feat = self.embedding(x)
|
||||
b, _, c = feat.size()
|
||||
feat = feat.view(b * t, -1, c).permute(0, 2, 1)
|
||||
feat = F.fold(feat,
|
||||
output_size=output_size,
|
||||
kernel_size=self.kernel_size,
|
||||
stride=self.stride,
|
||||
padding=self.padding)
|
||||
feat = self.bias_conv(feat)
|
||||
return feat
|
||||
|
||||
|
||||
class FusionFeedForward(nn.Module):
|
||||
def __init__(self, dim, hidden_dim=1960, t2t_params=None):
|
||||
super(FusionFeedForward, self).__init__()
|
||||
# We set hidden_dim as a default to 1960
|
||||
self.fc1 = nn.Sequential(nn.Linear(dim, hidden_dim))
|
||||
self.fc2 = nn.Sequential(nn.GELU(), nn.Linear(hidden_dim, dim))
|
||||
assert t2t_params is not None
|
||||
self.t2t_params = t2t_params
|
||||
self.kernel_shape = reduce((lambda x, y: x * y), t2t_params['kernel_size']) # 49
|
||||
|
||||
def forward(self, x, output_size):
|
||||
n_vecs = 1
|
||||
for i, d in enumerate(self.t2t_params['kernel_size']):
|
||||
n_vecs *= int((output_size[i] + 2 * self.t2t_params['padding'][i] -
|
||||
(d - 1) - 1) / self.t2t_params['stride'][i] + 1)
|
||||
|
||||
x = self.fc1(x)
|
||||
b, n, c = x.size()
|
||||
normalizer = x.new_ones(b, n, self.kernel_shape).view(-1, n_vecs, self.kernel_shape).permute(0, 2, 1)
|
||||
normalizer = F.fold(normalizer,
|
||||
output_size=output_size,
|
||||
kernel_size=self.t2t_params['kernel_size'],
|
||||
padding=self.t2t_params['padding'],
|
||||
stride=self.t2t_params['stride'])
|
||||
|
||||
x = F.fold(x.view(-1, n_vecs, c).permute(0, 2, 1),
|
||||
output_size=output_size,
|
||||
kernel_size=self.t2t_params['kernel_size'],
|
||||
padding=self.t2t_params['padding'],
|
||||
stride=self.t2t_params['stride'])
|
||||
|
||||
x = F.unfold(x / normalizer,
|
||||
kernel_size=self.t2t_params['kernel_size'],
|
||||
padding=self.t2t_params['padding'],
|
||||
stride=self.t2t_params['stride']).permute(
|
||||
0, 2, 1).contiguous().view(b, n, c)
|
||||
x = self.fc2(x)
|
||||
return x
|
||||
@@ -0,0 +1,54 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.nn import init as init
|
||||
from torch.nn.modules.utils import _pair, _single
|
||||
import math
|
||||
|
||||
class ModulatedDeformConv2d(nn.Module):
|
||||
def __init__(self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride=1,
|
||||
padding=0,
|
||||
dilation=1,
|
||||
groups=1,
|
||||
deform_groups=1,
|
||||
bias=True):
|
||||
super(ModulatedDeformConv2d, self).__init__()
|
||||
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = out_channels
|
||||
self.kernel_size = _pair(kernel_size)
|
||||
self.stride = stride
|
||||
self.padding = padding
|
||||
self.dilation = dilation
|
||||
self.groups = groups
|
||||
self.deform_groups = deform_groups
|
||||
self.with_bias = bias
|
||||
# enable compatibility with nn.Conv2d
|
||||
self.transposed = False
|
||||
self.output_padding = _single(0)
|
||||
|
||||
self.weight = nn.Parameter(torch.Tensor(out_channels, in_channels // groups, *self.kernel_size))
|
||||
if bias:
|
||||
self.bias = nn.Parameter(torch.Tensor(out_channels))
|
||||
else:
|
||||
self.register_parameter('bias', None)
|
||||
self.init_weights()
|
||||
|
||||
def init_weights(self):
|
||||
n = self.in_channels
|
||||
for k in self.kernel_size:
|
||||
n *= k
|
||||
stdv = 1. / math.sqrt(n)
|
||||
self.weight.data.uniform_(-stdv, stdv)
|
||||
if self.bias is not None:
|
||||
self.bias.data.zero_()
|
||||
|
||||
if hasattr(self, 'conv_offset'):
|
||||
self.conv_offset.weight.data.zero_()
|
||||
self.conv_offset.bias.data.zero_()
|
||||
|
||||
def forward(self, x, offset, mask):
|
||||
pass
|
||||
@@ -0,0 +1,265 @@
|
||||
import argparse
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from backend.inpaint.video.raft import RAFT
|
||||
from backend.inpaint.video.model.modules.flow_loss_utils import flow_warp, ternary_loss2
|
||||
|
||||
|
||||
def initialize_RAFT(model_path='weights/raft-things.pth', device='cuda'):
|
||||
"""Initializes the RAFT model.
|
||||
"""
|
||||
args = argparse.ArgumentParser()
|
||||
args.raft_model = model_path
|
||||
args.small = False
|
||||
args.mixed_precision = False
|
||||
args.alternate_corr = False
|
||||
model = torch.nn.DataParallel(RAFT(args))
|
||||
model.load_state_dict(torch.load(args.raft_model, map_location='cpu'))
|
||||
model = model.module
|
||||
|
||||
model.to(device)
|
||||
|
||||
return model
|
||||
|
||||
|
||||
class RAFT_bi(nn.Module):
|
||||
"""Flow completion loss"""
|
||||
def __init__(self, model_path='weights/raft-things.pth', device='cuda'):
|
||||
super().__init__()
|
||||
self.fix_raft = initialize_RAFT(model_path, device=device)
|
||||
|
||||
for p in self.fix_raft.parameters():
|
||||
p.requires_grad = False
|
||||
|
||||
self.l1_criterion = nn.L1Loss()
|
||||
self.eval()
|
||||
|
||||
def forward(self, gt_local_frames, iters=20):
|
||||
b, l_t, c, h, w = gt_local_frames.size()
|
||||
# print(gt_local_frames.shape)
|
||||
|
||||
with torch.no_grad():
|
||||
gtlf_1 = gt_local_frames[:, :-1, :, :, :].reshape(-1, c, h, w)
|
||||
gtlf_2 = gt_local_frames[:, 1:, :, :, :].reshape(-1, c, h, w)
|
||||
# print(gtlf_1.shape)
|
||||
|
||||
_, gt_flows_forward = self.fix_raft(gtlf_1, gtlf_2, iters=iters, test_mode=True)
|
||||
_, gt_flows_backward = self.fix_raft(gtlf_2, gtlf_1, iters=iters, test_mode=True)
|
||||
|
||||
|
||||
gt_flows_forward = gt_flows_forward.view(b, l_t-1, 2, h, w)
|
||||
gt_flows_backward = gt_flows_backward.view(b, l_t-1, 2, h, w)
|
||||
|
||||
return gt_flows_forward, gt_flows_backward
|
||||
|
||||
|
||||
##################################################################################
|
||||
def smoothness_loss(flow, cmask):
|
||||
delta_u, delta_v, mask = smoothness_deltas(flow)
|
||||
loss_u = charbonnier_loss(delta_u, cmask)
|
||||
loss_v = charbonnier_loss(delta_v, cmask)
|
||||
return loss_u + loss_v
|
||||
|
||||
|
||||
def smoothness_deltas(flow):
|
||||
"""
|
||||
flow: [b, c, h, w]
|
||||
"""
|
||||
mask_x = create_mask(flow, [[0, 0], [0, 1]])
|
||||
mask_y = create_mask(flow, [[0, 1], [0, 0]])
|
||||
mask = torch.cat((mask_x, mask_y), dim=1)
|
||||
mask = mask.to(flow.device)
|
||||
filter_x = torch.tensor([[0, 0, 0.], [0, 1, -1], [0, 0, 0]])
|
||||
filter_y = torch.tensor([[0, 0, 0.], [0, 1, 0], [0, -1, 0]])
|
||||
weights = torch.ones([2, 1, 3, 3])
|
||||
weights[0, 0] = filter_x
|
||||
weights[1, 0] = filter_y
|
||||
weights = weights.to(flow.device)
|
||||
|
||||
flow_u, flow_v = torch.split(flow, split_size_or_sections=1, dim=1)
|
||||
delta_u = F.conv2d(flow_u, weights, stride=1, padding=1)
|
||||
delta_v = F.conv2d(flow_v, weights, stride=1, padding=1)
|
||||
return delta_u, delta_v, mask
|
||||
|
||||
|
||||
def second_order_loss(flow, cmask):
|
||||
delta_u, delta_v, mask = second_order_deltas(flow)
|
||||
loss_u = charbonnier_loss(delta_u, cmask)
|
||||
loss_v = charbonnier_loss(delta_v, cmask)
|
||||
return loss_u + loss_v
|
||||
|
||||
|
||||
def charbonnier_loss(x, mask=None, truncate=None, alpha=0.45, beta=1.0, epsilon=0.001):
|
||||
"""
|
||||
Compute the generalized charbonnier loss of the difference tensor x
|
||||
All positions where mask == 0 are not taken into account
|
||||
x: a tensor of shape [b, c, h, w]
|
||||
mask: a mask of shape [b, mc, h, w], where mask channels must be either 1 or the same as
|
||||
the number of channels of x. Entries should be 0 or 1
|
||||
return: loss
|
||||
"""
|
||||
b, c, h, w = x.shape
|
||||
norm = b * c * h * w
|
||||
error = torch.pow(torch.square(x * beta) + torch.square(torch.tensor(epsilon)), alpha)
|
||||
if mask is not None:
|
||||
error = mask * error
|
||||
if truncate is not None:
|
||||
error = torch.min(error, truncate)
|
||||
return torch.sum(error) / norm
|
||||
|
||||
|
||||
def second_order_deltas(flow):
|
||||
"""
|
||||
consider the single flow first
|
||||
flow shape: [b, c, h, w]
|
||||
"""
|
||||
# create mask
|
||||
mask_x = create_mask(flow, [[0, 0], [1, 1]])
|
||||
mask_y = create_mask(flow, [[1, 1], [0, 0]])
|
||||
mask_diag = create_mask(flow, [[1, 1], [1, 1]])
|
||||
mask = torch.cat((mask_x, mask_y, mask_diag, mask_diag), dim=1)
|
||||
mask = mask.to(flow.device)
|
||||
|
||||
filter_x = torch.tensor([[0, 0, 0.], [1, -2, 1], [0, 0, 0]])
|
||||
filter_y = torch.tensor([[0, 1, 0.], [0, -2, 0], [0, 1, 0]])
|
||||
filter_diag1 = torch.tensor([[1, 0, 0.], [0, -2, 0], [0, 0, 1]])
|
||||
filter_diag2 = torch.tensor([[0, 0, 1.], [0, -2, 0], [1, 0, 0]])
|
||||
weights = torch.ones([4, 1, 3, 3])
|
||||
weights[0] = filter_x
|
||||
weights[1] = filter_y
|
||||
weights[2] = filter_diag1
|
||||
weights[3] = filter_diag2
|
||||
weights = weights.to(flow.device)
|
||||
|
||||
# split the flow into flow_u and flow_v, conv them with the weights
|
||||
flow_u, flow_v = torch.split(flow, split_size_or_sections=1, dim=1)
|
||||
delta_u = F.conv2d(flow_u, weights, stride=1, padding=1)
|
||||
delta_v = F.conv2d(flow_v, weights, stride=1, padding=1)
|
||||
return delta_u, delta_v, mask
|
||||
|
||||
def create_mask(tensor, paddings):
|
||||
"""
|
||||
tensor shape: [b, c, h, w]
|
||||
paddings: [2 x 2] shape list, the first row indicates up and down paddings
|
||||
the second row indicates left and right paddings
|
||||
| |
|
||||
| x |
|
||||
| x * x |
|
||||
| x |
|
||||
| |
|
||||
"""
|
||||
shape = tensor.shape
|
||||
inner_height = shape[2] - (paddings[0][0] + paddings[0][1])
|
||||
inner_width = shape[3] - (paddings[1][0] + paddings[1][1])
|
||||
inner = torch.ones([inner_height, inner_width])
|
||||
torch_paddings = [paddings[1][0], paddings[1][1], paddings[0][0], paddings[0][1]] # left, right, up and down
|
||||
mask2d = F.pad(inner, pad=torch_paddings)
|
||||
mask3d = mask2d.unsqueeze(0).repeat(shape[0], 1, 1)
|
||||
mask4d = mask3d.unsqueeze(1)
|
||||
return mask4d.detach()
|
||||
|
||||
def ternary_loss(flow_comp, flow_gt, mask, current_frame, shift_frame, scale_factor=1):
|
||||
if scale_factor != 1:
|
||||
current_frame = F.interpolate(current_frame, scale_factor=1 / scale_factor, mode='bilinear')
|
||||
shift_frame = F.interpolate(shift_frame, scale_factor=1 / scale_factor, mode='bilinear')
|
||||
warped_sc = flow_warp(shift_frame, flow_gt.permute(0, 2, 3, 1))
|
||||
noc_mask = torch.exp(-50. * torch.sum(torch.abs(current_frame - warped_sc), dim=1).pow(2)).unsqueeze(1)
|
||||
warped_comp_sc = flow_warp(shift_frame, flow_comp.permute(0, 2, 3, 1))
|
||||
loss = ternary_loss2(current_frame, warped_comp_sc, noc_mask, mask)
|
||||
return loss
|
||||
|
||||
class FlowLoss(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.l1_criterion = nn.L1Loss()
|
||||
|
||||
def forward(self, pred_flows, gt_flows, masks, frames):
|
||||
# pred_flows: b t-1 2 h w
|
||||
loss = 0
|
||||
warp_loss = 0
|
||||
h, w = pred_flows[0].shape[-2:]
|
||||
masks = [masks[:,:-1,...].contiguous(), masks[:, 1:, ...].contiguous()]
|
||||
frames0 = frames[:,:-1,...]
|
||||
frames1 = frames[:,1:,...]
|
||||
current_frames = [frames0, frames1]
|
||||
next_frames = [frames1, frames0]
|
||||
for i in range(len(pred_flows)):
|
||||
# print(pred_flows[i].shape)
|
||||
combined_flow = pred_flows[i] * masks[i] + gt_flows[i] * (1-masks[i])
|
||||
l1_loss = self.l1_criterion(pred_flows[i] * masks[i], gt_flows[i] * masks[i]) / torch.mean(masks[i])
|
||||
l1_loss += self.l1_criterion(pred_flows[i] * (1-masks[i]), gt_flows[i] * (1-masks[i])) / torch.mean((1-masks[i]))
|
||||
|
||||
smooth_loss = smoothness_loss(combined_flow.reshape(-1,2,h,w), masks[i].reshape(-1,1,h,w))
|
||||
smooth_loss2 = second_order_loss(combined_flow.reshape(-1,2,h,w), masks[i].reshape(-1,1,h,w))
|
||||
|
||||
warp_loss_i = ternary_loss(combined_flow.reshape(-1,2,h,w), gt_flows[i].reshape(-1,2,h,w),
|
||||
masks[i].reshape(-1,1,h,w), current_frames[i].reshape(-1,3,h,w), next_frames[i].reshape(-1,3,h,w))
|
||||
|
||||
loss += l1_loss + smooth_loss + smooth_loss2
|
||||
|
||||
warp_loss += warp_loss_i
|
||||
|
||||
return loss, warp_loss
|
||||
|
||||
|
||||
def edgeLoss(preds_edges, edges):
|
||||
"""
|
||||
|
||||
Args:
|
||||
preds_edges: with shape [b, c, h , w]
|
||||
edges: with shape [b, c, h, w]
|
||||
|
||||
Returns: Edge losses
|
||||
|
||||
"""
|
||||
mask = (edges > 0.5).float()
|
||||
b, c, h, w = mask.shape
|
||||
num_pos = torch.sum(mask, dim=[1, 2, 3]).float() # Shape: [b,].
|
||||
num_neg = c * h * w - num_pos # Shape: [b,].
|
||||
neg_weights = (num_neg / (num_pos + num_neg)).unsqueeze(1).unsqueeze(2).unsqueeze(3)
|
||||
pos_weights = (num_pos / (num_pos + num_neg)).unsqueeze(1).unsqueeze(2).unsqueeze(3)
|
||||
weight = neg_weights * mask + pos_weights * (1 - mask) # weight for debug
|
||||
losses = F.binary_cross_entropy_with_logits(preds_edges.float(), edges.float(), weight=weight, reduction='none')
|
||||
loss = torch.mean(losses)
|
||||
return loss
|
||||
|
||||
class EdgeLoss(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def forward(self, pred_edges, gt_edges, masks):
|
||||
# pred_flows: b t-1 1 h w
|
||||
loss = 0
|
||||
h, w = pred_edges[0].shape[-2:]
|
||||
masks = [masks[:,:-1,...].contiguous(), masks[:, 1:, ...].contiguous()]
|
||||
for i in range(len(pred_edges)):
|
||||
# print(f'edges_{i}', torch.sum(gt_edges[i])) # debug
|
||||
combined_edge = pred_edges[i] * masks[i] + gt_edges[i] * (1-masks[i])
|
||||
edge_loss = (edgeLoss(pred_edges[i].reshape(-1,1,h,w), gt_edges[i].reshape(-1,1,h,w)) \
|
||||
+ 5 * edgeLoss(combined_edge.reshape(-1,1,h,w), gt_edges[i].reshape(-1,1,h,w)))
|
||||
loss += edge_loss
|
||||
|
||||
return loss
|
||||
|
||||
|
||||
class FlowSimpleLoss(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.l1_criterion = nn.L1Loss()
|
||||
|
||||
def forward(self, pred_flows, gt_flows):
|
||||
# pred_flows: b t-1 2 h w
|
||||
loss = 0
|
||||
h, w = pred_flows[0].shape[-2:]
|
||||
h_orig, w_orig = gt_flows[0].shape[-2:]
|
||||
pred_flows = [f.view(-1, 2, h, w) for f in pred_flows]
|
||||
gt_flows = [f.view(-1, 2, h_orig, w_orig) for f in gt_flows]
|
||||
|
||||
ds_factor = 1.0*h/h_orig
|
||||
gt_flows = [F.interpolate(f, scale_factor=ds_factor, mode='area') * ds_factor for f in gt_flows]
|
||||
for i in range(len(pred_flows)):
|
||||
loss += self.l1_criterion(pred_flows[i], gt_flows[i])
|
||||
|
||||
return loss
|
||||
+142
@@ -0,0 +1,142 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
def flow_warp(x,
|
||||
flow,
|
||||
interpolation='bilinear',
|
||||
padding_mode='zeros',
|
||||
align_corners=True):
|
||||
"""Warp an image or a feature map with optical flow.
|
||||
Args:
|
||||
x (Tensor): Tensor with size (n, c, h, w).
|
||||
flow (Tensor): Tensor with size (n, h, w, 2). The last dimension is
|
||||
a two-channel, denoting the width and height relative offsets.
|
||||
Note that the values are not normalized to [-1, 1].
|
||||
interpolation (str): Interpolation mode: 'nearest' or 'bilinear'.
|
||||
Default: 'bilinear'.
|
||||
padding_mode (str): Padding mode: 'zeros' or 'border' or 'reflection'.
|
||||
Default: 'zeros'.
|
||||
align_corners (bool): Whether align corners. Default: True.
|
||||
Returns:
|
||||
Tensor: Warped image or feature map.
|
||||
"""
|
||||
if x.size()[-2:] != flow.size()[1:3]:
|
||||
raise ValueError(f'The spatial sizes of input ({x.size()[-2:]}) and '
|
||||
f'flow ({flow.size()[1:3]}) are not the same.')
|
||||
_, _, h, w = x.size()
|
||||
# create mesh grid
|
||||
device = flow.device
|
||||
grid_y, grid_x = torch.meshgrid(torch.arange(0, h, device=device), torch.arange(0, w, device=device))
|
||||
grid = torch.stack((grid_x, grid_y), 2).type_as(x) # (w, h, 2)
|
||||
grid.requires_grad = False
|
||||
|
||||
grid_flow = grid + flow
|
||||
# scale grid_flow to [-1,1]
|
||||
grid_flow_x = 2.0 * grid_flow[:, :, :, 0] / max(w - 1, 1) - 1.0
|
||||
grid_flow_y = 2.0 * grid_flow[:, :, :, 1] / max(h - 1, 1) - 1.0
|
||||
grid_flow = torch.stack((grid_flow_x, grid_flow_y), dim=3)
|
||||
output = F.grid_sample(x,
|
||||
grid_flow,
|
||||
mode=interpolation,
|
||||
padding_mode=padding_mode,
|
||||
align_corners=align_corners)
|
||||
return output
|
||||
|
||||
|
||||
# def image_warp(image, flow):
|
||||
# b, c, h, w = image.size()
|
||||
# device = image.device
|
||||
# flow = torch.cat([flow[:, 0:1, :, :] / ((w - 1.0) / 2.0), flow[:, 1:2, :, :] / ((h - 1.0) / 2.0)], dim=1) # normalize to [-1~1](from upper left to lower right
|
||||
# flow = flow.permute(0, 2, 3, 1) # if you wanna use grid_sample function, the channel(band) shape of show must be in the last dimension
|
||||
# x = np.linspace(-1, 1, w)
|
||||
# y = np.linspace(-1, 1, h)
|
||||
# X, Y = np.meshgrid(x, y)
|
||||
# grid = torch.cat((torch.from_numpy(X.astype('float32')).unsqueeze(0).unsqueeze(3),
|
||||
# torch.from_numpy(Y.astype('float32')).unsqueeze(0).unsqueeze(3)), 3).to(device)
|
||||
# output = torch.nn.functional.grid_sample(image, grid + flow, mode='bilinear', padding_mode='zeros')
|
||||
# return output
|
||||
|
||||
|
||||
def length_sq(x):
|
||||
return torch.sum(torch.square(x), dim=1, keepdim=True)
|
||||
|
||||
|
||||
def fbConsistencyCheck(flow_fw, flow_bw, alpha1=0.01, alpha2=0.5):
|
||||
flow_bw_warped = flow_warp(flow_bw, flow_fw.permute(0, 2, 3, 1)) # wb(wf(x))
|
||||
flow_fw_warped = flow_warp(flow_fw, flow_bw.permute(0, 2, 3, 1)) # wf(wb(x))
|
||||
flow_diff_fw = flow_fw + flow_bw_warped # wf + wb(wf(x))
|
||||
flow_diff_bw = flow_bw + flow_fw_warped # wb + wf(wb(x))
|
||||
|
||||
mag_sq_fw = length_sq(flow_fw) + length_sq(flow_bw_warped) # |wf| + |wb(wf(x))|
|
||||
mag_sq_bw = length_sq(flow_bw) + length_sq(flow_fw_warped) # |wb| + |wf(wb(x))|
|
||||
occ_thresh_fw = alpha1 * mag_sq_fw + alpha2
|
||||
occ_thresh_bw = alpha1 * mag_sq_bw + alpha2
|
||||
|
||||
fb_occ_fw = (length_sq(flow_diff_fw) > occ_thresh_fw).float()
|
||||
fb_occ_bw = (length_sq(flow_diff_bw) > occ_thresh_bw).float()
|
||||
|
||||
return fb_occ_fw, fb_occ_bw # fb_occ_fw -> frame2 area occluded by frame1, fb_occ_bw -> frame1 area occluded by frame2
|
||||
|
||||
|
||||
def rgb2gray(image):
|
||||
gray_image = image[:, 0] * 0.299 + image[:, 1] * 0.587 + 0.110 * image[:, 2]
|
||||
gray_image = gray_image.unsqueeze(1)
|
||||
return gray_image
|
||||
|
||||
|
||||
def ternary_transform(image, max_distance=1):
|
||||
device = image.device
|
||||
patch_size = 2 * max_distance + 1
|
||||
intensities = rgb2gray(image) * 255
|
||||
out_channels = patch_size * patch_size
|
||||
w = np.eye(out_channels).reshape(out_channels, 1, patch_size, patch_size)
|
||||
weights = torch.from_numpy(w).float().to(device)
|
||||
patches = F.conv2d(intensities, weights, stride=1, padding=1)
|
||||
transf = patches - intensities
|
||||
transf_norm = transf / torch.sqrt(0.81 + torch.square(transf))
|
||||
return transf_norm
|
||||
|
||||
|
||||
def hamming_distance(t1, t2):
|
||||
dist = torch.square(t1 - t2)
|
||||
dist_norm = dist / (0.1 + dist)
|
||||
dist_sum = torch.sum(dist_norm, dim=1, keepdim=True)
|
||||
return dist_sum
|
||||
|
||||
|
||||
def create_mask(mask, paddings):
|
||||
"""
|
||||
padding: [[top, bottom], [left, right]]
|
||||
"""
|
||||
shape = mask.shape
|
||||
inner_height = shape[2] - (paddings[0][0] + paddings[0][1])
|
||||
inner_width = shape[3] - (paddings[1][0] + paddings[1][1])
|
||||
inner = torch.ones([inner_height, inner_width])
|
||||
|
||||
mask2d = F.pad(inner, pad=[paddings[1][0], paddings[1][1], paddings[0][0], paddings[0][1]])
|
||||
mask3d = mask2d.unsqueeze(0)
|
||||
mask4d = mask3d.unsqueeze(0).repeat(shape[0], 1, 1, 1)
|
||||
return mask4d.detach()
|
||||
|
||||
|
||||
def ternary_loss2(frame1, warp_frame21, confMask, masks, max_distance=1):
|
||||
"""
|
||||
|
||||
Args:
|
||||
frame1: torch tensor, with shape [b * t, c, h, w]
|
||||
warp_frame21: torch tensor, with shape [b * t, c, h, w]
|
||||
confMask: confidence mask, with shape [b * t, c, h, w]
|
||||
masks: torch tensor, with shape [b * t, c, h, w]
|
||||
max_distance: maximum distance.
|
||||
|
||||
Returns: ternary loss
|
||||
|
||||
"""
|
||||
t1 = ternary_transform(frame1)
|
||||
t21 = ternary_transform(warp_frame21)
|
||||
dist = hamming_distance(t1, t21)
|
||||
loss = torch.mean(dist * confMask * masks) / torch.mean(masks)
|
||||
return loss
|
||||
|
||||
@@ -0,0 +1,344 @@
|
||||
import math
|
||||
from functools import reduce
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
class SoftSplit(nn.Module):
|
||||
def __init__(self, channel, hidden, kernel_size, stride, padding):
|
||||
super(SoftSplit, self).__init__()
|
||||
self.kernel_size = kernel_size
|
||||
self.stride = stride
|
||||
self.padding = padding
|
||||
self.t2t = nn.Unfold(kernel_size=kernel_size,
|
||||
stride=stride,
|
||||
padding=padding)
|
||||
c_in = reduce((lambda x, y: x * y), kernel_size) * channel
|
||||
self.embedding = nn.Linear(c_in, hidden)
|
||||
|
||||
def forward(self, x, b, output_size):
|
||||
f_h = int((output_size[0] + 2 * self.padding[0] -
|
||||
(self.kernel_size[0] - 1) - 1) / self.stride[0] + 1)
|
||||
f_w = int((output_size[1] + 2 * self.padding[1] -
|
||||
(self.kernel_size[1] - 1) - 1) / self.stride[1] + 1)
|
||||
|
||||
feat = self.t2t(x)
|
||||
feat = feat.permute(0, 2, 1)
|
||||
# feat shape [b*t, num_vec, ks*ks*c]
|
||||
feat = self.embedding(feat)
|
||||
# feat shape after embedding [b, t*num_vec, hidden]
|
||||
feat = feat.view(b, -1, f_h, f_w, feat.size(2))
|
||||
return feat
|
||||
|
||||
|
||||
class SoftComp(nn.Module):
|
||||
def __init__(self, channel, hidden, kernel_size, stride, padding):
|
||||
super(SoftComp, self).__init__()
|
||||
self.relu = nn.LeakyReLU(0.2, inplace=True)
|
||||
c_out = reduce((lambda x, y: x * y), kernel_size) * channel
|
||||
self.embedding = nn.Linear(hidden, c_out)
|
||||
self.kernel_size = kernel_size
|
||||
self.stride = stride
|
||||
self.padding = padding
|
||||
self.bias_conv = nn.Conv2d(channel,
|
||||
channel,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1)
|
||||
|
||||
def forward(self, x, t, output_size):
|
||||
b_, _, _, _, c_ = x.shape
|
||||
x = x.view(b_, -1, c_)
|
||||
feat = self.embedding(x)
|
||||
b, _, c = feat.size()
|
||||
feat = feat.view(b * t, -1, c).permute(0, 2, 1)
|
||||
feat = F.fold(feat,
|
||||
output_size=output_size,
|
||||
kernel_size=self.kernel_size,
|
||||
stride=self.stride,
|
||||
padding=self.padding)
|
||||
feat = self.bias_conv(feat)
|
||||
return feat
|
||||
|
||||
|
||||
class FusionFeedForward(nn.Module):
|
||||
def __init__(self, dim, hidden_dim=1960, t2t_params=None):
|
||||
super(FusionFeedForward, self).__init__()
|
||||
# We set hidden_dim as a default to 1960
|
||||
self.fc1 = nn.Sequential(nn.Linear(dim, hidden_dim))
|
||||
self.fc2 = nn.Sequential(nn.GELU(), nn.Linear(hidden_dim, dim))
|
||||
assert t2t_params is not None
|
||||
self.t2t_params = t2t_params
|
||||
self.kernel_shape = reduce((lambda x, y: x * y), t2t_params['kernel_size']) # 49
|
||||
|
||||
def forward(self, x, output_size):
|
||||
n_vecs = 1
|
||||
for i, d in enumerate(self.t2t_params['kernel_size']):
|
||||
n_vecs *= int((output_size[i] + 2 * self.t2t_params['padding'][i] -
|
||||
(d - 1) - 1) / self.t2t_params['stride'][i] + 1)
|
||||
|
||||
x = self.fc1(x)
|
||||
b, n, c = x.size()
|
||||
normalizer = x.new_ones(b, n, self.kernel_shape).view(-1, n_vecs, self.kernel_shape).permute(0, 2, 1)
|
||||
normalizer = F.fold(normalizer,
|
||||
output_size=output_size,
|
||||
kernel_size=self.t2t_params['kernel_size'],
|
||||
padding=self.t2t_params['padding'],
|
||||
stride=self.t2t_params['stride'])
|
||||
|
||||
x = F.fold(x.view(-1, n_vecs, c).permute(0, 2, 1),
|
||||
output_size=output_size,
|
||||
kernel_size=self.t2t_params['kernel_size'],
|
||||
padding=self.t2t_params['padding'],
|
||||
stride=self.t2t_params['stride'])
|
||||
|
||||
x = F.unfold(x / normalizer,
|
||||
kernel_size=self.t2t_params['kernel_size'],
|
||||
padding=self.t2t_params['padding'],
|
||||
stride=self.t2t_params['stride']).permute(
|
||||
0, 2, 1).contiguous().view(b, n, c)
|
||||
x = self.fc2(x)
|
||||
return x
|
||||
|
||||
|
||||
def window_partition(x, window_size, n_head):
|
||||
"""
|
||||
Args:
|
||||
x: shape is (B, T, H, W, C)
|
||||
window_size (tuple[int]): window size
|
||||
Returns:
|
||||
windows: (B, num_windows_h, num_windows_w, n_head, T, window_size, window_size, C//n_head)
|
||||
"""
|
||||
B, T, H, W, C = x.shape
|
||||
x = x.view(B, T, H // window_size[0], window_size[0], W // window_size[1], window_size[1], n_head, C//n_head)
|
||||
windows = x.permute(0, 2, 4, 6, 1, 3, 5, 7).contiguous()
|
||||
return windows
|
||||
|
||||
class SparseWindowAttention(nn.Module):
|
||||
def __init__(self, dim, n_head, window_size, pool_size=(4,4), qkv_bias=True, attn_drop=0., proj_drop=0.,
|
||||
pooling_token=True):
|
||||
super().__init__()
|
||||
assert dim % n_head == 0
|
||||
# key, query, value projections for all heads
|
||||
self.key = nn.Linear(dim, dim, qkv_bias)
|
||||
self.query = nn.Linear(dim, dim, qkv_bias)
|
||||
self.value = nn.Linear(dim, dim, qkv_bias)
|
||||
# regularization
|
||||
self.attn_drop = nn.Dropout(attn_drop)
|
||||
self.proj_drop = nn.Dropout(proj_drop)
|
||||
# output projection
|
||||
self.proj = nn.Linear(dim, dim)
|
||||
self.n_head = n_head
|
||||
self.window_size = window_size
|
||||
self.pooling_token = pooling_token
|
||||
if self.pooling_token:
|
||||
ks, stride = pool_size, pool_size
|
||||
self.pool_layer = nn.Conv2d(dim, dim, kernel_size=ks, stride=stride, padding=(0, 0), groups=dim)
|
||||
self.pool_layer.weight.data.fill_(1. / (pool_size[0] * pool_size[1]))
|
||||
self.pool_layer.bias.data.fill_(0)
|
||||
# self.expand_size = tuple(i // 2 for i in window_size)
|
||||
self.expand_size = tuple((i + 1) // 2 for i in window_size)
|
||||
|
||||
if any(i > 0 for i in self.expand_size):
|
||||
# get mask for rolled k and rolled v
|
||||
mask_tl = torch.ones(self.window_size[0], self.window_size[1])
|
||||
mask_tl[:-self.expand_size[0], :-self.expand_size[1]] = 0
|
||||
mask_tr = torch.ones(self.window_size[0], self.window_size[1])
|
||||
mask_tr[:-self.expand_size[0], self.expand_size[1]:] = 0
|
||||
mask_bl = torch.ones(self.window_size[0], self.window_size[1])
|
||||
mask_bl[self.expand_size[0]:, :-self.expand_size[1]] = 0
|
||||
mask_br = torch.ones(self.window_size[0], self.window_size[1])
|
||||
mask_br[self.expand_size[0]:, self.expand_size[1]:] = 0
|
||||
masrool_k = torch.stack((mask_tl, mask_tr, mask_bl, mask_br), 0).flatten(0)
|
||||
self.register_buffer("valid_ind_rolled", masrool_k.nonzero(as_tuple=False).view(-1))
|
||||
|
||||
self.max_pool = nn.MaxPool2d(window_size, window_size, (0, 0))
|
||||
|
||||
|
||||
def forward(self, x, mask=None, T_ind=None, attn_mask=None):
|
||||
b, t, h, w, c = x.shape # 20 36
|
||||
w_h, w_w = self.window_size[0], self.window_size[1]
|
||||
c_head = c // self.n_head
|
||||
n_wh = math.ceil(h / self.window_size[0])
|
||||
n_ww = math.ceil(w / self.window_size[1])
|
||||
new_h = n_wh * self.window_size[0] # 20
|
||||
new_w = n_ww * self.window_size[1] # 36
|
||||
pad_r = new_w - w
|
||||
pad_b = new_h - h
|
||||
# reverse order
|
||||
if pad_r > 0 or pad_b > 0:
|
||||
x = F.pad(x,(0, 0, 0, pad_r, 0, pad_b, 0, 0), mode='constant', value=0)
|
||||
mask = F.pad(mask,(0, 0, 0, pad_r, 0, pad_b, 0, 0), mode='constant', value=0)
|
||||
|
||||
# calculate query, key, values for all heads in batch and move head forward to be the batch dim
|
||||
q = self.query(x)
|
||||
k = self.key(x)
|
||||
v = self.value(x)
|
||||
win_q = window_partition(q.contiguous(), self.window_size, self.n_head).view(b, n_wh*n_ww, self.n_head, t, w_h*w_w, c_head)
|
||||
win_k = window_partition(k.contiguous(), self.window_size, self.n_head).view(b, n_wh*n_ww, self.n_head, t, w_h*w_w, c_head)
|
||||
win_v = window_partition(v.contiguous(), self.window_size, self.n_head).view(b, n_wh*n_ww, self.n_head, t, w_h*w_w, c_head)
|
||||
# roll_k and roll_v
|
||||
if any(i > 0 for i in self.expand_size):
|
||||
(k_tl, v_tl) = map(lambda a: torch.roll(a, shifts=(-self.expand_size[0], -self.expand_size[1]), dims=(2, 3)), (k, v))
|
||||
(k_tr, v_tr) = map(lambda a: torch.roll(a, shifts=(-self.expand_size[0], self.expand_size[1]), dims=(2, 3)), (k, v))
|
||||
(k_bl, v_bl) = map(lambda a: torch.roll(a, shifts=(self.expand_size[0], -self.expand_size[1]), dims=(2, 3)), (k, v))
|
||||
(k_br, v_br) = map(lambda a: torch.roll(a, shifts=(self.expand_size[0], self.expand_size[1]), dims=(2, 3)), (k, v))
|
||||
|
||||
(k_tl_windows, k_tr_windows, k_bl_windows, k_br_windows) = map(
|
||||
lambda a: window_partition(a, self.window_size, self.n_head).view(b, n_wh*n_ww, self.n_head, t, w_h*w_w, c_head),
|
||||
(k_tl, k_tr, k_bl, k_br))
|
||||
(v_tl_windows, v_tr_windows, v_bl_windows, v_br_windows) = map(
|
||||
lambda a: window_partition(a, self.window_size, self.n_head).view(b, n_wh*n_ww, self.n_head, t, w_h*w_w, c_head),
|
||||
(v_tl, v_tr, v_bl, v_br))
|
||||
rool_k = torch.cat((k_tl_windows, k_tr_windows, k_bl_windows, k_br_windows), 4).contiguous()
|
||||
rool_v = torch.cat((v_tl_windows, v_tr_windows, v_bl_windows, v_br_windows), 4).contiguous() # [b, n_wh*n_ww, n_head, t, w_h*w_w, c_head]
|
||||
# mask out tokens in current window
|
||||
rool_k = rool_k[:, :, :, :, self.valid_ind_rolled]
|
||||
rool_v = rool_v[:, :, :, :, self.valid_ind_rolled]
|
||||
roll_N = rool_k.shape[4]
|
||||
rool_k = rool_k.view(b, n_wh*n_ww, self.n_head, t, roll_N, c // self.n_head)
|
||||
rool_v = rool_v.view(b, n_wh*n_ww, self.n_head, t, roll_N, c // self.n_head)
|
||||
win_k = torch.cat((win_k, rool_k), dim=4)
|
||||
win_v = torch.cat((win_v, rool_v), dim=4)
|
||||
else:
|
||||
win_k = win_k
|
||||
win_v = win_v
|
||||
|
||||
# pool_k and pool_v
|
||||
if self.pooling_token:
|
||||
pool_x = self.pool_layer(x.view(b*t, new_h, new_w, c).permute(0,3,1,2))
|
||||
_, _, p_h, p_w = pool_x.shape
|
||||
pool_x = pool_x.permute(0,2,3,1).view(b, t, p_h, p_w, c)
|
||||
# pool_k
|
||||
pool_k = self.key(pool_x).unsqueeze(1).repeat(1, n_wh*n_ww, 1, 1, 1, 1) # [b, n_wh*n_ww, t, p_h, p_w, c]
|
||||
pool_k = pool_k.view(b, n_wh*n_ww, t, p_h, p_w, self.n_head, c_head).permute(0,1,5,2,3,4,6)
|
||||
pool_k = pool_k.contiguous().view(b, n_wh*n_ww, self.n_head, t, p_h*p_w, c_head)
|
||||
win_k = torch.cat((win_k, pool_k), dim=4)
|
||||
# pool_v
|
||||
pool_v = self.value(pool_x).unsqueeze(1).repeat(1, n_wh*n_ww, 1, 1, 1, 1) # [b, n_wh*n_ww, t, p_h, p_w, c]
|
||||
pool_v = pool_v.view(b, n_wh*n_ww, t, p_h, p_w, self.n_head, c_head).permute(0,1,5,2,3,4,6)
|
||||
pool_v = pool_v.contiguous().view(b, n_wh*n_ww, self.n_head, t, p_h*p_w, c_head)
|
||||
win_v = torch.cat((win_v, pool_v), dim=4)
|
||||
|
||||
# [b, n_wh*n_ww, n_head, t, w_h*w_w, c_head]
|
||||
out = torch.zeros_like(win_q)
|
||||
l_t = mask.size(1)
|
||||
|
||||
mask = self.max_pool(mask.view(b * l_t, new_h, new_w))
|
||||
mask = mask.view(b, l_t, n_wh*n_ww)
|
||||
mask = torch.sum(mask, dim=1) # [b, n_wh*n_ww]
|
||||
for i in range(win_q.shape[0]):
|
||||
### For masked windows
|
||||
mask_ind_i = mask[i].nonzero(as_tuple=False).view(-1)
|
||||
# mask out quary in current window
|
||||
# [b, n_wh*n_ww, n_head, t, w_h*w_w, c_head]
|
||||
mask_n = len(mask_ind_i)
|
||||
if mask_n > 0:
|
||||
win_q_t = win_q[i, mask_ind_i].view(mask_n, self.n_head, t*w_h*w_w, c_head)
|
||||
win_k_t = win_k[i, mask_ind_i]
|
||||
win_v_t = win_v[i, mask_ind_i]
|
||||
# mask out key and value
|
||||
if T_ind is not None:
|
||||
# key [n_wh*n_ww, n_head, t, w_h*w_w, c_head]
|
||||
win_k_t = win_k_t[:, :, T_ind.view(-1)].view(mask_n, self.n_head, -1, c_head)
|
||||
# value
|
||||
win_v_t = win_v_t[:, :, T_ind.view(-1)].view(mask_n, self.n_head, -1, c_head)
|
||||
else:
|
||||
win_k_t = win_k_t.view(n_wh*n_ww, self.n_head, t*w_h*w_w, c_head)
|
||||
win_v_t = win_v_t.view(n_wh*n_ww, self.n_head, t*w_h*w_w, c_head)
|
||||
|
||||
att_t = (win_q_t @ win_k_t.transpose(-2, -1)) * (1.0 / math.sqrt(win_q_t.size(-1)))
|
||||
att_t = F.softmax(att_t, dim=-1)
|
||||
att_t = self.attn_drop(att_t)
|
||||
y_t = att_t @ win_v_t
|
||||
|
||||
out[i, mask_ind_i] = y_t.view(-1, self.n_head, t, w_h*w_w, c_head)
|
||||
|
||||
### For unmasked windows
|
||||
unmask_ind_i = (mask[i] == 0).nonzero(as_tuple=False).view(-1)
|
||||
# mask out quary in current window
|
||||
# [b, n_wh*n_ww, n_head, t, w_h*w_w, c_head]
|
||||
win_q_s = win_q[i, unmask_ind_i]
|
||||
win_k_s = win_k[i, unmask_ind_i, :, :, :w_h*w_w]
|
||||
win_v_s = win_v[i, unmask_ind_i, :, :, :w_h*w_w]
|
||||
|
||||
att_s = (win_q_s @ win_k_s.transpose(-2, -1)) * (1.0 / math.sqrt(win_q_s.size(-1)))
|
||||
att_s = F.softmax(att_s, dim=-1)
|
||||
att_s = self.attn_drop(att_s)
|
||||
y_s = att_s @ win_v_s
|
||||
out[i, unmask_ind_i] = y_s
|
||||
|
||||
# re-assemble all head outputs side by side
|
||||
out = out.view(b, n_wh, n_ww, self.n_head, t, w_h, w_w, c_head)
|
||||
out = out.permute(0, 4, 1, 5, 2, 6, 3, 7).contiguous().view(b, t, new_h, new_w, c)
|
||||
|
||||
|
||||
if pad_r > 0 or pad_b > 0:
|
||||
out = out[:, :, :h, :w, :]
|
||||
|
||||
# output projection
|
||||
out = self.proj_drop(self.proj(out))
|
||||
return out
|
||||
|
||||
|
||||
class TemporalSparseTransformer(nn.Module):
|
||||
def __init__(self, dim, n_head, window_size, pool_size,
|
||||
norm_layer=nn.LayerNorm, t2t_params=None):
|
||||
super().__init__()
|
||||
self.window_size = window_size
|
||||
self.attention = SparseWindowAttention(dim, n_head, window_size, pool_size)
|
||||
self.norm1 = norm_layer(dim)
|
||||
self.norm2 = norm_layer(dim)
|
||||
self.mlp = FusionFeedForward(dim, t2t_params=t2t_params)
|
||||
|
||||
def forward(self, x, fold_x_size, mask=None, T_ind=None):
|
||||
"""
|
||||
Args:
|
||||
x: image tokens, shape [B T H W C]
|
||||
fold_x_size: fold feature size, shape [60 108]
|
||||
mask: mask tokens, shape [B T H W 1]
|
||||
Returns:
|
||||
out_tokens: shape [B T H W C]
|
||||
"""
|
||||
B, T, H, W, C = x.shape # 20 36
|
||||
|
||||
shortcut = x
|
||||
x = self.norm1(x)
|
||||
att_x = self.attention(x, mask, T_ind)
|
||||
|
||||
# FFN
|
||||
x = shortcut + att_x
|
||||
y = self.norm2(x)
|
||||
x = x + self.mlp(y.view(B, T * H * W, C), fold_x_size).view(B, T, H, W, C)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class TemporalSparseTransformerBlock(nn.Module):
|
||||
def __init__(self, dim, n_head, window_size, pool_size, depths, t2t_params=None):
|
||||
super().__init__()
|
||||
blocks = []
|
||||
for i in range(depths):
|
||||
blocks.append(
|
||||
TemporalSparseTransformer(dim, n_head, window_size, pool_size, t2t_params=t2t_params)
|
||||
)
|
||||
self.transformer = nn.Sequential(*blocks)
|
||||
self.depths = depths
|
||||
|
||||
def forward(self, x, fold_x_size, l_mask=None, t_dilation=2):
|
||||
"""
|
||||
Args:
|
||||
x: image tokens, shape [B T H W C]
|
||||
fold_x_size: fold feature size, shape [60 108]
|
||||
l_mask: local mask tokens, shape [B T H W 1]
|
||||
Returns:
|
||||
out_tokens: shape [B T H W C]
|
||||
"""
|
||||
assert self.depths % t_dilation == 0, 'wrong t_dilation input.'
|
||||
T = x.size(1)
|
||||
T_ind = [torch.arange(i, T, t_dilation) for i in range(t_dilation)] * (self.depths // t_dilation)
|
||||
|
||||
for i in range(0, self.depths):
|
||||
x = self.transformer[i](x, fold_x_size, l_mask, T_ind[i])
|
||||
|
||||
return x
|
||||
@@ -0,0 +1,288 @@
|
||||
"""
|
||||
Spectral Normalization from https://arxiv.org/abs/1802.05957
|
||||
"""
|
||||
import torch
|
||||
from torch.nn.functional import normalize
|
||||
|
||||
|
||||
class SpectralNorm(object):
|
||||
# Invariant before and after each forward call:
|
||||
# u = normalize(W @ v)
|
||||
# NB: At initialization, this invariant is not enforced
|
||||
|
||||
_version = 1
|
||||
|
||||
# At version 1:
|
||||
# made `W` not a buffer,
|
||||
# added `v` as a buffer, and
|
||||
# made eval mode use `W = u @ W_orig @ v` rather than the stored `W`.
|
||||
|
||||
def __init__(self, name='weight', n_power_iterations=1, dim=0, eps=1e-12):
|
||||
self.name = name
|
||||
self.dim = dim
|
||||
if n_power_iterations <= 0:
|
||||
raise ValueError(
|
||||
'Expected n_power_iterations to be positive, but '
|
||||
'got n_power_iterations={}'.format(n_power_iterations))
|
||||
self.n_power_iterations = n_power_iterations
|
||||
self.eps = eps
|
||||
|
||||
def reshape_weight_to_matrix(self, weight):
|
||||
weight_mat = weight
|
||||
if self.dim != 0:
|
||||
# permute dim to front
|
||||
weight_mat = weight_mat.permute(
|
||||
self.dim,
|
||||
*[d for d in range(weight_mat.dim()) if d != self.dim])
|
||||
height = weight_mat.size(0)
|
||||
return weight_mat.reshape(height, -1)
|
||||
|
||||
def compute_weight(self, module, do_power_iteration):
|
||||
# NB: If `do_power_iteration` is set, the `u` and `v` vectors are
|
||||
# updated in power iteration **in-place**. This is very important
|
||||
# because in `DataParallel` forward, the vectors (being buffers) are
|
||||
# broadcast from the parallelized module to each module replica,
|
||||
# which is a new module object created on the fly. And each replica
|
||||
# runs its own spectral norm power iteration. So simply assigning
|
||||
# the updated vectors to the module this function runs on will cause
|
||||
# the update to be lost forever. And the next time the parallelized
|
||||
# module is replicated, the same randomly initialized vectors are
|
||||
# broadcast and used!
|
||||
#
|
||||
# Therefore, to make the change propagate back, we rely on two
|
||||
# important behaviors (also enforced via tests):
|
||||
# 1. `DataParallel` doesn't clone storage if the broadcast tensor
|
||||
# is already on correct device; and it makes sure that the
|
||||
# parallelized module is already on `device[0]`.
|
||||
# 2. If the out tensor in `out=` kwarg has correct shape, it will
|
||||
# just fill in the values.
|
||||
# Therefore, since the same power iteration is performed on all
|
||||
# devices, simply updating the tensors in-place will make sure that
|
||||
# the module replica on `device[0]` will update the _u vector on the
|
||||
# parallized module (by shared storage).
|
||||
#
|
||||
# However, after we update `u` and `v` in-place, we need to **clone**
|
||||
# them before using them to normalize the weight. This is to support
|
||||
# backproping through two forward passes, e.g., the common pattern in
|
||||
# GAN training: loss = D(real) - D(fake). Otherwise, engine will
|
||||
# complain that variables needed to do backward for the first forward
|
||||
# (i.e., the `u` and `v` vectors) are changed in the second forward.
|
||||
weight = getattr(module, self.name + '_orig')
|
||||
u = getattr(module, self.name + '_u')
|
||||
v = getattr(module, self.name + '_v')
|
||||
weight_mat = self.reshape_weight_to_matrix(weight)
|
||||
|
||||
if do_power_iteration:
|
||||
with torch.no_grad():
|
||||
for _ in range(self.n_power_iterations):
|
||||
# Spectral norm of weight equals to `u^T W v`, where `u` and `v`
|
||||
# are the first left and right singular vectors.
|
||||
# This power iteration produces approximations of `u` and `v`.
|
||||
v = normalize(torch.mv(weight_mat.t(), u),
|
||||
dim=0,
|
||||
eps=self.eps,
|
||||
out=v)
|
||||
u = normalize(torch.mv(weight_mat, v),
|
||||
dim=0,
|
||||
eps=self.eps,
|
||||
out=u)
|
||||
if self.n_power_iterations > 0:
|
||||
# See above on why we need to clone
|
||||
u = u.clone()
|
||||
v = v.clone()
|
||||
|
||||
sigma = torch.dot(u, torch.mv(weight_mat, v))
|
||||
weight = weight / sigma
|
||||
return weight
|
||||
|
||||
def remove(self, module):
|
||||
with torch.no_grad():
|
||||
weight = self.compute_weight(module, do_power_iteration=False)
|
||||
delattr(module, self.name)
|
||||
delattr(module, self.name + '_u')
|
||||
delattr(module, self.name + '_v')
|
||||
delattr(module, self.name + '_orig')
|
||||
module.register_parameter(self.name,
|
||||
torch.nn.Parameter(weight.detach()))
|
||||
|
||||
def __call__(self, module, inputs):
|
||||
setattr(
|
||||
module, self.name,
|
||||
self.compute_weight(module, do_power_iteration=module.training))
|
||||
|
||||
def _solve_v_and_rescale(self, weight_mat, u, target_sigma):
|
||||
# Tries to returns a vector `v` s.t. `u = normalize(W @ v)`
|
||||
# (the invariant at top of this class) and `u @ W @ v = sigma`.
|
||||
# This uses pinverse in case W^T W is not invertible.
|
||||
v = torch.chain_matmul(weight_mat.t().mm(weight_mat).pinverse(),
|
||||
weight_mat.t(), u.unsqueeze(1)).squeeze(1)
|
||||
return v.mul_(target_sigma / torch.dot(u, torch.mv(weight_mat, v)))
|
||||
|
||||
@staticmethod
|
||||
def apply(module, name, n_power_iterations, dim, eps):
|
||||
for k, hook in module._forward_pre_hooks.items():
|
||||
if isinstance(hook, SpectralNorm) and hook.name == name:
|
||||
raise RuntimeError(
|
||||
"Cannot register two spectral_norm hooks on "
|
||||
"the same parameter {}".format(name))
|
||||
|
||||
fn = SpectralNorm(name, n_power_iterations, dim, eps)
|
||||
weight = module._parameters[name]
|
||||
|
||||
with torch.no_grad():
|
||||
weight_mat = fn.reshape_weight_to_matrix(weight)
|
||||
|
||||
h, w = weight_mat.size()
|
||||
# randomly initialize `u` and `v`
|
||||
u = normalize(weight.new_empty(h).normal_(0, 1), dim=0, eps=fn.eps)
|
||||
v = normalize(weight.new_empty(w).normal_(0, 1), dim=0, eps=fn.eps)
|
||||
|
||||
delattr(module, fn.name)
|
||||
module.register_parameter(fn.name + "_orig", weight)
|
||||
# We still need to assign weight back as fn.name because all sorts of
|
||||
# things may assume that it exists, e.g., when initializing weights.
|
||||
# However, we can't directly assign as it could be an nn.Parameter and
|
||||
# gets added as a parameter. Instead, we register weight.data as a plain
|
||||
# attribute.
|
||||
setattr(module, fn.name, weight.data)
|
||||
module.register_buffer(fn.name + "_u", u)
|
||||
module.register_buffer(fn.name + "_v", v)
|
||||
|
||||
module.register_forward_pre_hook(fn)
|
||||
|
||||
module._register_state_dict_hook(SpectralNormStateDictHook(fn))
|
||||
module._register_load_state_dict_pre_hook(
|
||||
SpectralNormLoadStateDictPreHook(fn))
|
||||
return fn
|
||||
|
||||
|
||||
# This is a top level class because Py2 pickle doesn't like inner class nor an
|
||||
# instancemethod.
|
||||
class SpectralNormLoadStateDictPreHook(object):
|
||||
# See docstring of SpectralNorm._version on the changes to spectral_norm.
|
||||
def __init__(self, fn):
|
||||
self.fn = fn
|
||||
|
||||
# For state_dict with version None, (assuming that it has gone through at
|
||||
# least one training forward), we have
|
||||
#
|
||||
# u = normalize(W_orig @ v)
|
||||
# W = W_orig / sigma, where sigma = u @ W_orig @ v
|
||||
#
|
||||
# To compute `v`, we solve `W_orig @ x = u`, and let
|
||||
# v = x / (u @ W_orig @ x) * (W / W_orig).
|
||||
def __call__(self, state_dict, prefix, local_metadata, strict,
|
||||
missing_keys, unexpected_keys, error_msgs):
|
||||
fn = self.fn
|
||||
version = local_metadata.get('spectral_norm',
|
||||
{}).get(fn.name + '.version', None)
|
||||
if version is None or version < 1:
|
||||
with torch.no_grad():
|
||||
weight_orig = state_dict[prefix + fn.name + '_orig']
|
||||
# weight = state_dict.pop(prefix + fn.name)
|
||||
# sigma = (weight_orig / weight).mean()
|
||||
weight_mat = fn.reshape_weight_to_matrix(weight_orig)
|
||||
u = state_dict[prefix + fn.name + '_u']
|
||||
# v = fn._solve_v_and_rescale(weight_mat, u, sigma)
|
||||
# state_dict[prefix + fn.name + '_v'] = v
|
||||
|
||||
|
||||
# This is a top level class because Py2 pickle doesn't like inner class nor an
|
||||
# instancemethod.
|
||||
class SpectralNormStateDictHook(object):
|
||||
# See docstring of SpectralNorm._version on the changes to spectral_norm.
|
||||
def __init__(self, fn):
|
||||
self.fn = fn
|
||||
|
||||
def __call__(self, module, state_dict, prefix, local_metadata):
|
||||
if 'spectral_norm' not in local_metadata:
|
||||
local_metadata['spectral_norm'] = {}
|
||||
key = self.fn.name + '.version'
|
||||
if key in local_metadata['spectral_norm']:
|
||||
raise RuntimeError(
|
||||
"Unexpected key in metadata['spectral_norm']: {}".format(key))
|
||||
local_metadata['spectral_norm'][key] = self.fn._version
|
||||
|
||||
|
||||
def spectral_norm(module,
|
||||
name='weight',
|
||||
n_power_iterations=1,
|
||||
eps=1e-12,
|
||||
dim=None):
|
||||
r"""Applies spectral normalization to a parameter in the given module.
|
||||
|
||||
.. math::
|
||||
\mathbf{W}_{SN} = \dfrac{\mathbf{W}}{\sigma(\mathbf{W})},
|
||||
\sigma(\mathbf{W}) = \max_{\mathbf{h}: \mathbf{h} \ne 0} \dfrac{\|\mathbf{W} \mathbf{h}\|_2}{\|\mathbf{h}\|_2}
|
||||
|
||||
Spectral normalization stabilizes the training of discriminators (critics)
|
||||
in Generative Adversarial Networks (GANs) by rescaling the weight tensor
|
||||
with spectral norm :math:`\sigma` of the weight matrix calculated using
|
||||
power iteration method. If the dimension of the weight tensor is greater
|
||||
than 2, it is reshaped to 2D in power iteration method to get spectral
|
||||
norm. This is implemented via a hook that calculates spectral norm and
|
||||
rescales weight before every :meth:`~Module.forward` call.
|
||||
|
||||
See `Spectral Normalization for Generative Adversarial Networks`_ .
|
||||
|
||||
.. _`Spectral Normalization for Generative Adversarial Networks`: https://arxiv.org/abs/1802.05957
|
||||
|
||||
Args:
|
||||
module (nn.Module): containing module
|
||||
name (str, optional): name of weight parameter
|
||||
n_power_iterations (int, optional): number of power iterations to
|
||||
calculate spectral norm
|
||||
eps (float, optional): epsilon for numerical stability in
|
||||
calculating norms
|
||||
dim (int, optional): dimension corresponding to number of outputs,
|
||||
the default is ``0``, except for modules that are instances of
|
||||
ConvTranspose{1,2,3}d, when it is ``1``
|
||||
|
||||
Returns:
|
||||
The original module with the spectral norm hook
|
||||
|
||||
Example::
|
||||
|
||||
>>> m = spectral_norm(nn.Linear(20, 40))
|
||||
>>> m
|
||||
Linear(in_features=20, out_features=40, bias=True)
|
||||
>>> m.weight_u.size()
|
||||
torch.Size([40])
|
||||
|
||||
"""
|
||||
if dim is None:
|
||||
if isinstance(module,
|
||||
(torch.nn.ConvTranspose1d, torch.nn.ConvTranspose2d,
|
||||
torch.nn.ConvTranspose3d)):
|
||||
dim = 1
|
||||
else:
|
||||
dim = 0
|
||||
SpectralNorm.apply(module, name, n_power_iterations, dim, eps)
|
||||
return module
|
||||
|
||||
|
||||
def remove_spectral_norm(module, name='weight'):
|
||||
r"""Removes the spectral normalization reparameterization from a module.
|
||||
|
||||
Args:
|
||||
module (Module): containing module
|
||||
name (str, optional): name of weight parameter
|
||||
|
||||
Example:
|
||||
>>> m = spectral_norm(nn.Linear(40, 10))
|
||||
>>> remove_spectral_norm(m)
|
||||
"""
|
||||
for k, hook in module._forward_pre_hooks.items():
|
||||
if isinstance(hook, SpectralNorm) and hook.name == name:
|
||||
hook.remove(module)
|
||||
del module._forward_pre_hooks[k]
|
||||
return module
|
||||
|
||||
raise ValueError("spectral_norm of '{}' not found in {}".format(
|
||||
name, module))
|
||||
|
||||
|
||||
def use_spectral_norm(module, use_sn=False):
|
||||
if use_sn:
|
||||
return spectral_norm(module)
|
||||
return module
|
||||
@@ -0,0 +1,539 @@
|
||||
''' Towards An End-to-End Framework for Video Inpainting
|
||||
'''
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torchvision
|
||||
|
||||
from einops import rearrange
|
||||
|
||||
from backend.inpaint.video.model.modules.base_module import BaseNetwork
|
||||
from backend.inpaint.video.model.modules.sparse_transformer import TemporalSparseTransformerBlock, SoftSplit, SoftComp
|
||||
from backend.inpaint.video.model.modules.spectral_norm import spectral_norm as _spectral_norm
|
||||
from backend.inpaint.video.model.modules.flow_loss_utils import flow_warp
|
||||
from backend.inpaint.video.model.modules.deformconv import ModulatedDeformConv2d
|
||||
|
||||
from .misc import constant_init
|
||||
|
||||
|
||||
def length_sq(x):
|
||||
return torch.sum(torch.square(x), dim=1, keepdim=True)
|
||||
|
||||
|
||||
def fbConsistencyCheck(flow_fw, flow_bw, alpha1=0.01, alpha2=0.5):
|
||||
flow_bw_warped = flow_warp(flow_bw, flow_fw.permute(0, 2, 3, 1)) # wb(wf(x))
|
||||
flow_diff_fw = flow_fw + flow_bw_warped # wf + wb(wf(x))
|
||||
|
||||
mag_sq_fw = length_sq(flow_fw) + length_sq(flow_bw_warped) # |wf| + |wb(wf(x))|
|
||||
occ_thresh_fw = alpha1 * mag_sq_fw + alpha2
|
||||
|
||||
# fb_valid_fw = (length_sq(flow_diff_fw) < occ_thresh_fw).float()
|
||||
fb_valid_fw = (length_sq(flow_diff_fw) < occ_thresh_fw).to(flow_fw)
|
||||
return fb_valid_fw
|
||||
|
||||
|
||||
class DeformableAlignment(ModulatedDeformConv2d):
|
||||
"""Second-order deformable alignment module."""
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
# self.max_residue_magnitude = kwargs.pop('max_residue_magnitude', 10)
|
||||
self.max_residue_magnitude = kwargs.pop('max_residue_magnitude', 3)
|
||||
|
||||
super(DeformableAlignment, self).__init__(*args, **kwargs)
|
||||
|
||||
self.conv_offset = nn.Sequential(
|
||||
nn.Conv2d(2 * self.out_channels + 2 + 1 + 2, self.out_channels, 3, 1, 1),
|
||||
nn.LeakyReLU(negative_slope=0.1, inplace=True),
|
||||
nn.Conv2d(self.out_channels, self.out_channels, 3, 1, 1),
|
||||
nn.LeakyReLU(negative_slope=0.1, inplace=True),
|
||||
nn.Conv2d(self.out_channels, self.out_channels, 3, 1, 1),
|
||||
nn.LeakyReLU(negative_slope=0.1, inplace=True),
|
||||
nn.Conv2d(self.out_channels, 27 * self.deform_groups, 3, 1, 1),
|
||||
)
|
||||
self.init_offset()
|
||||
|
||||
def init_offset(self):
|
||||
constant_init(self.conv_offset[-1], val=0, bias=0)
|
||||
|
||||
def forward(self, x, cond_feat, flow):
|
||||
out = self.conv_offset(cond_feat)
|
||||
o1, o2, mask = torch.chunk(out, 3, dim=1)
|
||||
|
||||
# offset
|
||||
offset = self.max_residue_magnitude * torch.tanh(torch.cat((o1, o2), dim=1))
|
||||
offset = offset + flow.flip(1).repeat(1, offset.size(1) // 2, 1, 1)
|
||||
|
||||
# mask
|
||||
mask = torch.sigmoid(mask)
|
||||
|
||||
return torchvision.ops.deform_conv2d(x, offset, self.weight, self.bias,
|
||||
self.stride, self.padding,
|
||||
self.dilation, mask)
|
||||
|
||||
|
||||
class BidirectionalPropagation(nn.Module):
|
||||
def __init__(self, channel, learnable=True):
|
||||
super(BidirectionalPropagation, self).__init__()
|
||||
self.deform_align = nn.ModuleDict()
|
||||
self.backbone = nn.ModuleDict()
|
||||
self.channel = channel
|
||||
self.prop_list = ['backward_1', 'forward_1']
|
||||
self.learnable = learnable
|
||||
|
||||
if self.learnable:
|
||||
for i, module in enumerate(self.prop_list):
|
||||
self.deform_align[module] = DeformableAlignment(
|
||||
channel, channel, 3, padding=1, deform_groups=16)
|
||||
|
||||
self.backbone[module] = nn.Sequential(
|
||||
nn.Conv2d(2 * channel + 2, channel, 3, 1, 1),
|
||||
nn.LeakyReLU(negative_slope=0.2, inplace=True),
|
||||
nn.Conv2d(channel, channel, 3, 1, 1),
|
||||
)
|
||||
|
||||
self.fuse = nn.Sequential(
|
||||
nn.Conv2d(2 * channel + 2, channel, 3, 1, 1),
|
||||
nn.LeakyReLU(negative_slope=0.2, inplace=True),
|
||||
nn.Conv2d(channel, channel, 3, 1, 1),
|
||||
)
|
||||
|
||||
def binary_mask(self, mask, th=0.1):
|
||||
mask[mask > th] = 1
|
||||
mask[mask <= th] = 0
|
||||
# return mask.float()
|
||||
return mask.to(mask)
|
||||
|
||||
def forward(self, x, flows_forward, flows_backward, mask, interpolation='bilinear'):
|
||||
"""
|
||||
x shape : [b, t, c, h, w]
|
||||
return [b, t, c, h, w]
|
||||
"""
|
||||
|
||||
# For backward warping
|
||||
# pred_flows_forward for backward feature propagation
|
||||
# pred_flows_backward for forward feature propagation
|
||||
b, t, c, h, w = x.shape
|
||||
feats, masks = {}, {}
|
||||
feats['input'] = [x[:, i, :, :, :] for i in range(0, t)]
|
||||
masks['input'] = [mask[:, i, :, :, :] for i in range(0, t)]
|
||||
|
||||
prop_list = ['backward_1', 'forward_1']
|
||||
cache_list = ['input'] + prop_list
|
||||
|
||||
for p_i, module_name in enumerate(prop_list):
|
||||
feats[module_name] = []
|
||||
masks[module_name] = []
|
||||
|
||||
if 'backward' in module_name:
|
||||
frame_idx = range(0, t)
|
||||
frame_idx = frame_idx[::-1]
|
||||
flow_idx = frame_idx
|
||||
flows_for_prop = flows_forward
|
||||
flows_for_check = flows_backward
|
||||
else:
|
||||
frame_idx = range(0, t)
|
||||
flow_idx = range(-1, t - 1)
|
||||
flows_for_prop = flows_backward
|
||||
flows_for_check = flows_forward
|
||||
|
||||
for i, idx in enumerate(frame_idx):
|
||||
feat_current = feats[cache_list[p_i]][idx]
|
||||
mask_current = masks[cache_list[p_i]][idx]
|
||||
|
||||
if i == 0:
|
||||
feat_prop = feat_current
|
||||
mask_prop = mask_current
|
||||
else:
|
||||
flow_prop = flows_for_prop[:, flow_idx[i], :, :, :]
|
||||
flow_check = flows_for_check[:, flow_idx[i], :, :, :]
|
||||
flow_vaild_mask = fbConsistencyCheck(flow_prop, flow_check)
|
||||
feat_warped = flow_warp(feat_prop, flow_prop.permute(0, 2, 3, 1), interpolation)
|
||||
|
||||
if self.learnable:
|
||||
cond = torch.cat([feat_current, feat_warped, flow_prop, flow_vaild_mask, mask_current], dim=1)
|
||||
feat_prop = self.deform_align[module_name](feat_prop, cond, flow_prop)
|
||||
mask_prop = mask_current
|
||||
else:
|
||||
mask_prop_valid = flow_warp(mask_prop, flow_prop.permute(0, 2, 3, 1))
|
||||
mask_prop_valid = self.binary_mask(mask_prop_valid)
|
||||
|
||||
union_vaild_mask = self.binary_mask(mask_current * flow_vaild_mask * (1 - mask_prop_valid))
|
||||
feat_prop = union_vaild_mask * feat_warped + (1 - union_vaild_mask) * feat_current
|
||||
# update mask
|
||||
mask_prop = self.binary_mask(mask_current * (1 - (flow_vaild_mask * (1 - mask_prop_valid))))
|
||||
|
||||
# refine
|
||||
if self.learnable:
|
||||
feat = torch.cat([feat_current, feat_prop, mask_current], dim=1)
|
||||
feat_prop = feat_prop + self.backbone[module_name](feat)
|
||||
# feat_prop = self.backbone[module_name](feat_prop)
|
||||
|
||||
feats[module_name].append(feat_prop)
|
||||
masks[module_name].append(mask_prop)
|
||||
|
||||
# end for
|
||||
if 'backward' in module_name:
|
||||
feats[module_name] = feats[module_name][::-1]
|
||||
masks[module_name] = masks[module_name][::-1]
|
||||
|
||||
outputs_b = torch.stack(feats['backward_1'], dim=1).view(-1, c, h, w)
|
||||
outputs_f = torch.stack(feats['forward_1'], dim=1).view(-1, c, h, w)
|
||||
|
||||
if self.learnable:
|
||||
mask_in = mask.view(-1, 2, h, w)
|
||||
masks_b, masks_f = None, None
|
||||
outputs = self.fuse(torch.cat([outputs_b, outputs_f, mask_in], dim=1)) + x.view(-1, c, h, w)
|
||||
else:
|
||||
masks_b = torch.stack(masks['backward_1'], dim=1)
|
||||
masks_f = torch.stack(masks['forward_1'], dim=1)
|
||||
outputs = outputs_f
|
||||
|
||||
return outputs_b.view(b, -1, c, h, w), outputs_f.view(b, -1, c, h, w), \
|
||||
outputs.view(b, -1, c, h, w), masks_f
|
||||
|
||||
|
||||
class Encoder(nn.Module):
|
||||
def __init__(self):
|
||||
super(Encoder, self).__init__()
|
||||
self.group = [1, 2, 4, 8, 1]
|
||||
self.layers = nn.ModuleList([
|
||||
nn.Conv2d(5, 64, kernel_size=3, stride=2, padding=1),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
nn.Conv2d(256, 384, kernel_size=3, stride=1, padding=1, groups=1),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
nn.Conv2d(640, 512, kernel_size=3, stride=1, padding=1, groups=2),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
nn.Conv2d(768, 384, kernel_size=3, stride=1, padding=1, groups=4),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
nn.Conv2d(640, 256, kernel_size=3, stride=1, padding=1, groups=8),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
nn.Conv2d(512, 128, kernel_size=3, stride=1, padding=1, groups=1),
|
||||
nn.LeakyReLU(0.2, inplace=True)
|
||||
])
|
||||
|
||||
def forward(self, x):
|
||||
bt, c, _, _ = x.size()
|
||||
# h, w = h//4, w//4
|
||||
out = x
|
||||
for i, layer in enumerate(self.layers):
|
||||
if i == 8:
|
||||
x0 = out
|
||||
_, _, h, w = x0.size()
|
||||
if i > 8 and i % 2 == 0:
|
||||
g = self.group[(i - 8) // 2]
|
||||
x = x0.view(bt, g, -1, h, w)
|
||||
o = out.view(bt, g, -1, h, w)
|
||||
out = torch.cat([x, o], 2).view(bt, -1, h, w)
|
||||
out = layer(out)
|
||||
return out
|
||||
|
||||
|
||||
class deconv(nn.Module):
|
||||
def __init__(self,
|
||||
input_channel,
|
||||
output_channel,
|
||||
kernel_size=3,
|
||||
padding=0):
|
||||
super().__init__()
|
||||
self.conv = nn.Conv2d(input_channel,
|
||||
output_channel,
|
||||
kernel_size=kernel_size,
|
||||
stride=1,
|
||||
padding=padding)
|
||||
|
||||
def forward(self, x):
|
||||
x = F.interpolate(x,
|
||||
scale_factor=2,
|
||||
mode='bilinear',
|
||||
align_corners=True)
|
||||
return self.conv(x)
|
||||
|
||||
|
||||
class InpaintGenerator(BaseNetwork):
|
||||
def __init__(self, init_weights=True, model_path=None):
|
||||
super(InpaintGenerator, self).__init__()
|
||||
channel = 128
|
||||
hidden = 512
|
||||
|
||||
# encoder
|
||||
self.encoder = Encoder()
|
||||
|
||||
# decoder
|
||||
self.decoder = nn.Sequential(
|
||||
deconv(channel, 128, kernel_size=3, padding=1),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
nn.Conv2d(128, 64, kernel_size=3, stride=1, padding=1),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
deconv(64, 64, kernel_size=3, padding=1),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
nn.Conv2d(64, 3, kernel_size=3, stride=1, padding=1))
|
||||
|
||||
# soft split and soft composition
|
||||
kernel_size = (7, 7)
|
||||
padding = (3, 3)
|
||||
stride = (3, 3)
|
||||
t2t_params = {
|
||||
'kernel_size': kernel_size,
|
||||
'stride': stride,
|
||||
'padding': padding
|
||||
}
|
||||
self.ss = SoftSplit(channel, hidden, kernel_size, stride, padding)
|
||||
self.sc = SoftComp(channel, hidden, kernel_size, stride, padding)
|
||||
self.max_pool = nn.MaxPool2d(kernel_size, stride, padding)
|
||||
|
||||
# feature propagation module
|
||||
self.img_prop_module = BidirectionalPropagation(3, learnable=False)
|
||||
self.feat_prop_module = BidirectionalPropagation(128, learnable=True)
|
||||
|
||||
depths = 8
|
||||
num_heads = 4
|
||||
window_size = (5, 9)
|
||||
pool_size = (4, 4)
|
||||
self.transformers = TemporalSparseTransformerBlock(dim=hidden,
|
||||
n_head=num_heads,
|
||||
window_size=window_size,
|
||||
pool_size=pool_size,
|
||||
depths=depths,
|
||||
t2t_params=t2t_params)
|
||||
if init_weights:
|
||||
self.init_weights()
|
||||
|
||||
if model_path is not None:
|
||||
print('Pretrained ProPainter has loaded...')
|
||||
ckpt = torch.load(model_path, map_location='cpu')
|
||||
self.load_state_dict(ckpt, strict=True)
|
||||
|
||||
# print network parameter number
|
||||
self.print_network()
|
||||
|
||||
def img_propagation(self, masked_frames, completed_flows, masks, interpolation='nearest'):
|
||||
_, _, prop_frames, updated_masks = self.img_prop_module(masked_frames, completed_flows[0], completed_flows[1],
|
||||
masks, interpolation)
|
||||
return prop_frames, updated_masks
|
||||
|
||||
def forward(self, masked_frames, completed_flows, masks_in, masks_updated, num_local_frames,
|
||||
interpolation='bilinear', t_dilation=2):
|
||||
"""
|
||||
Args:
|
||||
masks_in: original mask
|
||||
masks_updated: updated mask after image propagation
|
||||
"""
|
||||
|
||||
l_t = num_local_frames
|
||||
b, t, _, ori_h, ori_w = masked_frames.size()
|
||||
|
||||
# extracting features
|
||||
enc_feat = self.encoder(torch.cat([masked_frames.view(b * t, 3, ori_h, ori_w),
|
||||
masks_in.view(b * t, 1, ori_h, ori_w),
|
||||
masks_updated.view(b * t, 1, ori_h, ori_w)], dim=1))
|
||||
_, c, h, w = enc_feat.size()
|
||||
local_feat = enc_feat.view(b, t, c, h, w)[:, :l_t, ...]
|
||||
ref_feat = enc_feat.view(b, t, c, h, w)[:, l_t:, ...]
|
||||
fold_feat_size = (h, w)
|
||||
|
||||
ds_flows_f = F.interpolate(completed_flows[0].view(-1, 2, ori_h, ori_w), scale_factor=1 / 4, mode='bilinear',
|
||||
align_corners=False).view(b, l_t - 1, 2, h, w) / 4.0
|
||||
ds_flows_b = F.interpolate(completed_flows[1].view(-1, 2, ori_h, ori_w), scale_factor=1 / 4, mode='bilinear',
|
||||
align_corners=False).view(b, l_t - 1, 2, h, w) / 4.0
|
||||
ds_mask_in = F.interpolate(masks_in.reshape(-1, 1, ori_h, ori_w), scale_factor=1 / 4, mode='nearest').view(b, t,
|
||||
1, h,
|
||||
w)
|
||||
ds_mask_in_local = ds_mask_in[:, :l_t]
|
||||
ds_mask_updated_local = F.interpolate(masks_updated[:, :l_t].reshape(-1, 1, ori_h, ori_w), scale_factor=1 / 4,
|
||||
mode='nearest').view(b, l_t, 1, h, w)
|
||||
|
||||
if self.training:
|
||||
mask_pool_l = self.max_pool(ds_mask_in.view(-1, 1, h, w))
|
||||
mask_pool_l = mask_pool_l.view(b, t, 1, mask_pool_l.size(-2), mask_pool_l.size(-1))
|
||||
else:
|
||||
mask_pool_l = self.max_pool(ds_mask_in_local.view(-1, 1, h, w))
|
||||
mask_pool_l = mask_pool_l.view(b, l_t, 1, mask_pool_l.size(-2), mask_pool_l.size(-1))
|
||||
|
||||
prop_mask_in = torch.cat([ds_mask_in_local, ds_mask_updated_local], dim=2)
|
||||
_, _, local_feat, _ = self.feat_prop_module(local_feat, ds_flows_f, ds_flows_b, prop_mask_in, interpolation)
|
||||
enc_feat = torch.cat((local_feat, ref_feat), dim=1)
|
||||
|
||||
trans_feat = self.ss(enc_feat.view(-1, c, h, w), b, fold_feat_size)
|
||||
mask_pool_l = rearrange(mask_pool_l, 'b t c h w -> b t h w c').contiguous()
|
||||
trans_feat = self.transformers(trans_feat, fold_feat_size, mask_pool_l, t_dilation=t_dilation)
|
||||
trans_feat = self.sc(trans_feat, t, fold_feat_size)
|
||||
trans_feat = trans_feat.view(b, t, -1, h, w)
|
||||
|
||||
enc_feat = enc_feat + trans_feat
|
||||
|
||||
if self.training:
|
||||
output = self.decoder(enc_feat.view(-1, c, h, w))
|
||||
output = torch.tanh(output).view(b, t, 3, ori_h, ori_w)
|
||||
else:
|
||||
output = self.decoder(enc_feat[:, :l_t].view(-1, c, h, w))
|
||||
output = torch.tanh(output).view(b, l_t, 3, ori_h, ori_w)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
# ######################################################################
|
||||
# Discriminator for Temporal Patch GAN
|
||||
# ######################################################################
|
||||
class Discriminator(BaseNetwork):
|
||||
def __init__(self,
|
||||
in_channels=3,
|
||||
use_sigmoid=False,
|
||||
use_spectral_norm=True,
|
||||
init_weights=True):
|
||||
super(Discriminator, self).__init__()
|
||||
self.use_sigmoid = use_sigmoid
|
||||
nf = 32
|
||||
|
||||
self.conv = nn.Sequential(
|
||||
spectral_norm(
|
||||
nn.Conv3d(in_channels=in_channels,
|
||||
out_channels=nf * 1,
|
||||
kernel_size=(3, 5, 5),
|
||||
stride=(1, 2, 2),
|
||||
padding=1,
|
||||
bias=not use_spectral_norm), use_spectral_norm),
|
||||
# nn.InstanceNorm2d(64, track_running_stats=False),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
spectral_norm(
|
||||
nn.Conv3d(nf * 1,
|
||||
nf * 2,
|
||||
kernel_size=(3, 5, 5),
|
||||
stride=(1, 2, 2),
|
||||
padding=(1, 2, 2),
|
||||
bias=not use_spectral_norm), use_spectral_norm),
|
||||
# nn.InstanceNorm2d(128, track_running_stats=False),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
spectral_norm(
|
||||
nn.Conv3d(nf * 2,
|
||||
nf * 4,
|
||||
kernel_size=(3, 5, 5),
|
||||
stride=(1, 2, 2),
|
||||
padding=(1, 2, 2),
|
||||
bias=not use_spectral_norm), use_spectral_norm),
|
||||
# nn.InstanceNorm2d(256, track_running_stats=False),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
spectral_norm(
|
||||
nn.Conv3d(nf * 4,
|
||||
nf * 4,
|
||||
kernel_size=(3, 5, 5),
|
||||
stride=(1, 2, 2),
|
||||
padding=(1, 2, 2),
|
||||
bias=not use_spectral_norm), use_spectral_norm),
|
||||
# nn.InstanceNorm2d(256, track_running_stats=False),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
spectral_norm(
|
||||
nn.Conv3d(nf * 4,
|
||||
nf * 4,
|
||||
kernel_size=(3, 5, 5),
|
||||
stride=(1, 2, 2),
|
||||
padding=(1, 2, 2),
|
||||
bias=not use_spectral_norm), use_spectral_norm),
|
||||
# nn.InstanceNorm2d(256, track_running_stats=False),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
nn.Conv3d(nf * 4,
|
||||
nf * 4,
|
||||
kernel_size=(3, 5, 5),
|
||||
stride=(1, 2, 2),
|
||||
padding=(1, 2, 2)))
|
||||
|
||||
if init_weights:
|
||||
self.init_weights()
|
||||
|
||||
def forward(self, xs):
|
||||
# T, C, H, W = xs.shape (old)
|
||||
# B, T, C, H, W (new)
|
||||
xs_t = torch.transpose(xs, 1, 2)
|
||||
feat = self.conv(xs_t)
|
||||
if self.use_sigmoid:
|
||||
feat = torch.sigmoid(feat)
|
||||
out = torch.transpose(feat, 1, 2) # B, T, C, H, W
|
||||
return out
|
||||
|
||||
|
||||
class Discriminator_2D(BaseNetwork):
|
||||
def __init__(self,
|
||||
in_channels=3,
|
||||
use_sigmoid=False,
|
||||
use_spectral_norm=True,
|
||||
init_weights=True):
|
||||
super(Discriminator_2D, self).__init__()
|
||||
self.use_sigmoid = use_sigmoid
|
||||
nf = 32
|
||||
|
||||
self.conv = nn.Sequential(
|
||||
spectral_norm(
|
||||
nn.Conv3d(in_channels=in_channels,
|
||||
out_channels=nf * 1,
|
||||
kernel_size=(1, 5, 5),
|
||||
stride=(1, 2, 2),
|
||||
padding=(0, 2, 2),
|
||||
bias=not use_spectral_norm), use_spectral_norm),
|
||||
# nn.InstanceNorm2d(64, track_running_stats=False),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
spectral_norm(
|
||||
nn.Conv3d(nf * 1,
|
||||
nf * 2,
|
||||
kernel_size=(1, 5, 5),
|
||||
stride=(1, 2, 2),
|
||||
padding=(0, 2, 2),
|
||||
bias=not use_spectral_norm), use_spectral_norm),
|
||||
# nn.InstanceNorm2d(128, track_running_stats=False),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
spectral_norm(
|
||||
nn.Conv3d(nf * 2,
|
||||
nf * 4,
|
||||
kernel_size=(1, 5, 5),
|
||||
stride=(1, 2, 2),
|
||||
padding=(0, 2, 2),
|
||||
bias=not use_spectral_norm), use_spectral_norm),
|
||||
# nn.InstanceNorm2d(256, track_running_stats=False),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
spectral_norm(
|
||||
nn.Conv3d(nf * 4,
|
||||
nf * 4,
|
||||
kernel_size=(1, 5, 5),
|
||||
stride=(1, 2, 2),
|
||||
padding=(0, 2, 2),
|
||||
bias=not use_spectral_norm), use_spectral_norm),
|
||||
# nn.InstanceNorm2d(256, track_running_stats=False),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
spectral_norm(
|
||||
nn.Conv3d(nf * 4,
|
||||
nf * 4,
|
||||
kernel_size=(1, 5, 5),
|
||||
stride=(1, 2, 2),
|
||||
padding=(0, 2, 2),
|
||||
bias=not use_spectral_norm), use_spectral_norm),
|
||||
# nn.InstanceNorm2d(256, track_running_stats=False),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
nn.Conv3d(nf * 4,
|
||||
nf * 4,
|
||||
kernel_size=(1, 5, 5),
|
||||
stride=(1, 2, 2),
|
||||
padding=(0, 2, 2)))
|
||||
|
||||
if init_weights:
|
||||
self.init_weights()
|
||||
|
||||
def forward(self, xs):
|
||||
# T, C, H, W = xs.shape (old)
|
||||
# B, T, C, H, W (new)
|
||||
xs_t = torch.transpose(xs, 1, 2)
|
||||
feat = self.conv(xs_t)
|
||||
if self.use_sigmoid:
|
||||
feat = torch.sigmoid(feat)
|
||||
out = torch.transpose(feat, 1, 2) # B, T, C, H, W
|
||||
return out
|
||||
|
||||
|
||||
def spectral_norm(module, mode=True):
|
||||
if mode:
|
||||
return _spectral_norm(module)
|
||||
return module
|
||||
@@ -0,0 +1,348 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torchvision
|
||||
|
||||
from backend.inpaint.video.model.modules.deformconv import ModulatedDeformConv2d
|
||||
from .misc import constant_init
|
||||
|
||||
|
||||
class SecondOrderDeformableAlignment(ModulatedDeformConv2d):
|
||||
"""Second-order deformable alignment module."""
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
self.max_residue_magnitude = kwargs.pop('max_residue_magnitude', 5)
|
||||
|
||||
super(SecondOrderDeformableAlignment, self).__init__(*args, **kwargs)
|
||||
|
||||
self.conv_offset = nn.Sequential(
|
||||
nn.Conv2d(3 * self.out_channels, self.out_channels, 3, 1, 1),
|
||||
nn.LeakyReLU(negative_slope=0.1, inplace=True),
|
||||
nn.Conv2d(self.out_channels, self.out_channels, 3, 1, 1),
|
||||
nn.LeakyReLU(negative_slope=0.1, inplace=True),
|
||||
nn.Conv2d(self.out_channels, self.out_channels, 3, 1, 1),
|
||||
nn.LeakyReLU(negative_slope=0.1, inplace=True),
|
||||
nn.Conv2d(self.out_channels, 27 * self.deform_groups, 3, 1, 1),
|
||||
)
|
||||
self.init_offset()
|
||||
|
||||
def init_offset(self):
|
||||
constant_init(self.conv_offset[-1], val=0, bias=0)
|
||||
|
||||
def forward(self, x, extra_feat):
|
||||
out = self.conv_offset(extra_feat)
|
||||
o1, o2, mask = torch.chunk(out, 3, dim=1)
|
||||
|
||||
# offset
|
||||
offset = self.max_residue_magnitude * torch.tanh(torch.cat((o1, o2), dim=1))
|
||||
offset_1, offset_2 = torch.chunk(offset, 2, dim=1)
|
||||
offset = torch.cat([offset_1, offset_2], dim=1)
|
||||
|
||||
# mask
|
||||
mask = torch.sigmoid(mask)
|
||||
|
||||
return torchvision.ops.deform_conv2d(x, offset, self.weight, self.bias,
|
||||
self.stride, self.padding,
|
||||
self.dilation, mask)
|
||||
|
||||
|
||||
class BidirectionalPropagation(nn.Module):
|
||||
def __init__(self, channel):
|
||||
super(BidirectionalPropagation, self).__init__()
|
||||
modules = ['backward_', 'forward_']
|
||||
self.deform_align = nn.ModuleDict()
|
||||
self.backbone = nn.ModuleDict()
|
||||
self.channel = channel
|
||||
|
||||
for i, module in enumerate(modules):
|
||||
self.deform_align[module] = SecondOrderDeformableAlignment(
|
||||
2 * channel, channel, 3, padding=1, deform_groups=16)
|
||||
|
||||
self.backbone[module] = nn.Sequential(
|
||||
nn.Conv2d((2 + i) * channel, channel, 3, 1, 1),
|
||||
nn.LeakyReLU(negative_slope=0.1, inplace=True),
|
||||
nn.Conv2d(channel, channel, 3, 1, 1),
|
||||
)
|
||||
|
||||
self.fusion = nn.Conv2d(2 * channel, channel, 1, 1, 0)
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
x shape : [b, t, c, h, w]
|
||||
return [b, t, c, h, w]
|
||||
"""
|
||||
b, t, c, h, w = x.shape
|
||||
feats = {}
|
||||
feats['spatial'] = [x[:, i, :, :, :] for i in range(0, t)]
|
||||
|
||||
for module_name in ['backward_', 'forward_']:
|
||||
|
||||
feats[module_name] = []
|
||||
|
||||
frame_idx = range(0, t)
|
||||
mapping_idx = list(range(0, len(feats['spatial'])))
|
||||
mapping_idx += mapping_idx[::-1]
|
||||
|
||||
if 'backward' in module_name:
|
||||
frame_idx = frame_idx[::-1]
|
||||
|
||||
feat_prop = x.new_zeros(b, self.channel, h, w)
|
||||
for i, idx in enumerate(frame_idx):
|
||||
feat_current = feats['spatial'][mapping_idx[idx]]
|
||||
if i > 0:
|
||||
cond_n1 = feat_prop
|
||||
|
||||
# initialize second-order features
|
||||
feat_n2 = torch.zeros_like(feat_prop)
|
||||
cond_n2 = torch.zeros_like(cond_n1)
|
||||
if i > 1: # second-order features
|
||||
feat_n2 = feats[module_name][-2]
|
||||
cond_n2 = feat_n2
|
||||
|
||||
cond = torch.cat([cond_n1, feat_current, cond_n2],
|
||||
dim=1) # condition information, cond(flow warped 1st/2nd feature)
|
||||
feat_prop = torch.cat([feat_prop, feat_n2], dim=1) # two order feat_prop -1 & -2
|
||||
feat_prop = self.deform_align[module_name](feat_prop, cond)
|
||||
|
||||
# fuse current features
|
||||
feat = [feat_current] + \
|
||||
[feats[k][idx] for k in feats if k not in ['spatial', module_name]] \
|
||||
+ [feat_prop]
|
||||
|
||||
feat = torch.cat(feat, dim=1)
|
||||
# embed current features
|
||||
feat_prop = feat_prop + self.backbone[module_name](feat)
|
||||
|
||||
feats[module_name].append(feat_prop)
|
||||
|
||||
# end for
|
||||
if 'backward' in module_name:
|
||||
feats[module_name] = feats[module_name][::-1]
|
||||
|
||||
outputs = []
|
||||
for i in range(0, t):
|
||||
align_feats = [feats[k].pop(0) for k in feats if k != 'spatial']
|
||||
align_feats = torch.cat(align_feats, dim=1)
|
||||
outputs.append(self.fusion(align_feats))
|
||||
|
||||
return torch.stack(outputs, dim=1) + x
|
||||
|
||||
|
||||
class deconv(nn.Module):
|
||||
def __init__(self,
|
||||
input_channel,
|
||||
output_channel,
|
||||
kernel_size=3,
|
||||
padding=0):
|
||||
super().__init__()
|
||||
self.conv = nn.Conv2d(input_channel,
|
||||
output_channel,
|
||||
kernel_size=kernel_size,
|
||||
stride=1,
|
||||
padding=padding)
|
||||
|
||||
def forward(self, x):
|
||||
x = F.interpolate(x,
|
||||
scale_factor=2,
|
||||
mode='bilinear',
|
||||
align_corners=True)
|
||||
return self.conv(x)
|
||||
|
||||
|
||||
class P3DBlock(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, kernel_size, stride, padding, use_residual=0, bias=True):
|
||||
super().__init__()
|
||||
self.conv1 = nn.Sequential(
|
||||
nn.Conv3d(in_channels, out_channels, kernel_size=(1, kernel_size, kernel_size),
|
||||
stride=(1, stride, stride), padding=(0, padding, padding), bias=bias),
|
||||
nn.LeakyReLU(0.2, inplace=True)
|
||||
)
|
||||
self.conv2 = nn.Sequential(
|
||||
nn.Conv3d(out_channels, out_channels, kernel_size=(3, 1, 1), stride=(1, 1, 1),
|
||||
padding=(2, 0, 0), dilation=(2, 1, 1), bias=bias)
|
||||
)
|
||||
self.use_residual = use_residual
|
||||
|
||||
def forward(self, feats):
|
||||
feat1 = self.conv1(feats)
|
||||
feat2 = self.conv2(feat1)
|
||||
if self.use_residual:
|
||||
output = feats + feat2
|
||||
else:
|
||||
output = feat2
|
||||
return output
|
||||
|
||||
|
||||
class EdgeDetection(nn.Module):
|
||||
def __init__(self, in_ch=2, out_ch=1, mid_ch=16):
|
||||
super().__init__()
|
||||
self.projection = nn.Sequential(
|
||||
nn.Conv2d(in_ch, mid_ch, 3, 1, 1),
|
||||
nn.LeakyReLU(0.2, inplace=True)
|
||||
)
|
||||
|
||||
self.mid_layer_1 = nn.Sequential(
|
||||
nn.Conv2d(mid_ch, mid_ch, 3, 1, 1),
|
||||
nn.LeakyReLU(0.2, inplace=True)
|
||||
)
|
||||
|
||||
self.mid_layer_2 = nn.Sequential(
|
||||
nn.Conv2d(mid_ch, mid_ch, 3, 1, 1)
|
||||
)
|
||||
|
||||
self.l_relu = nn.LeakyReLU(0.01, inplace=True)
|
||||
|
||||
self.out_layer = nn.Conv2d(mid_ch, out_ch, 1, 1, 0)
|
||||
|
||||
def forward(self, flow):
|
||||
flow = self.projection(flow)
|
||||
edge = self.mid_layer_1(flow)
|
||||
edge = self.mid_layer_2(edge)
|
||||
edge = self.l_relu(flow + edge)
|
||||
edge = self.out_layer(edge)
|
||||
edge = torch.sigmoid(edge)
|
||||
return edge
|
||||
|
||||
|
||||
class RecurrentFlowCompleteNet(nn.Module):
|
||||
def __init__(self, model_path=None):
|
||||
super().__init__()
|
||||
self.downsample = nn.Sequential(
|
||||
nn.Conv3d(3, 32, kernel_size=(1, 5, 5), stride=(1, 2, 2),
|
||||
padding=(0, 2, 2), padding_mode='replicate'),
|
||||
nn.LeakyReLU(0.2, inplace=True)
|
||||
)
|
||||
|
||||
self.encoder1 = nn.Sequential(
|
||||
P3DBlock(32, 32, 3, 1, 1),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
P3DBlock(32, 64, 3, 2, 1),
|
||||
nn.LeakyReLU(0.2, inplace=True)
|
||||
) # 4x
|
||||
|
||||
self.encoder2 = nn.Sequential(
|
||||
P3DBlock(64, 64, 3, 1, 1),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
P3DBlock(64, 128, 3, 2, 1),
|
||||
nn.LeakyReLU(0.2, inplace=True)
|
||||
) # 8x
|
||||
|
||||
self.mid_dilation = nn.Sequential(
|
||||
nn.Conv3d(128, 128, (1, 3, 3), (1, 1, 1), padding=(0, 3, 3), dilation=(1, 3, 3)), # p = d*(k-1)/2
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
nn.Conv3d(128, 128, (1, 3, 3), (1, 1, 1), padding=(0, 2, 2), dilation=(1, 2, 2)),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
nn.Conv3d(128, 128, (1, 3, 3), (1, 1, 1), padding=(0, 1, 1), dilation=(1, 1, 1)),
|
||||
nn.LeakyReLU(0.2, inplace=True)
|
||||
)
|
||||
|
||||
# feature propagation module
|
||||
self.feat_prop_module = BidirectionalPropagation(128)
|
||||
|
||||
self.decoder2 = nn.Sequential(
|
||||
nn.Conv2d(128, 128, 3, 1, 1),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
deconv(128, 64, 3, 1),
|
||||
nn.LeakyReLU(0.2, inplace=True)
|
||||
) # 4x
|
||||
|
||||
self.decoder1 = nn.Sequential(
|
||||
nn.Conv2d(64, 64, 3, 1, 1),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
deconv(64, 32, 3, 1),
|
||||
nn.LeakyReLU(0.2, inplace=True)
|
||||
) # 2x
|
||||
|
||||
self.upsample = nn.Sequential(
|
||||
nn.Conv2d(32, 32, 3, padding=1),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
deconv(32, 2, 3, 1)
|
||||
)
|
||||
|
||||
# edge loss
|
||||
self.edgeDetector = EdgeDetection(in_ch=2, out_ch=1, mid_ch=16)
|
||||
|
||||
# Need to initial the weights of MSDeformAttn specifically
|
||||
for m in self.modules():
|
||||
if isinstance(m, SecondOrderDeformableAlignment):
|
||||
m.init_offset()
|
||||
|
||||
if model_path is not None:
|
||||
print('Pretrained flow completion model has loaded...')
|
||||
ckpt = torch.load(model_path, map_location='cpu')
|
||||
self.load_state_dict(ckpt, strict=True)
|
||||
|
||||
def forward(self, masked_flows, masks):
|
||||
# masked_flows: b t-1 2 h w
|
||||
# masks: b t-1 2 h w
|
||||
b, t, _, h, w = masked_flows.size()
|
||||
masked_flows = masked_flows.permute(0, 2, 1, 3, 4)
|
||||
masks = masks.permute(0, 2, 1, 3, 4)
|
||||
|
||||
inputs = torch.cat((masked_flows, masks), dim=1)
|
||||
|
||||
x = self.downsample(inputs)
|
||||
|
||||
feat_e1 = self.encoder1(x)
|
||||
feat_e2 = self.encoder2(feat_e1) # b c t h w
|
||||
feat_mid = self.mid_dilation(feat_e2) # b c t h w
|
||||
feat_mid = feat_mid.permute(0, 2, 1, 3, 4) # b t c h w
|
||||
|
||||
feat_prop = self.feat_prop_module(feat_mid)
|
||||
feat_prop = feat_prop.view(-1, 128, h // 8, w // 8) # b*t c h w
|
||||
|
||||
_, c, _, h_f, w_f = feat_e1.shape
|
||||
feat_e1 = feat_e1.permute(0, 2, 1, 3, 4).contiguous().view(-1, c, h_f, w_f) # b*t c h w
|
||||
feat_d2 = self.decoder2(feat_prop) + feat_e1
|
||||
|
||||
_, c, _, h_f, w_f = x.shape
|
||||
x = x.permute(0, 2, 1, 3, 4).contiguous().view(-1, c, h_f, w_f) # b*t c h w
|
||||
|
||||
feat_d1 = self.decoder1(feat_d2)
|
||||
|
||||
flow = self.upsample(feat_d1)
|
||||
if self.training:
|
||||
edge = self.edgeDetector(flow)
|
||||
edge = edge.view(b, t, 1, h, w)
|
||||
else:
|
||||
edge = None
|
||||
|
||||
flow = flow.view(b, t, 2, h, w)
|
||||
|
||||
return flow, edge
|
||||
|
||||
def forward_bidirect_flow(self, masked_flows_bi, masks):
|
||||
"""
|
||||
Args:
|
||||
masked_flows_bi: [masked_flows_f, masked_flows_b] | (b t-1 2 h w), (b t-1 2 h w)
|
||||
masks: b t 1 h w
|
||||
"""
|
||||
masks_forward = masks[:, :-1, ...].contiguous()
|
||||
masks_backward = masks[:, 1:, ...].contiguous()
|
||||
|
||||
# mask flow
|
||||
masked_flows_forward = masked_flows_bi[0] * (1 - masks_forward)
|
||||
masked_flows_backward = masked_flows_bi[1] * (1 - masks_backward)
|
||||
|
||||
# -- completion --
|
||||
# forward
|
||||
pred_flows_forward, pred_edges_forward = self.forward(masked_flows_forward, masks_forward)
|
||||
|
||||
# backward
|
||||
masked_flows_backward = torch.flip(masked_flows_backward, dims=[1])
|
||||
masks_backward = torch.flip(masks_backward, dims=[1])
|
||||
pred_flows_backward, pred_edges_backward = self.forward(masked_flows_backward, masks_backward)
|
||||
pred_flows_backward = torch.flip(pred_flows_backward, dims=[1])
|
||||
if self.training:
|
||||
pred_edges_backward = torch.flip(pred_edges_backward, dims=[1])
|
||||
|
||||
return [pred_flows_forward, pred_flows_backward], [pred_edges_forward, pred_edges_backward]
|
||||
|
||||
def combine_flow(self, masked_flows_bi, pred_flows_bi, masks):
|
||||
masks_forward = masks[:, :-1, ...].contiguous()
|
||||
masks_backward = masks[:, 1:, ...].contiguous()
|
||||
|
||||
pred_flows_forward = pred_flows_bi[0] * masks_forward + masked_flows_bi[0] * (1 - masks_forward)
|
||||
pred_flows_backward = pred_flows_bi[1] * masks_backward + masked_flows_bi[1] * (1 - masks_backward)
|
||||
|
||||
return pred_flows_forward, pred_flows_backward
|
||||
@@ -0,0 +1,157 @@
|
||||
import os
|
||||
import torch
|
||||
from collections import OrderedDict
|
||||
from torch import nn as nn
|
||||
from torchvision.models import vgg as vgg
|
||||
|
||||
VGG_PRETRAIN_PATH = 'experiments/pretrained_models/vgg19-dcbb9e9d.pth'
|
||||
NAMES = {
|
||||
'vgg11': [
|
||||
'conv1_1', 'relu1_1', 'pool1', 'conv2_1', 'relu2_1', 'pool2', 'conv3_1', 'relu3_1', 'conv3_2', 'relu3_2',
|
||||
'pool3', 'conv4_1', 'relu4_1', 'conv4_2', 'relu4_2', 'pool4', 'conv5_1', 'relu5_1', 'conv5_2', 'relu5_2',
|
||||
'pool5'
|
||||
],
|
||||
'vgg13': [
|
||||
'conv1_1', 'relu1_1', 'conv1_2', 'relu1_2', 'pool1', 'conv2_1', 'relu2_1', 'conv2_2', 'relu2_2', 'pool2',
|
||||
'conv3_1', 'relu3_1', 'conv3_2', 'relu3_2', 'pool3', 'conv4_1', 'relu4_1', 'conv4_2', 'relu4_2', 'pool4',
|
||||
'conv5_1', 'relu5_1', 'conv5_2', 'relu5_2', 'pool5'
|
||||
],
|
||||
'vgg16': [
|
||||
'conv1_1', 'relu1_1', 'conv1_2', 'relu1_2', 'pool1', 'conv2_1', 'relu2_1', 'conv2_2', 'relu2_2', 'pool2',
|
||||
'conv3_1', 'relu3_1', 'conv3_2', 'relu3_2', 'conv3_3', 'relu3_3', 'pool3', 'conv4_1', 'relu4_1', 'conv4_2',
|
||||
'relu4_2', 'conv4_3', 'relu4_3', 'pool4', 'conv5_1', 'relu5_1', 'conv5_2', 'relu5_2', 'conv5_3', 'relu5_3',
|
||||
'pool5'
|
||||
],
|
||||
'vgg19': [
|
||||
'conv1_1', 'relu1_1', 'conv1_2', 'relu1_2', 'pool1', 'conv2_1', 'relu2_1', 'conv2_2', 'relu2_2', 'pool2',
|
||||
'conv3_1', 'relu3_1', 'conv3_2', 'relu3_2', 'conv3_3', 'relu3_3', 'conv3_4', 'relu3_4', 'pool3', 'conv4_1',
|
||||
'relu4_1', 'conv4_2', 'relu4_2', 'conv4_3', 'relu4_3', 'conv4_4', 'relu4_4', 'pool4', 'conv5_1', 'relu5_1',
|
||||
'conv5_2', 'relu5_2', 'conv5_3', 'relu5_3', 'conv5_4', 'relu5_4', 'pool5'
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
def insert_bn(names):
|
||||
"""Insert bn layer after each conv.
|
||||
|
||||
Args:
|
||||
names (list): The list of layer names.
|
||||
|
||||
Returns:
|
||||
list: The list of layer names with bn layers.
|
||||
"""
|
||||
names_bn = []
|
||||
for name in names:
|
||||
names_bn.append(name)
|
||||
if 'conv' in name:
|
||||
position = name.replace('conv', '')
|
||||
names_bn.append('bn' + position)
|
||||
return names_bn
|
||||
|
||||
class VGGFeatureExtractor(nn.Module):
|
||||
"""VGG network for feature extraction.
|
||||
|
||||
In this implementation, we allow users to choose whether use normalization
|
||||
in the input feature and the type of vgg network. Note that the pretrained
|
||||
path must fit the vgg type.
|
||||
|
||||
Args:
|
||||
layer_name_list (list[str]): Forward function returns the corresponding
|
||||
features according to the layer_name_list.
|
||||
Example: {'relu1_1', 'relu2_1', 'relu3_1'}.
|
||||
vgg_type (str): Set the type of vgg network. Default: 'vgg19'.
|
||||
use_input_norm (bool): If True, normalize the input image. Importantly,
|
||||
the input feature must in the range [0, 1]. Default: True.
|
||||
range_norm (bool): If True, norm images with range [-1, 1] to [0, 1].
|
||||
Default: False.
|
||||
requires_grad (bool): If true, the parameters of VGG network will be
|
||||
optimized. Default: False.
|
||||
remove_pooling (bool): If true, the max pooling operations in VGG net
|
||||
will be removed. Default: False.
|
||||
pooling_stride (int): The stride of max pooling operation. Default: 2.
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
layer_name_list,
|
||||
vgg_type='vgg19',
|
||||
use_input_norm=True,
|
||||
range_norm=False,
|
||||
requires_grad=False,
|
||||
remove_pooling=False,
|
||||
pooling_stride=2):
|
||||
super(VGGFeatureExtractor, self).__init__()
|
||||
|
||||
self.layer_name_list = layer_name_list
|
||||
self.use_input_norm = use_input_norm
|
||||
self.range_norm = range_norm
|
||||
|
||||
self.names = NAMES[vgg_type.replace('_bn', '')]
|
||||
if 'bn' in vgg_type:
|
||||
self.names = insert_bn(self.names)
|
||||
|
||||
# only borrow layers that will be used to avoid unused params
|
||||
max_idx = 0
|
||||
for v in layer_name_list:
|
||||
idx = self.names.index(v)
|
||||
if idx > max_idx:
|
||||
max_idx = idx
|
||||
|
||||
if os.path.exists(VGG_PRETRAIN_PATH):
|
||||
vgg_net = getattr(vgg, vgg_type)(pretrained=False)
|
||||
state_dict = torch.load(VGG_PRETRAIN_PATH, map_location=lambda storage, loc: storage)
|
||||
vgg_net.load_state_dict(state_dict)
|
||||
else:
|
||||
vgg_net = getattr(vgg, vgg_type)(pretrained=True)
|
||||
|
||||
features = vgg_net.features[:max_idx + 1]
|
||||
|
||||
modified_net = OrderedDict()
|
||||
for k, v in zip(self.names, features):
|
||||
if 'pool' in k:
|
||||
# if remove_pooling is true, pooling operation will be removed
|
||||
if remove_pooling:
|
||||
continue
|
||||
else:
|
||||
# in some cases, we may want to change the default stride
|
||||
modified_net[k] = nn.MaxPool2d(kernel_size=2, stride=pooling_stride)
|
||||
else:
|
||||
modified_net[k] = v
|
||||
|
||||
self.vgg_net = nn.Sequential(modified_net)
|
||||
|
||||
if not requires_grad:
|
||||
self.vgg_net.eval()
|
||||
for param in self.parameters():
|
||||
param.requires_grad = False
|
||||
else:
|
||||
self.vgg_net.train()
|
||||
for param in self.parameters():
|
||||
param.requires_grad = True
|
||||
|
||||
if self.use_input_norm:
|
||||
# the mean is for image with range [0, 1]
|
||||
self.register_buffer('mean', torch.Tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1))
|
||||
# the std is for image with range [0, 1]
|
||||
self.register_buffer('std', torch.Tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1))
|
||||
|
||||
def forward(self, x):
|
||||
"""Forward function.
|
||||
|
||||
Args:
|
||||
x (Tensor): Input tensor with shape (n, c, h, w).
|
||||
|
||||
Returns:
|
||||
Tensor: Forward results.
|
||||
"""
|
||||
if self.range_norm:
|
||||
x = (x + 1) / 2
|
||||
if self.use_input_norm:
|
||||
x = (x - self.mean) / self.std
|
||||
output = {}
|
||||
|
||||
for key, layer in self.vgg_net._modules.items():
|
||||
x = layer(x)
|
||||
if key in self.layer_name_list:
|
||||
output[key] = x.clone()
|
||||
|
||||
return output
|
||||
Executable
+2
@@ -0,0 +1,2 @@
|
||||
# from .demo import RAFT_infer
|
||||
from .raft import RAFT
|
||||
Executable
+91
@@ -0,0 +1,91 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from .utils.utils import bilinear_sampler, coords_grid
|
||||
|
||||
try:
|
||||
import alt_cuda_corr
|
||||
except:
|
||||
# alt_cuda_corr is not compiled
|
||||
pass
|
||||
|
||||
|
||||
class CorrBlock:
|
||||
def __init__(self, fmap1, fmap2, num_levels=4, radius=4):
|
||||
self.num_levels = num_levels
|
||||
self.radius = radius
|
||||
self.corr_pyramid = []
|
||||
|
||||
# all pairs correlation
|
||||
corr = CorrBlock.corr(fmap1, fmap2)
|
||||
|
||||
batch, h1, w1, dim, h2, w2 = corr.shape
|
||||
corr = corr.reshape(batch*h1*w1, dim, h2, w2)
|
||||
|
||||
self.corr_pyramid.append(corr)
|
||||
for i in range(self.num_levels-1):
|
||||
corr = F.avg_pool2d(corr, 2, stride=2)
|
||||
self.corr_pyramid.append(corr)
|
||||
|
||||
def __call__(self, coords):
|
||||
r = self.radius
|
||||
coords = coords.permute(0, 2, 3, 1)
|
||||
batch, h1, w1, _ = coords.shape
|
||||
|
||||
out_pyramid = []
|
||||
for i in range(self.num_levels):
|
||||
corr = self.corr_pyramid[i]
|
||||
dx = torch.linspace(-r, r, 2*r+1)
|
||||
dy = torch.linspace(-r, r, 2*r+1)
|
||||
delta = torch.stack(torch.meshgrid(dy, dx), axis=-1).to(coords.device)
|
||||
|
||||
centroid_lvl = coords.reshape(batch*h1*w1, 1, 1, 2) / 2**i
|
||||
delta_lvl = delta.view(1, 2*r+1, 2*r+1, 2)
|
||||
coords_lvl = centroid_lvl + delta_lvl
|
||||
|
||||
corr = bilinear_sampler(corr, coords_lvl)
|
||||
corr = corr.view(batch, h1, w1, -1)
|
||||
out_pyramid.append(corr)
|
||||
|
||||
out = torch.cat(out_pyramid, dim=-1)
|
||||
return out.permute(0, 3, 1, 2).contiguous().float()
|
||||
|
||||
@staticmethod
|
||||
def corr(fmap1, fmap2):
|
||||
batch, dim, ht, wd = fmap1.shape
|
||||
fmap1 = fmap1.view(batch, dim, ht*wd)
|
||||
fmap2 = fmap2.view(batch, dim, ht*wd)
|
||||
|
||||
corr = torch.matmul(fmap1.transpose(1,2), fmap2)
|
||||
corr = corr.view(batch, ht, wd, 1, ht, wd)
|
||||
return corr / torch.sqrt(torch.tensor(dim).float())
|
||||
|
||||
|
||||
class AlternateCorrBlock:
|
||||
def __init__(self, fmap1, fmap2, num_levels=4, radius=4):
|
||||
self.num_levels = num_levels
|
||||
self.radius = radius
|
||||
|
||||
self.pyramid = [(fmap1, fmap2)]
|
||||
for i in range(self.num_levels):
|
||||
fmap1 = F.avg_pool2d(fmap1, 2, stride=2)
|
||||
fmap2 = F.avg_pool2d(fmap2, 2, stride=2)
|
||||
self.pyramid.append((fmap1, fmap2))
|
||||
|
||||
def __call__(self, coords):
|
||||
|
||||
coords = coords.permute(0, 2, 3, 1)
|
||||
B, H, W, _ = coords.shape
|
||||
|
||||
corr_list = []
|
||||
for i in range(self.num_levels):
|
||||
r = self.radius
|
||||
fmap1_i = self.pyramid[0][0].permute(0, 2, 3, 1)
|
||||
fmap2_i = self.pyramid[i][1].permute(0, 2, 3, 1)
|
||||
|
||||
coords_i = (coords / 2**i).reshape(B, 1, H, W, 2).contiguous()
|
||||
corr = alt_cuda_corr(fmap1_i, fmap2_i, coords_i, r)
|
||||
corr_list.append(corr.squeeze(1))
|
||||
|
||||
corr = torch.stack(corr_list, dim=1)
|
||||
corr = corr.reshape(B, -1, H, W)
|
||||
return corr / 16.0
|
||||
Executable
+235
@@ -0,0 +1,235 @@
|
||||
# Data loading based on https://github.com/NVIDIA/flownet2-pytorch
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.utils.data as data
|
||||
import torch.nn.functional as F
|
||||
|
||||
import os
|
||||
import math
|
||||
import random
|
||||
from glob import glob
|
||||
import os.path as osp
|
||||
|
||||
from utils import frame_utils
|
||||
from utils.augmentor import FlowAugmentor, SparseFlowAugmentor
|
||||
|
||||
|
||||
class FlowDataset(data.Dataset):
|
||||
def __init__(self, aug_params=None, sparse=False):
|
||||
self.augmentor = None
|
||||
self.sparse = sparse
|
||||
if aug_params is not None:
|
||||
if sparse:
|
||||
self.augmentor = SparseFlowAugmentor(**aug_params)
|
||||
else:
|
||||
self.augmentor = FlowAugmentor(**aug_params)
|
||||
|
||||
self.is_test = False
|
||||
self.init_seed = False
|
||||
self.flow_list = []
|
||||
self.image_list = []
|
||||
self.extra_info = []
|
||||
|
||||
def __getitem__(self, index):
|
||||
|
||||
if self.is_test:
|
||||
img1 = frame_utils.read_gen(self.image_list[index][0])
|
||||
img2 = frame_utils.read_gen(self.image_list[index][1])
|
||||
img1 = np.array(img1).astype(np.uint8)[..., :3]
|
||||
img2 = np.array(img2).astype(np.uint8)[..., :3]
|
||||
img1 = torch.from_numpy(img1).permute(2, 0, 1).float()
|
||||
img2 = torch.from_numpy(img2).permute(2, 0, 1).float()
|
||||
return img1, img2, self.extra_info[index]
|
||||
|
||||
if not self.init_seed:
|
||||
worker_info = torch.utils.data.get_worker_info()
|
||||
if worker_info is not None:
|
||||
torch.manual_seed(worker_info.id)
|
||||
np.random.seed(worker_info.id)
|
||||
random.seed(worker_info.id)
|
||||
self.init_seed = True
|
||||
|
||||
index = index % len(self.image_list)
|
||||
valid = None
|
||||
if self.sparse:
|
||||
flow, valid = frame_utils.readFlowKITTI(self.flow_list[index])
|
||||
else:
|
||||
flow = frame_utils.read_gen(self.flow_list[index])
|
||||
|
||||
img1 = frame_utils.read_gen(self.image_list[index][0])
|
||||
img2 = frame_utils.read_gen(self.image_list[index][1])
|
||||
|
||||
flow = np.array(flow).astype(np.float32)
|
||||
img1 = np.array(img1).astype(np.uint8)
|
||||
img2 = np.array(img2).astype(np.uint8)
|
||||
|
||||
# grayscale images
|
||||
if len(img1.shape) == 2:
|
||||
img1 = np.tile(img1[...,None], (1, 1, 3))
|
||||
img2 = np.tile(img2[...,None], (1, 1, 3))
|
||||
else:
|
||||
img1 = img1[..., :3]
|
||||
img2 = img2[..., :3]
|
||||
|
||||
if self.augmentor is not None:
|
||||
if self.sparse:
|
||||
img1, img2, flow, valid = self.augmentor(img1, img2, flow, valid)
|
||||
else:
|
||||
img1, img2, flow = self.augmentor(img1, img2, flow)
|
||||
|
||||
img1 = torch.from_numpy(img1).permute(2, 0, 1).float()
|
||||
img2 = torch.from_numpy(img2).permute(2, 0, 1).float()
|
||||
flow = torch.from_numpy(flow).permute(2, 0, 1).float()
|
||||
|
||||
if valid is not None:
|
||||
valid = torch.from_numpy(valid)
|
||||
else:
|
||||
valid = (flow[0].abs() < 1000) & (flow[1].abs() < 1000)
|
||||
|
||||
return img1, img2, flow, valid.float()
|
||||
|
||||
|
||||
def __rmul__(self, v):
|
||||
self.flow_list = v * self.flow_list
|
||||
self.image_list = v * self.image_list
|
||||
return self
|
||||
|
||||
def __len__(self):
|
||||
return len(self.image_list)
|
||||
|
||||
|
||||
class MpiSintel(FlowDataset):
|
||||
def __init__(self, aug_params=None, split='training', root='datasets/Sintel', dstype='clean'):
|
||||
super(MpiSintel, self).__init__(aug_params)
|
||||
flow_root = osp.join(root, split, 'flow')
|
||||
image_root = osp.join(root, split, dstype)
|
||||
|
||||
if split == 'test':
|
||||
self.is_test = True
|
||||
|
||||
for scene in os.listdir(image_root):
|
||||
image_list = sorted(glob(osp.join(image_root, scene, '*.png')))
|
||||
for i in range(len(image_list)-1):
|
||||
self.image_list += [ [image_list[i], image_list[i+1]] ]
|
||||
self.extra_info += [ (scene, i) ] # scene and frame_id
|
||||
|
||||
if split != 'test':
|
||||
self.flow_list += sorted(glob(osp.join(flow_root, scene, '*.flo')))
|
||||
|
||||
|
||||
class FlyingChairs(FlowDataset):
|
||||
def __init__(self, aug_params=None, split='train', root='datasets/FlyingChairs_release/data'):
|
||||
super(FlyingChairs, self).__init__(aug_params)
|
||||
|
||||
images = sorted(glob(osp.join(root, '*.ppm')))
|
||||
flows = sorted(glob(osp.join(root, '*.flo')))
|
||||
assert (len(images)//2 == len(flows))
|
||||
|
||||
split_list = np.loadtxt('chairs_split.txt', dtype=np.int32)
|
||||
for i in range(len(flows)):
|
||||
xid = split_list[i]
|
||||
if (split=='training' and xid==1) or (split=='validation' and xid==2):
|
||||
self.flow_list += [ flows[i] ]
|
||||
self.image_list += [ [images[2*i], images[2*i+1]] ]
|
||||
|
||||
|
||||
class FlyingThings3D(FlowDataset):
|
||||
def __init__(self, aug_params=None, root='datasets/FlyingThings3D', dstype='frames_cleanpass'):
|
||||
super(FlyingThings3D, self).__init__(aug_params)
|
||||
|
||||
for cam in ['left']:
|
||||
for direction in ['into_future', 'into_past']:
|
||||
image_dirs = sorted(glob(osp.join(root, dstype, 'TRAIN/*/*')))
|
||||
image_dirs = sorted([osp.join(f, cam) for f in image_dirs])
|
||||
|
||||
flow_dirs = sorted(glob(osp.join(root, 'optical_flow/TRAIN/*/*')))
|
||||
flow_dirs = sorted([osp.join(f, direction, cam) for f in flow_dirs])
|
||||
|
||||
for idir, fdir in zip(image_dirs, flow_dirs):
|
||||
images = sorted(glob(osp.join(idir, '*.png')) )
|
||||
flows = sorted(glob(osp.join(fdir, '*.pfm')) )
|
||||
for i in range(len(flows)-1):
|
||||
if direction == 'into_future':
|
||||
self.image_list += [ [images[i], images[i+1]] ]
|
||||
self.flow_list += [ flows[i] ]
|
||||
elif direction == 'into_past':
|
||||
self.image_list += [ [images[i+1], images[i]] ]
|
||||
self.flow_list += [ flows[i+1] ]
|
||||
|
||||
|
||||
class KITTI(FlowDataset):
|
||||
def __init__(self, aug_params=None, split='training', root='datasets/KITTI'):
|
||||
super(KITTI, self).__init__(aug_params, sparse=True)
|
||||
if split == 'testing':
|
||||
self.is_test = True
|
||||
|
||||
root = osp.join(root, split)
|
||||
images1 = sorted(glob(osp.join(root, 'image_2/*_10.png')))
|
||||
images2 = sorted(glob(osp.join(root, 'image_2/*_11.png')))
|
||||
|
||||
for img1, img2 in zip(images1, images2):
|
||||
frame_id = img1.split('/')[-1]
|
||||
self.extra_info += [ [frame_id] ]
|
||||
self.image_list += [ [img1, img2] ]
|
||||
|
||||
if split == 'training':
|
||||
self.flow_list = sorted(glob(osp.join(root, 'flow_occ/*_10.png')))
|
||||
|
||||
|
||||
class HD1K(FlowDataset):
|
||||
def __init__(self, aug_params=None, root='datasets/HD1k'):
|
||||
super(HD1K, self).__init__(aug_params, sparse=True)
|
||||
|
||||
seq_ix = 0
|
||||
while 1:
|
||||
flows = sorted(glob(os.path.join(root, 'hd1k_flow_gt', 'flow_occ/%06d_*.png' % seq_ix)))
|
||||
images = sorted(glob(os.path.join(root, 'hd1k_input', 'image_2/%06d_*.png' % seq_ix)))
|
||||
|
||||
if len(flows) == 0:
|
||||
break
|
||||
|
||||
for i in range(len(flows)-1):
|
||||
self.flow_list += [flows[i]]
|
||||
self.image_list += [ [images[i], images[i+1]] ]
|
||||
|
||||
seq_ix += 1
|
||||
|
||||
|
||||
def fetch_dataloader(args, TRAIN_DS='C+T+K+S+H'):
|
||||
""" Create the data loader for the corresponding trainign set """
|
||||
|
||||
if args.stage == 'chairs':
|
||||
aug_params = {'crop_size': args.image_size, 'min_scale': -0.1, 'max_scale': 1.0, 'do_flip': True}
|
||||
train_dataset = FlyingChairs(aug_params, split='training')
|
||||
|
||||
elif args.stage == 'things':
|
||||
aug_params = {'crop_size': args.image_size, 'min_scale': -0.4, 'max_scale': 0.8, 'do_flip': True}
|
||||
clean_dataset = FlyingThings3D(aug_params, dstype='frames_cleanpass')
|
||||
final_dataset = FlyingThings3D(aug_params, dstype='frames_finalpass')
|
||||
train_dataset = clean_dataset + final_dataset
|
||||
|
||||
elif args.stage == 'sintel':
|
||||
aug_params = {'crop_size': args.image_size, 'min_scale': -0.2, 'max_scale': 0.6, 'do_flip': True}
|
||||
things = FlyingThings3D(aug_params, dstype='frames_cleanpass')
|
||||
sintel_clean = MpiSintel(aug_params, split='training', dstype='clean')
|
||||
sintel_final = MpiSintel(aug_params, split='training', dstype='final')
|
||||
|
||||
if TRAIN_DS == 'C+T+K+S+H':
|
||||
kitti = KITTI({'crop_size': args.image_size, 'min_scale': -0.3, 'max_scale': 0.5, 'do_flip': True})
|
||||
hd1k = HD1K({'crop_size': args.image_size, 'min_scale': -0.5, 'max_scale': 0.2, 'do_flip': True})
|
||||
train_dataset = 100*sintel_clean + 100*sintel_final + 200*kitti + 5*hd1k + things
|
||||
|
||||
elif TRAIN_DS == 'C+T+K/S':
|
||||
train_dataset = 100*sintel_clean + 100*sintel_final + things
|
||||
|
||||
elif args.stage == 'kitti':
|
||||
aug_params = {'crop_size': args.image_size, 'min_scale': -0.2, 'max_scale': 0.4, 'do_flip': False}
|
||||
train_dataset = KITTI(aug_params, split='training')
|
||||
|
||||
train_loader = data.DataLoader(train_dataset, batch_size=args.batch_size,
|
||||
pin_memory=False, shuffle=True, num_workers=4, drop_last=True)
|
||||
|
||||
print('Training with %d image pairs' % len(train_dataset))
|
||||
return train_loader
|
||||
|
||||
Executable
+79
@@ -0,0 +1,79 @@
|
||||
import sys
|
||||
import argparse
|
||||
import os
|
||||
import cv2
|
||||
import glob
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from .raft import RAFT
|
||||
from .utils import flow_viz
|
||||
from .utils.utils import InputPadder
|
||||
|
||||
|
||||
|
||||
DEVICE = 'cuda'
|
||||
|
||||
def load_image(imfile):
|
||||
img = np.array(Image.open(imfile)).astype(np.uint8)
|
||||
img = torch.from_numpy(img).permute(2, 0, 1).float()
|
||||
return img
|
||||
|
||||
|
||||
def load_image_list(image_files):
|
||||
images = []
|
||||
for imfile in sorted(image_files):
|
||||
images.append(load_image(imfile))
|
||||
|
||||
images = torch.stack(images, dim=0)
|
||||
images = images.to(DEVICE)
|
||||
|
||||
padder = InputPadder(images.shape)
|
||||
return padder.pad(images)[0]
|
||||
|
||||
|
||||
def viz(img, flo):
|
||||
img = img[0].permute(1,2,0).cpu().numpy()
|
||||
flo = flo[0].permute(1,2,0).cpu().numpy()
|
||||
|
||||
# map flow to rgb image
|
||||
flo = flow_viz.flow_to_image(flo)
|
||||
# img_flo = np.concatenate([img, flo], axis=0)
|
||||
img_flo = flo
|
||||
|
||||
cv2.imwrite('/home/chengao/test/flow.png', img_flo[:, :, [2,1,0]])
|
||||
# cv2.imshow('image', img_flo[:, :, [2,1,0]]/255.0)
|
||||
# cv2.waitKey()
|
||||
|
||||
|
||||
def demo(args):
|
||||
model = torch.nn.DataParallel(RAFT(args))
|
||||
model.load_state_dict(torch.load(args.model))
|
||||
|
||||
model = model.module
|
||||
model.to(DEVICE)
|
||||
model.eval()
|
||||
|
||||
with torch.no_grad():
|
||||
images = glob.glob(os.path.join(args.path, '*.png')) + \
|
||||
glob.glob(os.path.join(args.path, '*.jpg'))
|
||||
|
||||
images = load_image_list(images)
|
||||
for i in range(images.shape[0]-1):
|
||||
image1 = images[i,None]
|
||||
image2 = images[i+1,None]
|
||||
|
||||
flow_low, flow_up = model(image1, image2, iters=20, test_mode=True)
|
||||
viz(image1, flow_up)
|
||||
|
||||
|
||||
def RAFT_infer(args):
|
||||
model = torch.nn.DataParallel(RAFT(args))
|
||||
model.load_state_dict(torch.load(args.model))
|
||||
|
||||
model = model.module
|
||||
model.to(DEVICE)
|
||||
model.eval()
|
||||
|
||||
return model
|
||||
Executable
+267
@@ -0,0 +1,267 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class ResidualBlock(nn.Module):
|
||||
def __init__(self, in_planes, planes, norm_fn='group', stride=1):
|
||||
super(ResidualBlock, self).__init__()
|
||||
|
||||
self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, padding=1, stride=stride)
|
||||
self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, padding=1)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
|
||||
num_groups = planes // 8
|
||||
|
||||
if norm_fn == 'group':
|
||||
self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
|
||||
self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
|
||||
if not stride == 1:
|
||||
self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
|
||||
|
||||
elif norm_fn == 'batch':
|
||||
self.norm1 = nn.BatchNorm2d(planes)
|
||||
self.norm2 = nn.BatchNorm2d(planes)
|
||||
if not stride == 1:
|
||||
self.norm3 = nn.BatchNorm2d(planes)
|
||||
|
||||
elif norm_fn == 'instance':
|
||||
self.norm1 = nn.InstanceNorm2d(planes)
|
||||
self.norm2 = nn.InstanceNorm2d(planes)
|
||||
if not stride == 1:
|
||||
self.norm3 = nn.InstanceNorm2d(planes)
|
||||
|
||||
elif norm_fn == 'none':
|
||||
self.norm1 = nn.Sequential()
|
||||
self.norm2 = nn.Sequential()
|
||||
if not stride == 1:
|
||||
self.norm3 = nn.Sequential()
|
||||
|
||||
if stride == 1:
|
||||
self.downsample = None
|
||||
|
||||
else:
|
||||
self.downsample = nn.Sequential(
|
||||
nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm3)
|
||||
|
||||
|
||||
def forward(self, x):
|
||||
y = x
|
||||
y = self.relu(self.norm1(self.conv1(y)))
|
||||
y = self.relu(self.norm2(self.conv2(y)))
|
||||
|
||||
if self.downsample is not None:
|
||||
x = self.downsample(x)
|
||||
|
||||
return self.relu(x+y)
|
||||
|
||||
|
||||
|
||||
class BottleneckBlock(nn.Module):
|
||||
def __init__(self, in_planes, planes, norm_fn='group', stride=1):
|
||||
super(BottleneckBlock, self).__init__()
|
||||
|
||||
self.conv1 = nn.Conv2d(in_planes, planes//4, kernel_size=1, padding=0)
|
||||
self.conv2 = nn.Conv2d(planes//4, planes//4, kernel_size=3, padding=1, stride=stride)
|
||||
self.conv3 = nn.Conv2d(planes//4, planes, kernel_size=1, padding=0)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
|
||||
num_groups = planes // 8
|
||||
|
||||
if norm_fn == 'group':
|
||||
self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes//4)
|
||||
self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes//4)
|
||||
self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
|
||||
if not stride == 1:
|
||||
self.norm4 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
|
||||
|
||||
elif norm_fn == 'batch':
|
||||
self.norm1 = nn.BatchNorm2d(planes//4)
|
||||
self.norm2 = nn.BatchNorm2d(planes//4)
|
||||
self.norm3 = nn.BatchNorm2d(planes)
|
||||
if not stride == 1:
|
||||
self.norm4 = nn.BatchNorm2d(planes)
|
||||
|
||||
elif norm_fn == 'instance':
|
||||
self.norm1 = nn.InstanceNorm2d(planes//4)
|
||||
self.norm2 = nn.InstanceNorm2d(planes//4)
|
||||
self.norm3 = nn.InstanceNorm2d(planes)
|
||||
if not stride == 1:
|
||||
self.norm4 = nn.InstanceNorm2d(planes)
|
||||
|
||||
elif norm_fn == 'none':
|
||||
self.norm1 = nn.Sequential()
|
||||
self.norm2 = nn.Sequential()
|
||||
self.norm3 = nn.Sequential()
|
||||
if not stride == 1:
|
||||
self.norm4 = nn.Sequential()
|
||||
|
||||
if stride == 1:
|
||||
self.downsample = None
|
||||
|
||||
else:
|
||||
self.downsample = nn.Sequential(
|
||||
nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm4)
|
||||
|
||||
|
||||
def forward(self, x):
|
||||
y = x
|
||||
y = self.relu(self.norm1(self.conv1(y)))
|
||||
y = self.relu(self.norm2(self.conv2(y)))
|
||||
y = self.relu(self.norm3(self.conv3(y)))
|
||||
|
||||
if self.downsample is not None:
|
||||
x = self.downsample(x)
|
||||
|
||||
return self.relu(x+y)
|
||||
|
||||
class BasicEncoder(nn.Module):
|
||||
def __init__(self, output_dim=128, norm_fn='batch', dropout=0.0):
|
||||
super(BasicEncoder, self).__init__()
|
||||
self.norm_fn = norm_fn
|
||||
|
||||
if self.norm_fn == 'group':
|
||||
self.norm1 = nn.GroupNorm(num_groups=8, num_channels=64)
|
||||
|
||||
elif self.norm_fn == 'batch':
|
||||
self.norm1 = nn.BatchNorm2d(64)
|
||||
|
||||
elif self.norm_fn == 'instance':
|
||||
self.norm1 = nn.InstanceNorm2d(64)
|
||||
|
||||
elif self.norm_fn == 'none':
|
||||
self.norm1 = nn.Sequential()
|
||||
|
||||
self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3)
|
||||
self.relu1 = nn.ReLU(inplace=True)
|
||||
|
||||
self.in_planes = 64
|
||||
self.layer1 = self._make_layer(64, stride=1)
|
||||
self.layer2 = self._make_layer(96, stride=2)
|
||||
self.layer3 = self._make_layer(128, stride=2)
|
||||
|
||||
# output convolution
|
||||
self.conv2 = nn.Conv2d(128, output_dim, kernel_size=1)
|
||||
|
||||
self.dropout = None
|
||||
if dropout > 0:
|
||||
self.dropout = nn.Dropout2d(p=dropout)
|
||||
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
|
||||
elif isinstance(m, (nn.BatchNorm2d, nn.InstanceNorm2d, nn.GroupNorm)):
|
||||
if m.weight is not None:
|
||||
nn.init.constant_(m.weight, 1)
|
||||
if m.bias is not None:
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
def _make_layer(self, dim, stride=1):
|
||||
layer1 = ResidualBlock(self.in_planes, dim, self.norm_fn, stride=stride)
|
||||
layer2 = ResidualBlock(dim, dim, self.norm_fn, stride=1)
|
||||
layers = (layer1, layer2)
|
||||
|
||||
self.in_planes = dim
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
|
||||
def forward(self, x):
|
||||
|
||||
# if input is list, combine batch dimension
|
||||
is_list = isinstance(x, tuple) or isinstance(x, list)
|
||||
if is_list:
|
||||
batch_dim = x[0].shape[0]
|
||||
x = torch.cat(x, dim=0)
|
||||
|
||||
x = self.conv1(x)
|
||||
x = self.norm1(x)
|
||||
x = self.relu1(x)
|
||||
|
||||
x = self.layer1(x)
|
||||
x = self.layer2(x)
|
||||
x = self.layer3(x)
|
||||
|
||||
x = self.conv2(x)
|
||||
|
||||
if self.training and self.dropout is not None:
|
||||
x = self.dropout(x)
|
||||
|
||||
if is_list:
|
||||
x = torch.split(x, [batch_dim, batch_dim], dim=0)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class SmallEncoder(nn.Module):
|
||||
def __init__(self, output_dim=128, norm_fn='batch', dropout=0.0):
|
||||
super(SmallEncoder, self).__init__()
|
||||
self.norm_fn = norm_fn
|
||||
|
||||
if self.norm_fn == 'group':
|
||||
self.norm1 = nn.GroupNorm(num_groups=8, num_channels=32)
|
||||
|
||||
elif self.norm_fn == 'batch':
|
||||
self.norm1 = nn.BatchNorm2d(32)
|
||||
|
||||
elif self.norm_fn == 'instance':
|
||||
self.norm1 = nn.InstanceNorm2d(32)
|
||||
|
||||
elif self.norm_fn == 'none':
|
||||
self.norm1 = nn.Sequential()
|
||||
|
||||
self.conv1 = nn.Conv2d(3, 32, kernel_size=7, stride=2, padding=3)
|
||||
self.relu1 = nn.ReLU(inplace=True)
|
||||
|
||||
self.in_planes = 32
|
||||
self.layer1 = self._make_layer(32, stride=1)
|
||||
self.layer2 = self._make_layer(64, stride=2)
|
||||
self.layer3 = self._make_layer(96, stride=2)
|
||||
|
||||
self.dropout = None
|
||||
if dropout > 0:
|
||||
self.dropout = nn.Dropout2d(p=dropout)
|
||||
|
||||
self.conv2 = nn.Conv2d(96, output_dim, kernel_size=1)
|
||||
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
|
||||
elif isinstance(m, (nn.BatchNorm2d, nn.InstanceNorm2d, nn.GroupNorm)):
|
||||
if m.weight is not None:
|
||||
nn.init.constant_(m.weight, 1)
|
||||
if m.bias is not None:
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
def _make_layer(self, dim, stride=1):
|
||||
layer1 = BottleneckBlock(self.in_planes, dim, self.norm_fn, stride=stride)
|
||||
layer2 = BottleneckBlock(dim, dim, self.norm_fn, stride=1)
|
||||
layers = (layer1, layer2)
|
||||
|
||||
self.in_planes = dim
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
|
||||
def forward(self, x):
|
||||
|
||||
# if input is list, combine batch dimension
|
||||
is_list = isinstance(x, tuple) or isinstance(x, list)
|
||||
if is_list:
|
||||
batch_dim = x[0].shape[0]
|
||||
x = torch.cat(x, dim=0)
|
||||
|
||||
x = self.conv1(x)
|
||||
x = self.norm1(x)
|
||||
x = self.relu1(x)
|
||||
|
||||
x = self.layer1(x)
|
||||
x = self.layer2(x)
|
||||
x = self.layer3(x)
|
||||
x = self.conv2(x)
|
||||
|
||||
if self.training and self.dropout is not None:
|
||||
x = self.dropout(x)
|
||||
|
||||
if is_list:
|
||||
x = torch.split(x, [batch_dim, batch_dim], dim=0)
|
||||
|
||||
return x
|
||||
Executable
+146
@@ -0,0 +1,146 @@
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from .update import BasicUpdateBlock, SmallUpdateBlock
|
||||
from .extractor import BasicEncoder, SmallEncoder
|
||||
from .corr import CorrBlock, AlternateCorrBlock
|
||||
from .utils.utils import bilinear_sampler, coords_grid, upflow8
|
||||
|
||||
try:
|
||||
autocast = torch.cuda.amp.autocast
|
||||
except:
|
||||
# dummy autocast for PyTorch < 1.6
|
||||
class autocast:
|
||||
def __init__(self, enabled):
|
||||
pass
|
||||
def __enter__(self):
|
||||
pass
|
||||
def __exit__(self, *args):
|
||||
pass
|
||||
|
||||
|
||||
class RAFT(nn.Module):
|
||||
def __init__(self, args):
|
||||
super(RAFT, self).__init__()
|
||||
self.args = args
|
||||
|
||||
if args.small:
|
||||
self.hidden_dim = hdim = 96
|
||||
self.context_dim = cdim = 64
|
||||
args.corr_levels = 4
|
||||
args.corr_radius = 3
|
||||
|
||||
else:
|
||||
self.hidden_dim = hdim = 128
|
||||
self.context_dim = cdim = 128
|
||||
args.corr_levels = 4
|
||||
args.corr_radius = 4
|
||||
|
||||
if 'dropout' not in args._get_kwargs():
|
||||
args.dropout = 0
|
||||
|
||||
if 'alternate_corr' not in args._get_kwargs():
|
||||
args.alternate_corr = False
|
||||
|
||||
# feature network, context network, and update block
|
||||
if args.small:
|
||||
self.fnet = SmallEncoder(output_dim=128, norm_fn='instance', dropout=args.dropout)
|
||||
self.cnet = SmallEncoder(output_dim=hdim+cdim, norm_fn='none', dropout=args.dropout)
|
||||
self.update_block = SmallUpdateBlock(self.args, hidden_dim=hdim)
|
||||
|
||||
else:
|
||||
self.fnet = BasicEncoder(output_dim=256, norm_fn='instance', dropout=args.dropout)
|
||||
self.cnet = BasicEncoder(output_dim=hdim+cdim, norm_fn='batch', dropout=args.dropout)
|
||||
self.update_block = BasicUpdateBlock(self.args, hidden_dim=hdim)
|
||||
|
||||
|
||||
def freeze_bn(self):
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.BatchNorm2d):
|
||||
m.eval()
|
||||
|
||||
def initialize_flow(self, img):
|
||||
""" Flow is represented as difference between two coordinate grids flow = coords1 - coords0"""
|
||||
N, C, H, W = img.shape
|
||||
coords0 = coords_grid(N, H//8, W//8).to(img.device)
|
||||
coords1 = coords_grid(N, H//8, W//8).to(img.device)
|
||||
|
||||
# optical flow computed as difference: flow = coords1 - coords0
|
||||
return coords0, coords1
|
||||
|
||||
def upsample_flow(self, flow, mask):
|
||||
""" Upsample flow field [H/8, W/8, 2] -> [H, W, 2] using convex combination """
|
||||
N, _, H, W = flow.shape
|
||||
mask = mask.view(N, 1, 9, 8, 8, H, W)
|
||||
mask = torch.softmax(mask, dim=2)
|
||||
|
||||
up_flow = F.unfold(8 * flow, [3,3], padding=1)
|
||||
up_flow = up_flow.view(N, 2, 9, 1, 1, H, W)
|
||||
|
||||
up_flow = torch.sum(mask * up_flow, dim=2)
|
||||
up_flow = up_flow.permute(0, 1, 4, 2, 5, 3)
|
||||
return up_flow.reshape(N, 2, 8*H, 8*W)
|
||||
|
||||
|
||||
def forward(self, image1, image2, iters=12, flow_init=None, test_mode=True):
|
||||
""" Estimate optical flow between pair of frames """
|
||||
|
||||
# image1 = 2 * (image1 / 255.0) - 1.0
|
||||
# image2 = 2 * (image2 / 255.0) - 1.0
|
||||
|
||||
image1 = image1.contiguous()
|
||||
image2 = image2.contiguous()
|
||||
|
||||
hdim = self.hidden_dim
|
||||
cdim = self.context_dim
|
||||
|
||||
# run the feature network
|
||||
with autocast(enabled=self.args.mixed_precision):
|
||||
fmap1, fmap2 = self.fnet([image1, image2])
|
||||
|
||||
fmap1 = fmap1.float()
|
||||
fmap2 = fmap2.float()
|
||||
|
||||
if self.args.alternate_corr:
|
||||
corr_fn = AlternateCorrBlock(fmap1, fmap2, radius=self.args.corr_radius)
|
||||
else:
|
||||
corr_fn = CorrBlock(fmap1, fmap2, radius=self.args.corr_radius)
|
||||
|
||||
# run the context network
|
||||
with autocast(enabled=self.args.mixed_precision):
|
||||
cnet = self.cnet(image1)
|
||||
net, inp = torch.split(cnet, [hdim, cdim], dim=1)
|
||||
net = torch.tanh(net)
|
||||
inp = torch.relu(inp)
|
||||
|
||||
coords0, coords1 = self.initialize_flow(image1)
|
||||
|
||||
if flow_init is not None:
|
||||
coords1 = coords1 + flow_init
|
||||
|
||||
flow_predictions = []
|
||||
for itr in range(iters):
|
||||
coords1 = coords1.detach()
|
||||
corr = corr_fn(coords1) # index correlation volume
|
||||
|
||||
flow = coords1 - coords0
|
||||
with autocast(enabled=self.args.mixed_precision):
|
||||
net, up_mask, delta_flow = self.update_block(net, inp, corr, flow)
|
||||
|
||||
# F(t+1) = F(t) + \Delta(t)
|
||||
coords1 = coords1 + delta_flow
|
||||
|
||||
# upsample predictions
|
||||
if up_mask is None:
|
||||
flow_up = upflow8(coords1 - coords0)
|
||||
else:
|
||||
flow_up = self.upsample_flow(coords1 - coords0, up_mask)
|
||||
|
||||
flow_predictions.append(flow_up)
|
||||
|
||||
if test_mode:
|
||||
return coords1 - coords0, flow_up
|
||||
|
||||
return flow_predictions
|
||||
Executable
+139
@@ -0,0 +1,139 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class FlowHead(nn.Module):
|
||||
def __init__(self, input_dim=128, hidden_dim=256):
|
||||
super(FlowHead, self).__init__()
|
||||
self.conv1 = nn.Conv2d(input_dim, hidden_dim, 3, padding=1)
|
||||
self.conv2 = nn.Conv2d(hidden_dim, 2, 3, padding=1)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
|
||||
def forward(self, x):
|
||||
return self.conv2(self.relu(self.conv1(x)))
|
||||
|
||||
class ConvGRU(nn.Module):
|
||||
def __init__(self, hidden_dim=128, input_dim=192+128):
|
||||
super(ConvGRU, self).__init__()
|
||||
self.convz = nn.Conv2d(hidden_dim+input_dim, hidden_dim, 3, padding=1)
|
||||
self.convr = nn.Conv2d(hidden_dim+input_dim, hidden_dim, 3, padding=1)
|
||||
self.convq = nn.Conv2d(hidden_dim+input_dim, hidden_dim, 3, padding=1)
|
||||
|
||||
def forward(self, h, x):
|
||||
hx = torch.cat([h, x], dim=1)
|
||||
|
||||
z = torch.sigmoid(self.convz(hx))
|
||||
r = torch.sigmoid(self.convr(hx))
|
||||
q = torch.tanh(self.convq(torch.cat([r*h, x], dim=1)))
|
||||
|
||||
h = (1-z) * h + z * q
|
||||
return h
|
||||
|
||||
class SepConvGRU(nn.Module):
|
||||
def __init__(self, hidden_dim=128, input_dim=192+128):
|
||||
super(SepConvGRU, self).__init__()
|
||||
self.convz1 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (1,5), padding=(0,2))
|
||||
self.convr1 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (1,5), padding=(0,2))
|
||||
self.convq1 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (1,5), padding=(0,2))
|
||||
|
||||
self.convz2 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (5,1), padding=(2,0))
|
||||
self.convr2 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (5,1), padding=(2,0))
|
||||
self.convq2 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (5,1), padding=(2,0))
|
||||
|
||||
|
||||
def forward(self, h, x):
|
||||
# horizontal
|
||||
hx = torch.cat([h, x], dim=1)
|
||||
z = torch.sigmoid(self.convz1(hx))
|
||||
r = torch.sigmoid(self.convr1(hx))
|
||||
q = torch.tanh(self.convq1(torch.cat([r*h, x], dim=1)))
|
||||
h = (1-z) * h + z * q
|
||||
|
||||
# vertical
|
||||
hx = torch.cat([h, x], dim=1)
|
||||
z = torch.sigmoid(self.convz2(hx))
|
||||
r = torch.sigmoid(self.convr2(hx))
|
||||
q = torch.tanh(self.convq2(torch.cat([r*h, x], dim=1)))
|
||||
h = (1-z) * h + z * q
|
||||
|
||||
return h
|
||||
|
||||
class SmallMotionEncoder(nn.Module):
|
||||
def __init__(self, args):
|
||||
super(SmallMotionEncoder, self).__init__()
|
||||
cor_planes = args.corr_levels * (2*args.corr_radius + 1)**2
|
||||
self.convc1 = nn.Conv2d(cor_planes, 96, 1, padding=0)
|
||||
self.convf1 = nn.Conv2d(2, 64, 7, padding=3)
|
||||
self.convf2 = nn.Conv2d(64, 32, 3, padding=1)
|
||||
self.conv = nn.Conv2d(128, 80, 3, padding=1)
|
||||
|
||||
def forward(self, flow, corr):
|
||||
cor = F.relu(self.convc1(corr))
|
||||
flo = F.relu(self.convf1(flow))
|
||||
flo = F.relu(self.convf2(flo))
|
||||
cor_flo = torch.cat([cor, flo], dim=1)
|
||||
out = F.relu(self.conv(cor_flo))
|
||||
return torch.cat([out, flow], dim=1)
|
||||
|
||||
class BasicMotionEncoder(nn.Module):
|
||||
def __init__(self, args):
|
||||
super(BasicMotionEncoder, self).__init__()
|
||||
cor_planes = args.corr_levels * (2*args.corr_radius + 1)**2
|
||||
self.convc1 = nn.Conv2d(cor_planes, 256, 1, padding=0)
|
||||
self.convc2 = nn.Conv2d(256, 192, 3, padding=1)
|
||||
self.convf1 = nn.Conv2d(2, 128, 7, padding=3)
|
||||
self.convf2 = nn.Conv2d(128, 64, 3, padding=1)
|
||||
self.conv = nn.Conv2d(64+192, 128-2, 3, padding=1)
|
||||
|
||||
def forward(self, flow, corr):
|
||||
cor = F.relu(self.convc1(corr))
|
||||
cor = F.relu(self.convc2(cor))
|
||||
flo = F.relu(self.convf1(flow))
|
||||
flo = F.relu(self.convf2(flo))
|
||||
|
||||
cor_flo = torch.cat([cor, flo], dim=1)
|
||||
out = F.relu(self.conv(cor_flo))
|
||||
return torch.cat([out, flow], dim=1)
|
||||
|
||||
class SmallUpdateBlock(nn.Module):
|
||||
def __init__(self, args, hidden_dim=96):
|
||||
super(SmallUpdateBlock, self).__init__()
|
||||
self.encoder = SmallMotionEncoder(args)
|
||||
self.gru = ConvGRU(hidden_dim=hidden_dim, input_dim=82+64)
|
||||
self.flow_head = FlowHead(hidden_dim, hidden_dim=128)
|
||||
|
||||
def forward(self, net, inp, corr, flow):
|
||||
motion_features = self.encoder(flow, corr)
|
||||
inp = torch.cat([inp, motion_features], dim=1)
|
||||
net = self.gru(net, inp)
|
||||
delta_flow = self.flow_head(net)
|
||||
|
||||
return net, None, delta_flow
|
||||
|
||||
class BasicUpdateBlock(nn.Module):
|
||||
def __init__(self, args, hidden_dim=128, input_dim=128):
|
||||
super(BasicUpdateBlock, self).__init__()
|
||||
self.args = args
|
||||
self.encoder = BasicMotionEncoder(args)
|
||||
self.gru = SepConvGRU(hidden_dim=hidden_dim, input_dim=128+hidden_dim)
|
||||
self.flow_head = FlowHead(hidden_dim, hidden_dim=256)
|
||||
|
||||
self.mask = nn.Sequential(
|
||||
nn.Conv2d(128, 256, 3, padding=1),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Conv2d(256, 64*9, 1, padding=0))
|
||||
|
||||
def forward(self, net, inp, corr, flow, upsample=True):
|
||||
motion_features = self.encoder(flow, corr)
|
||||
inp = torch.cat([inp, motion_features], dim=1)
|
||||
|
||||
net = self.gru(net, inp)
|
||||
delta_flow = self.flow_head(net)
|
||||
|
||||
# scale mask to balence gradients
|
||||
mask = .25 * self.mask(net)
|
||||
return net, mask, delta_flow
|
||||
|
||||
|
||||
|
||||
Executable
+2
@@ -0,0 +1,2 @@
|
||||
from .flow_viz import flow_to_image
|
||||
from .frame_utils import writeFlow
|
||||
+246
@@ -0,0 +1,246 @@
|
||||
import numpy as np
|
||||
import random
|
||||
import math
|
||||
from PIL import Image
|
||||
|
||||
import cv2
|
||||
cv2.setNumThreads(0)
|
||||
cv2.ocl.setUseOpenCL(False)
|
||||
|
||||
import torch
|
||||
from torchvision.transforms import ColorJitter
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class FlowAugmentor:
|
||||
def __init__(self, crop_size, min_scale=-0.2, max_scale=0.5, do_flip=True):
|
||||
|
||||
# spatial augmentation params
|
||||
self.crop_size = crop_size
|
||||
self.min_scale = min_scale
|
||||
self.max_scale = max_scale
|
||||
self.spatial_aug_prob = 0.8
|
||||
self.stretch_prob = 0.8
|
||||
self.max_stretch = 0.2
|
||||
|
||||
# flip augmentation params
|
||||
self.do_flip = do_flip
|
||||
self.h_flip_prob = 0.5
|
||||
self.v_flip_prob = 0.1
|
||||
|
||||
# photometric augmentation params
|
||||
self.photo_aug = ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0.5/3.14)
|
||||
self.asymmetric_color_aug_prob = 0.2
|
||||
self.eraser_aug_prob = 0.5
|
||||
|
||||
def color_transform(self, img1, img2):
|
||||
""" Photometric augmentation """
|
||||
|
||||
# asymmetric
|
||||
if np.random.rand() < self.asymmetric_color_aug_prob:
|
||||
img1 = np.array(self.photo_aug(Image.fromarray(img1)), dtype=np.uint8)
|
||||
img2 = np.array(self.photo_aug(Image.fromarray(img2)), dtype=np.uint8)
|
||||
|
||||
# symmetric
|
||||
else:
|
||||
image_stack = np.concatenate([img1, img2], axis=0)
|
||||
image_stack = np.array(self.photo_aug(Image.fromarray(image_stack)), dtype=np.uint8)
|
||||
img1, img2 = np.split(image_stack, 2, axis=0)
|
||||
|
||||
return img1, img2
|
||||
|
||||
def eraser_transform(self, img1, img2, bounds=[50, 100]):
|
||||
""" Occlusion augmentation """
|
||||
|
||||
ht, wd = img1.shape[:2]
|
||||
if np.random.rand() < self.eraser_aug_prob:
|
||||
mean_color = np.mean(img2.reshape(-1, 3), axis=0)
|
||||
for _ in range(np.random.randint(1, 3)):
|
||||
x0 = np.random.randint(0, wd)
|
||||
y0 = np.random.randint(0, ht)
|
||||
dx = np.random.randint(bounds[0], bounds[1])
|
||||
dy = np.random.randint(bounds[0], bounds[1])
|
||||
img2[y0:y0+dy, x0:x0+dx, :] = mean_color
|
||||
|
||||
return img1, img2
|
||||
|
||||
def spatial_transform(self, img1, img2, flow):
|
||||
# randomly sample scale
|
||||
ht, wd = img1.shape[:2]
|
||||
min_scale = np.maximum(
|
||||
(self.crop_size[0] + 8) / float(ht),
|
||||
(self.crop_size[1] + 8) / float(wd))
|
||||
|
||||
scale = 2 ** np.random.uniform(self.min_scale, self.max_scale)
|
||||
scale_x = scale
|
||||
scale_y = scale
|
||||
if np.random.rand() < self.stretch_prob:
|
||||
scale_x *= 2 ** np.random.uniform(-self.max_stretch, self.max_stretch)
|
||||
scale_y *= 2 ** np.random.uniform(-self.max_stretch, self.max_stretch)
|
||||
|
||||
scale_x = np.clip(scale_x, min_scale, None)
|
||||
scale_y = np.clip(scale_y, min_scale, None)
|
||||
|
||||
if np.random.rand() < self.spatial_aug_prob:
|
||||
# rescale the images
|
||||
img1 = cv2.resize(img1, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR)
|
||||
img2 = cv2.resize(img2, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR)
|
||||
flow = cv2.resize(flow, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR)
|
||||
flow = flow * [scale_x, scale_y]
|
||||
|
||||
if self.do_flip:
|
||||
if np.random.rand() < self.h_flip_prob: # h-flip
|
||||
img1 = img1[:, ::-1]
|
||||
img2 = img2[:, ::-1]
|
||||
flow = flow[:, ::-1] * [-1.0, 1.0]
|
||||
|
||||
if np.random.rand() < self.v_flip_prob: # v-flip
|
||||
img1 = img1[::-1, :]
|
||||
img2 = img2[::-1, :]
|
||||
flow = flow[::-1, :] * [1.0, -1.0]
|
||||
|
||||
y0 = np.random.randint(0, img1.shape[0] - self.crop_size[0])
|
||||
x0 = np.random.randint(0, img1.shape[1] - self.crop_size[1])
|
||||
|
||||
img1 = img1[y0:y0+self.crop_size[0], x0:x0+self.crop_size[1]]
|
||||
img2 = img2[y0:y0+self.crop_size[0], x0:x0+self.crop_size[1]]
|
||||
flow = flow[y0:y0+self.crop_size[0], x0:x0+self.crop_size[1]]
|
||||
|
||||
return img1, img2, flow
|
||||
|
||||
def __call__(self, img1, img2, flow):
|
||||
img1, img2 = self.color_transform(img1, img2)
|
||||
img1, img2 = self.eraser_transform(img1, img2)
|
||||
img1, img2, flow = self.spatial_transform(img1, img2, flow)
|
||||
|
||||
img1 = np.ascontiguousarray(img1)
|
||||
img2 = np.ascontiguousarray(img2)
|
||||
flow = np.ascontiguousarray(flow)
|
||||
|
||||
return img1, img2, flow
|
||||
|
||||
class SparseFlowAugmentor:
|
||||
def __init__(self, crop_size, min_scale=-0.2, max_scale=0.5, do_flip=False):
|
||||
# spatial augmentation params
|
||||
self.crop_size = crop_size
|
||||
self.min_scale = min_scale
|
||||
self.max_scale = max_scale
|
||||
self.spatial_aug_prob = 0.8
|
||||
self.stretch_prob = 0.8
|
||||
self.max_stretch = 0.2
|
||||
|
||||
# flip augmentation params
|
||||
self.do_flip = do_flip
|
||||
self.h_flip_prob = 0.5
|
||||
self.v_flip_prob = 0.1
|
||||
|
||||
# photometric augmentation params
|
||||
self.photo_aug = ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3, hue=0.3/3.14)
|
||||
self.asymmetric_color_aug_prob = 0.2
|
||||
self.eraser_aug_prob = 0.5
|
||||
|
||||
def color_transform(self, img1, img2):
|
||||
image_stack = np.concatenate([img1, img2], axis=0)
|
||||
image_stack = np.array(self.photo_aug(Image.fromarray(image_stack)), dtype=np.uint8)
|
||||
img1, img2 = np.split(image_stack, 2, axis=0)
|
||||
return img1, img2
|
||||
|
||||
def eraser_transform(self, img1, img2):
|
||||
ht, wd = img1.shape[:2]
|
||||
if np.random.rand() < self.eraser_aug_prob:
|
||||
mean_color = np.mean(img2.reshape(-1, 3), axis=0)
|
||||
for _ in range(np.random.randint(1, 3)):
|
||||
x0 = np.random.randint(0, wd)
|
||||
y0 = np.random.randint(0, ht)
|
||||
dx = np.random.randint(50, 100)
|
||||
dy = np.random.randint(50, 100)
|
||||
img2[y0:y0+dy, x0:x0+dx, :] = mean_color
|
||||
|
||||
return img1, img2
|
||||
|
||||
def resize_sparse_flow_map(self, flow, valid, fx=1.0, fy=1.0):
|
||||
ht, wd = flow.shape[:2]
|
||||
coords = np.meshgrid(np.arange(wd), np.arange(ht))
|
||||
coords = np.stack(coords, axis=-1)
|
||||
|
||||
coords = coords.reshape(-1, 2).astype(np.float32)
|
||||
flow = flow.reshape(-1, 2).astype(np.float32)
|
||||
valid = valid.reshape(-1).astype(np.float32)
|
||||
|
||||
coords0 = coords[valid>=1]
|
||||
flow0 = flow[valid>=1]
|
||||
|
||||
ht1 = int(round(ht * fy))
|
||||
wd1 = int(round(wd * fx))
|
||||
|
||||
coords1 = coords0 * [fx, fy]
|
||||
flow1 = flow0 * [fx, fy]
|
||||
|
||||
xx = np.round(coords1[:,0]).astype(np.int32)
|
||||
yy = np.round(coords1[:,1]).astype(np.int32)
|
||||
|
||||
v = (xx > 0) & (xx < wd1) & (yy > 0) & (yy < ht1)
|
||||
xx = xx[v]
|
||||
yy = yy[v]
|
||||
flow1 = flow1[v]
|
||||
|
||||
flow_img = np.zeros([ht1, wd1, 2], dtype=np.float32)
|
||||
valid_img = np.zeros([ht1, wd1], dtype=np.int32)
|
||||
|
||||
flow_img[yy, xx] = flow1
|
||||
valid_img[yy, xx] = 1
|
||||
|
||||
return flow_img, valid_img
|
||||
|
||||
def spatial_transform(self, img1, img2, flow, valid):
|
||||
# randomly sample scale
|
||||
|
||||
ht, wd = img1.shape[:2]
|
||||
min_scale = np.maximum(
|
||||
(self.crop_size[0] + 1) / float(ht),
|
||||
(self.crop_size[1] + 1) / float(wd))
|
||||
|
||||
scale = 2 ** np.random.uniform(self.min_scale, self.max_scale)
|
||||
scale_x = np.clip(scale, min_scale, None)
|
||||
scale_y = np.clip(scale, min_scale, None)
|
||||
|
||||
if np.random.rand() < self.spatial_aug_prob:
|
||||
# rescale the images
|
||||
img1 = cv2.resize(img1, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR)
|
||||
img2 = cv2.resize(img2, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR)
|
||||
flow, valid = self.resize_sparse_flow_map(flow, valid, fx=scale_x, fy=scale_y)
|
||||
|
||||
if self.do_flip:
|
||||
if np.random.rand() < 0.5: # h-flip
|
||||
img1 = img1[:, ::-1]
|
||||
img2 = img2[:, ::-1]
|
||||
flow = flow[:, ::-1] * [-1.0, 1.0]
|
||||
valid = valid[:, ::-1]
|
||||
|
||||
margin_y = 20
|
||||
margin_x = 50
|
||||
|
||||
y0 = np.random.randint(0, img1.shape[0] - self.crop_size[0] + margin_y)
|
||||
x0 = np.random.randint(-margin_x, img1.shape[1] - self.crop_size[1] + margin_x)
|
||||
|
||||
y0 = np.clip(y0, 0, img1.shape[0] - self.crop_size[0])
|
||||
x0 = np.clip(x0, 0, img1.shape[1] - self.crop_size[1])
|
||||
|
||||
img1 = img1[y0:y0+self.crop_size[0], x0:x0+self.crop_size[1]]
|
||||
img2 = img2[y0:y0+self.crop_size[0], x0:x0+self.crop_size[1]]
|
||||
flow = flow[y0:y0+self.crop_size[0], x0:x0+self.crop_size[1]]
|
||||
valid = valid[y0:y0+self.crop_size[0], x0:x0+self.crop_size[1]]
|
||||
return img1, img2, flow, valid
|
||||
|
||||
|
||||
def __call__(self, img1, img2, flow, valid):
|
||||
img1, img2 = self.color_transform(img1, img2)
|
||||
img1, img2 = self.eraser_transform(img1, img2)
|
||||
img1, img2, flow, valid = self.spatial_transform(img1, img2, flow, valid)
|
||||
|
||||
img1 = np.ascontiguousarray(img1)
|
||||
img2 = np.ascontiguousarray(img2)
|
||||
flow = np.ascontiguousarray(flow)
|
||||
valid = np.ascontiguousarray(valid)
|
||||
|
||||
return img1, img2, flow, valid
|
||||
Executable
+132
@@ -0,0 +1,132 @@
|
||||
# Flow visualization code used from https://github.com/tomrunia/OpticalFlow_Visualization
|
||||
|
||||
|
||||
# MIT License
|
||||
#
|
||||
# Copyright (c) 2018 Tom Runia
|
||||
#
|
||||
# Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
# of this software and associated documentation files (the "Software"), to deal
|
||||
# in the Software without restriction, including without limitation the rights
|
||||
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
# copies of the Software, and to permit persons to whom the Software is
|
||||
# furnished to do so, subject to conditions.
|
||||
#
|
||||
# Author: Tom Runia
|
||||
# Date Created: 2018-08-03
|
||||
|
||||
import numpy as np
|
||||
|
||||
def make_colorwheel():
|
||||
"""
|
||||
Generates a color wheel for optical flow visualization as presented in:
|
||||
Baker et al. "A Database and Evaluation Methodology for Optical Flow" (ICCV, 2007)
|
||||
URL: http://vision.middlebury.edu/flow/flowEval-iccv07.pdf
|
||||
|
||||
Code follows the original C++ source code of Daniel Scharstein.
|
||||
Code follows the the Matlab source code of Deqing Sun.
|
||||
|
||||
Returns:
|
||||
np.ndarray: Color wheel
|
||||
"""
|
||||
|
||||
RY = 15
|
||||
YG = 6
|
||||
GC = 4
|
||||
CB = 11
|
||||
BM = 13
|
||||
MR = 6
|
||||
|
||||
ncols = RY + YG + GC + CB + BM + MR
|
||||
colorwheel = np.zeros((ncols, 3))
|
||||
col = 0
|
||||
|
||||
# RY
|
||||
colorwheel[0:RY, 0] = 255
|
||||
colorwheel[0:RY, 1] = np.floor(255*np.arange(0,RY)/RY)
|
||||
col = col+RY
|
||||
# YG
|
||||
colorwheel[col:col+YG, 0] = 255 - np.floor(255*np.arange(0,YG)/YG)
|
||||
colorwheel[col:col+YG, 1] = 255
|
||||
col = col+YG
|
||||
# GC
|
||||
colorwheel[col:col+GC, 1] = 255
|
||||
colorwheel[col:col+GC, 2] = np.floor(255*np.arange(0,GC)/GC)
|
||||
col = col+GC
|
||||
# CB
|
||||
colorwheel[col:col+CB, 1] = 255 - np.floor(255*np.arange(CB)/CB)
|
||||
colorwheel[col:col+CB, 2] = 255
|
||||
col = col+CB
|
||||
# BM
|
||||
colorwheel[col:col+BM, 2] = 255
|
||||
colorwheel[col:col+BM, 0] = np.floor(255*np.arange(0,BM)/BM)
|
||||
col = col+BM
|
||||
# MR
|
||||
colorwheel[col:col+MR, 2] = 255 - np.floor(255*np.arange(MR)/MR)
|
||||
colorwheel[col:col+MR, 0] = 255
|
||||
return colorwheel
|
||||
|
||||
|
||||
def flow_uv_to_colors(u, v, convert_to_bgr=False):
|
||||
"""
|
||||
Applies the flow color wheel to (possibly clipped) flow components u and v.
|
||||
|
||||
According to the C++ source code of Daniel Scharstein
|
||||
According to the Matlab source code of Deqing Sun
|
||||
|
||||
Args:
|
||||
u (np.ndarray): Input horizontal flow of shape [H,W]
|
||||
v (np.ndarray): Input vertical flow of shape [H,W]
|
||||
convert_to_bgr (bool, optional): Convert output image to BGR. Defaults to False.
|
||||
|
||||
Returns:
|
||||
np.ndarray: Flow visualization image of shape [H,W,3]
|
||||
"""
|
||||
flow_image = np.zeros((u.shape[0], u.shape[1], 3), np.uint8)
|
||||
colorwheel = make_colorwheel() # shape [55x3]
|
||||
ncols = colorwheel.shape[0]
|
||||
rad = np.sqrt(np.square(u) + np.square(v))
|
||||
a = np.arctan2(-v, -u)/np.pi
|
||||
fk = (a+1) / 2*(ncols-1)
|
||||
k0 = np.floor(fk).astype(np.int32)
|
||||
k1 = k0 + 1
|
||||
k1[k1 == ncols] = 0
|
||||
f = fk - k0
|
||||
for i in range(colorwheel.shape[1]):
|
||||
tmp = colorwheel[:,i]
|
||||
col0 = tmp[k0] / 255.0
|
||||
col1 = tmp[k1] / 255.0
|
||||
col = (1-f)*col0 + f*col1
|
||||
idx = (rad <= 1)
|
||||
col[idx] = 1 - rad[idx] * (1-col[idx])
|
||||
col[~idx] = col[~idx] * 0.75 # out of range
|
||||
# Note the 2-i => BGR instead of RGB
|
||||
ch_idx = 2-i if convert_to_bgr else i
|
||||
flow_image[:,:,ch_idx] = np.floor(255 * col)
|
||||
return flow_image
|
||||
|
||||
|
||||
def flow_to_image(flow_uv, clip_flow=None, convert_to_bgr=False):
|
||||
"""
|
||||
Expects a two dimensional flow image of shape.
|
||||
|
||||
Args:
|
||||
flow_uv (np.ndarray): Flow UV image of shape [H,W,2]
|
||||
clip_flow (float, optional): Clip maximum of flow values. Defaults to None.
|
||||
convert_to_bgr (bool, optional): Convert output image to BGR. Defaults to False.
|
||||
|
||||
Returns:
|
||||
np.ndarray: Flow visualization image of shape [H,W,3]
|
||||
"""
|
||||
assert flow_uv.ndim == 3, 'input flow must have three dimensions'
|
||||
assert flow_uv.shape[2] == 2, 'input flow must have shape [H,W,2]'
|
||||
if clip_flow is not None:
|
||||
flow_uv = np.clip(flow_uv, 0, clip_flow)
|
||||
u = flow_uv[:,:,0]
|
||||
v = flow_uv[:,:,1]
|
||||
rad = np.sqrt(np.square(u) + np.square(v))
|
||||
rad_max = np.max(rad)
|
||||
epsilon = 1e-5
|
||||
u = u / (rad_max + epsilon)
|
||||
v = v / (rad_max + epsilon)
|
||||
return flow_uv_to_colors(u, v, convert_to_bgr)
|
||||
@@ -0,0 +1,118 @@
|
||||
# Flow visualization code adapted from https://github.com/tomrunia/OpticalFlow_Visualization
|
||||
import torch
|
||||
torch.pi = torch.acos(torch.zeros(1)).item() * 2 # which is 3.1415927410125732
|
||||
|
||||
@torch.no_grad()
|
||||
def flow_to_image(flow: torch.Tensor) -> torch.Tensor:
|
||||
|
||||
"""
|
||||
Converts a flow to an RGB image.
|
||||
|
||||
Args:
|
||||
flow (Tensor): Flow of shape (N, 2, H, W) or (2, H, W) and dtype torch.float.
|
||||
|
||||
Returns:
|
||||
img (Tensor): Image Tensor of dtype uint8 where each color corresponds
|
||||
to a given flow direction. Shape is (N, 3, H, W) or (3, H, W) depending on the input.
|
||||
"""
|
||||
|
||||
if flow.dtype != torch.float:
|
||||
raise ValueError(f"Flow should be of dtype torch.float, got {flow.dtype}.")
|
||||
|
||||
orig_shape = flow.shape
|
||||
if flow.ndim == 3:
|
||||
flow = flow[None] # Add batch dim
|
||||
|
||||
if flow.ndim != 4 or flow.shape[1] != 2:
|
||||
raise ValueError(f"Input flow should have shape (2, H, W) or (N, 2, H, W), got {orig_shape}.")
|
||||
|
||||
max_norm = torch.sum(flow**2, dim=1).sqrt().max()
|
||||
epsilon = torch.finfo((flow).dtype).eps
|
||||
normalized_flow = flow / (max_norm + epsilon)
|
||||
img = _normalized_flow_to_image(normalized_flow)
|
||||
|
||||
if len(orig_shape) == 3:
|
||||
img = img[0] # Remove batch dim
|
||||
return img
|
||||
|
||||
@torch.no_grad()
|
||||
def _normalized_flow_to_image(normalized_flow: torch.Tensor) -> torch.Tensor:
|
||||
|
||||
"""
|
||||
Converts a batch of normalized flow to an RGB image.
|
||||
|
||||
Args:
|
||||
normalized_flow (torch.Tensor): Normalized flow tensor of shape (N, 2, H, W)
|
||||
Returns:
|
||||
img (Tensor(N, 3, H, W)): Flow visualization image of dtype uint8.
|
||||
"""
|
||||
|
||||
N, _, H, W = normalized_flow.shape
|
||||
device = normalized_flow.device
|
||||
flow_image = torch.zeros((N, 3, H, W), dtype=torch.uint8, device=device)
|
||||
colorwheel = _make_colorwheel().to(device) # shape [55x3]
|
||||
num_cols = colorwheel.shape[0]
|
||||
norm = torch.sum(normalized_flow**2, dim=1).sqrt()
|
||||
a = torch.atan2(-normalized_flow[:, 1, :, :], -normalized_flow[:, 0, :, :]) / torch.pi
|
||||
fk = (a + 1) / 2 * (num_cols - 1)
|
||||
k0 = torch.floor(fk).to(torch.long)
|
||||
k1 = k0 + 1
|
||||
k1[k1 == num_cols] = 0
|
||||
f = fk - k0
|
||||
|
||||
for c in range(colorwheel.shape[1]):
|
||||
tmp = colorwheel[:, c]
|
||||
col0 = tmp[k0] / 255.0
|
||||
col1 = tmp[k1] / 255.0
|
||||
col = (1 - f) * col0 + f * col1
|
||||
col = 1 - norm * (1 - col)
|
||||
flow_image[:, c, :, :] = torch.floor(255. * col)
|
||||
return flow_image
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def _make_colorwheel() -> torch.Tensor:
|
||||
"""
|
||||
Generates a color wheel for optical flow visualization as presented in:
|
||||
Baker et al. "A Database and Evaluation Methodology for Optical Flow" (ICCV, 2007)
|
||||
URL: http://vision.middlebury.edu/flow/flowEval-iccv07.pdf.
|
||||
|
||||
Returns:
|
||||
colorwheel (Tensor[55, 3]): Colorwheel Tensor.
|
||||
"""
|
||||
|
||||
RY = 15
|
||||
YG = 6
|
||||
GC = 4
|
||||
CB = 11
|
||||
BM = 13
|
||||
MR = 6
|
||||
|
||||
ncols = RY + YG + GC + CB + BM + MR
|
||||
colorwheel = torch.zeros((ncols, 3))
|
||||
col = 0
|
||||
|
||||
# RY
|
||||
colorwheel[0:RY, 0] = 255
|
||||
colorwheel[0:RY, 1] = torch.floor(255. * torch.arange(0., RY) / RY)
|
||||
col = col + RY
|
||||
# YG
|
||||
colorwheel[col : col + YG, 0] = 255 - torch.floor(255. * torch.arange(0., YG) / YG)
|
||||
colorwheel[col : col + YG, 1] = 255
|
||||
col = col + YG
|
||||
# GC
|
||||
colorwheel[col : col + GC, 1] = 255
|
||||
colorwheel[col : col + GC, 2] = torch.floor(255. * torch.arange(0., GC) / GC)
|
||||
col = col + GC
|
||||
# CB
|
||||
colorwheel[col : col + CB, 1] = 255 - torch.floor(255. * torch.arange(CB) / CB)
|
||||
colorwheel[col : col + CB, 2] = 255
|
||||
col = col + CB
|
||||
# BM
|
||||
colorwheel[col : col + BM, 2] = 255
|
||||
colorwheel[col : col + BM, 0] = torch.floor(255. * torch.arange(0., BM) / BM)
|
||||
col = col + BM
|
||||
# MR
|
||||
colorwheel[col : col + MR, 2] = 255 - torch.floor(255. * torch.arange(MR) / MR)
|
||||
colorwheel[col : col + MR, 0] = 255
|
||||
return colorwheel
|
||||
+137
@@ -0,0 +1,137 @@
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from os.path import *
|
||||
import re
|
||||
|
||||
import cv2
|
||||
cv2.setNumThreads(0)
|
||||
cv2.ocl.setUseOpenCL(False)
|
||||
|
||||
TAG_CHAR = np.array([202021.25], np.float32)
|
||||
|
||||
def readFlow(fn):
|
||||
""" Read .flo file in Middlebury format"""
|
||||
# Code adapted from:
|
||||
# http://stackoverflow.com/questions/28013200/reading-middlebury-flow-files-with-python-bytes-array-numpy
|
||||
|
||||
# WARNING: this will work on little-endian architectures (eg Intel x86) only!
|
||||
# print 'fn = %s'%(fn)
|
||||
with open(fn, 'rb') as f:
|
||||
magic = np.fromfile(f, np.float32, count=1)
|
||||
if 202021.25 != magic:
|
||||
print('Magic number incorrect. Invalid .flo file')
|
||||
return None
|
||||
else:
|
||||
w = np.fromfile(f, np.int32, count=1)
|
||||
h = np.fromfile(f, np.int32, count=1)
|
||||
# print 'Reading %d x %d flo file\n' % (w, h)
|
||||
data = np.fromfile(f, np.float32, count=2*int(w)*int(h))
|
||||
# Reshape data into 3D array (columns, rows, bands)
|
||||
# The reshape here is for visualization, the original code is (w,h,2)
|
||||
return np.resize(data, (int(h), int(w), 2))
|
||||
|
||||
def readPFM(file):
|
||||
file = open(file, 'rb')
|
||||
|
||||
color = None
|
||||
width = None
|
||||
height = None
|
||||
scale = None
|
||||
endian = None
|
||||
|
||||
header = file.readline().rstrip()
|
||||
if header == b'PF':
|
||||
color = True
|
||||
elif header == b'Pf':
|
||||
color = False
|
||||
else:
|
||||
raise Exception('Not a PFM file.')
|
||||
|
||||
dim_match = re.match(rb'^(\d+)\s(\d+)\s$', file.readline())
|
||||
if dim_match:
|
||||
width, height = map(int, dim_match.groups())
|
||||
else:
|
||||
raise Exception('Malformed PFM header.')
|
||||
|
||||
scale = float(file.readline().rstrip())
|
||||
if scale < 0: # little-endian
|
||||
endian = '<'
|
||||
scale = -scale
|
||||
else:
|
||||
endian = '>' # big-endian
|
||||
|
||||
data = np.fromfile(file, endian + 'f')
|
||||
shape = (height, width, 3) if color else (height, width)
|
||||
|
||||
data = np.reshape(data, shape)
|
||||
data = np.flipud(data)
|
||||
return data
|
||||
|
||||
def writeFlow(filename,uv,v=None):
|
||||
""" Write optical flow to file.
|
||||
|
||||
If v is None, uv is assumed to contain both u and v channels,
|
||||
stacked in depth.
|
||||
Original code by Deqing Sun, adapted from Daniel Scharstein.
|
||||
"""
|
||||
nBands = 2
|
||||
|
||||
if v is None:
|
||||
assert(uv.ndim == 3)
|
||||
assert(uv.shape[2] == 2)
|
||||
u = uv[:,:,0]
|
||||
v = uv[:,:,1]
|
||||
else:
|
||||
u = uv
|
||||
|
||||
assert(u.shape == v.shape)
|
||||
height,width = u.shape
|
||||
f = open(filename,'wb')
|
||||
# write the header
|
||||
f.write(TAG_CHAR)
|
||||
np.array(width).astype(np.int32).tofile(f)
|
||||
np.array(height).astype(np.int32).tofile(f)
|
||||
# arrange into matrix form
|
||||
tmp = np.zeros((height, width*nBands))
|
||||
tmp[:,np.arange(width)*2] = u
|
||||
tmp[:,np.arange(width)*2 + 1] = v
|
||||
tmp.astype(np.float32).tofile(f)
|
||||
f.close()
|
||||
|
||||
|
||||
def readFlowKITTI(filename):
|
||||
flow = cv2.imread(filename, cv2.IMREAD_ANYDEPTH|cv2.IMREAD_COLOR)
|
||||
flow = flow[:,:,::-1].astype(np.float32)
|
||||
flow, valid = flow[:, :, :2], flow[:, :, 2]
|
||||
flow = (flow - 2**15) / 64.0
|
||||
return flow, valid
|
||||
|
||||
def readDispKITTI(filename):
|
||||
disp = cv2.imread(filename, cv2.IMREAD_ANYDEPTH) / 256.0
|
||||
valid = disp > 0.0
|
||||
flow = np.stack([-disp, np.zeros_like(disp)], -1)
|
||||
return flow, valid
|
||||
|
||||
|
||||
def writeFlowKITTI(filename, uv):
|
||||
uv = 64.0 * uv + 2**15
|
||||
valid = np.ones([uv.shape[0], uv.shape[1], 1])
|
||||
uv = np.concatenate([uv, valid], axis=-1).astype(np.uint16)
|
||||
cv2.imwrite(filename, uv[..., ::-1])
|
||||
|
||||
|
||||
def read_gen(file_name, pil=False):
|
||||
ext = splitext(file_name)[-1]
|
||||
if ext == '.png' or ext == '.jpeg' or ext == '.ppm' or ext == '.jpg':
|
||||
return Image.open(file_name)
|
||||
elif ext == '.bin' or ext == '.raw':
|
||||
return np.load(file_name)
|
||||
elif ext == '.flo':
|
||||
return readFlow(file_name).astype(np.float32)
|
||||
elif ext == '.pfm':
|
||||
flow = readPFM(file_name).astype(np.float32)
|
||||
if len(flow.shape) == 2:
|
||||
return flow
|
||||
else:
|
||||
return flow[:, :, :-1]
|
||||
return []
|
||||
Executable
+82
@@ -0,0 +1,82 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
from scipy import interpolate
|
||||
|
||||
|
||||
class InputPadder:
|
||||
""" Pads images such that dimensions are divisible by 8 """
|
||||
def __init__(self, dims, mode='sintel'):
|
||||
self.ht, self.wd = dims[-2:]
|
||||
pad_ht = (((self.ht // 8) + 1) * 8 - self.ht) % 8
|
||||
pad_wd = (((self.wd // 8) + 1) * 8 - self.wd) % 8
|
||||
if mode == 'sintel':
|
||||
self._pad = [pad_wd//2, pad_wd - pad_wd//2, pad_ht//2, pad_ht - pad_ht//2]
|
||||
else:
|
||||
self._pad = [pad_wd//2, pad_wd - pad_wd//2, 0, pad_ht]
|
||||
|
||||
def pad(self, *inputs):
|
||||
return [F.pad(x, self._pad, mode='replicate') for x in inputs]
|
||||
|
||||
def unpad(self,x):
|
||||
ht, wd = x.shape[-2:]
|
||||
c = [self._pad[2], ht-self._pad[3], self._pad[0], wd-self._pad[1]]
|
||||
return x[..., c[0]:c[1], c[2]:c[3]]
|
||||
|
||||
def forward_interpolate(flow):
|
||||
flow = flow.detach().cpu().numpy()
|
||||
dx, dy = flow[0], flow[1]
|
||||
|
||||
ht, wd = dx.shape
|
||||
x0, y0 = np.meshgrid(np.arange(wd), np.arange(ht))
|
||||
|
||||
x1 = x0 + dx
|
||||
y1 = y0 + dy
|
||||
|
||||
x1 = x1.reshape(-1)
|
||||
y1 = y1.reshape(-1)
|
||||
dx = dx.reshape(-1)
|
||||
dy = dy.reshape(-1)
|
||||
|
||||
valid = (x1 > 0) & (x1 < wd) & (y1 > 0) & (y1 < ht)
|
||||
x1 = x1[valid]
|
||||
y1 = y1[valid]
|
||||
dx = dx[valid]
|
||||
dy = dy[valid]
|
||||
|
||||
flow_x = interpolate.griddata(
|
||||
(x1, y1), dx, (x0, y0), method='nearest', fill_value=0)
|
||||
|
||||
flow_y = interpolate.griddata(
|
||||
(x1, y1), dy, (x0, y0), method='nearest', fill_value=0)
|
||||
|
||||
flow = np.stack([flow_x, flow_y], axis=0)
|
||||
return torch.from_numpy(flow).float()
|
||||
|
||||
|
||||
def bilinear_sampler(img, coords, mode='bilinear', mask=False):
|
||||
""" Wrapper for grid_sample, uses pixel coordinates """
|
||||
H, W = img.shape[-2:]
|
||||
xgrid, ygrid = coords.split([1,1], dim=-1)
|
||||
xgrid = 2*xgrid/(W-1) - 1
|
||||
ygrid = 2*ygrid/(H-1) - 1
|
||||
|
||||
grid = torch.cat([xgrid, ygrid], dim=-1)
|
||||
img = F.grid_sample(img, grid, align_corners=True)
|
||||
|
||||
if mask:
|
||||
mask = (xgrid > -1) & (ygrid > -1) & (xgrid < 1) & (ygrid < 1)
|
||||
return img, mask.float()
|
||||
|
||||
return img
|
||||
|
||||
|
||||
def coords_grid(batch, ht, wd):
|
||||
coords = torch.meshgrid(torch.arange(ht), torch.arange(wd))
|
||||
coords = torch.stack(coords[::-1], dim=0).float()
|
||||
return coords[None].repeat(batch, 1, 1, 1)
|
||||
|
||||
|
||||
def upflow8(flow, mode='bilinear'):
|
||||
new_size = (8 * flow.shape[2], 8 * flow.shape[3])
|
||||
return 8 * F.interpolate(flow, size=new_size, mode=mode, align_corners=True)
|
||||
@@ -0,0 +1,378 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
import os
|
||||
import cv2
|
||||
import numpy as np
|
||||
import scipy.ndimage
|
||||
from PIL import Image
|
||||
|
||||
import torch
|
||||
import torchvision
|
||||
|
||||
from backend import config
|
||||
from backend.inpaint.video.model.modules.flow_comp_raft import RAFT_bi
|
||||
from backend.inpaint.video.model.recurrent_flow_completion import RecurrentFlowCompleteNet
|
||||
from backend.inpaint.video.model.propainter import InpaintGenerator
|
||||
from backend.inpaint.video.core.utils import to_tensors
|
||||
from backend.inpaint.video.model.misc import get_device
|
||||
|
||||
import warnings
|
||||
|
||||
warnings.filterwarnings("ignore")
|
||||
|
||||
|
||||
def binary_mask(mask, th=0.1):
|
||||
mask[mask > th] = 1
|
||||
mask[mask <= th] = 0
|
||||
return mask
|
||||
|
||||
|
||||
# read frame-wise masks
|
||||
def read_mask(mpath, length, size, flow_mask_dilates=8, mask_dilates=5):
|
||||
masks_img = []
|
||||
masks_dilated = []
|
||||
flow_masks = []
|
||||
# 如果传入的直接为numpy array
|
||||
if isinstance(mpath, np.ndarray):
|
||||
masks_img = [Image.fromarray(mpath)]
|
||||
# input single img path
|
||||
else:
|
||||
if isinstance(mpath, str):
|
||||
if mpath.endswith(('jpg', 'jpeg', 'png', 'JPG', 'JPEG', 'PNG')):
|
||||
masks_img = [Image.open(mpath)]
|
||||
else:
|
||||
mnames = sorted(os.listdir(mpath))
|
||||
for mp in mnames:
|
||||
masks_img.append(Image.open(os.path.join(mpath, mp)))
|
||||
|
||||
for mask_img in masks_img:
|
||||
mask_img = np.array(mask_img.convert('L'))
|
||||
|
||||
# Dilate 8 pixel so that all known pixel is trustworthy
|
||||
if flow_mask_dilates > 0:
|
||||
flow_mask_img = scipy.ndimage.binary_dilation(mask_img, iterations=flow_mask_dilates).astype(np.uint8)
|
||||
else:
|
||||
flow_mask_img = binary_mask(mask_img).astype(np.uint8)
|
||||
# Close the small holes inside the foreground objects
|
||||
# flow_mask_img = cv2.morphologyEx(flow_mask_img, cv2.MORPH_CLOSE, np.ones((21, 21),np.uint8)).astype(bool)
|
||||
# flow_mask_img = scipy.ndimage.binary_fill_holes(flow_mask_img).astype(np.uint8)
|
||||
flow_masks.append(Image.fromarray(flow_mask_img * 255))
|
||||
|
||||
if mask_dilates > 0:
|
||||
mask_img = scipy.ndimage.binary_dilation(mask_img, iterations=mask_dilates).astype(np.uint8)
|
||||
else:
|
||||
mask_img = binary_mask(mask_img).astype(np.uint8)
|
||||
masks_dilated.append(Image.fromarray(mask_img * 255))
|
||||
|
||||
if len(masks_img) == 1:
|
||||
flow_masks = flow_masks * length
|
||||
masks_dilated = masks_dilated * length
|
||||
|
||||
return flow_masks, masks_dilated
|
||||
|
||||
|
||||
def extrapolation(video_ori, scale):
|
||||
"""Prepares the data for video outpainting.
|
||||
"""
|
||||
nFrame = len(video_ori)
|
||||
imgW, imgH = video_ori[0].size
|
||||
|
||||
# Defines new FOV.
|
||||
imgH_extr = int(scale[0] * imgH)
|
||||
imgW_extr = int(scale[1] * imgW)
|
||||
imgH_extr = imgH_extr - imgH_extr % 8
|
||||
imgW_extr = imgW_extr - imgW_extr % 8
|
||||
H_start = int((imgH_extr - imgH) / 2)
|
||||
W_start = int((imgW_extr - imgW) / 2)
|
||||
|
||||
# Extrapolates the FOV for video.
|
||||
frames = []
|
||||
for v in video_ori:
|
||||
frame = np.zeros((imgH_extr, imgW_extr, 3), dtype=np.uint8)
|
||||
frame[H_start: H_start + imgH, W_start: W_start + imgW, :] = v
|
||||
frames.append(Image.fromarray(frame))
|
||||
|
||||
# Generates the mask for missing region.
|
||||
masks_dilated = []
|
||||
flow_masks = []
|
||||
|
||||
dilate_h = 4 if H_start > 10 else 0
|
||||
dilate_w = 4 if W_start > 10 else 0
|
||||
mask = np.ones(((imgH_extr, imgW_extr)), dtype=np.uint8)
|
||||
|
||||
mask[H_start + dilate_h: H_start + imgH - dilate_h,
|
||||
W_start + dilate_w: W_start + imgW - dilate_w] = 0
|
||||
flow_masks.append(Image.fromarray(mask * 255))
|
||||
|
||||
mask[H_start: H_start + imgH, W_start: W_start + imgW] = 0
|
||||
masks_dilated.append(Image.fromarray(mask * 255))
|
||||
|
||||
flow_masks = flow_masks * nFrame
|
||||
masks_dilated = masks_dilated * nFrame
|
||||
|
||||
return frames, flow_masks, masks_dilated, (imgW_extr, imgH_extr)
|
||||
|
||||
|
||||
def get_ref_index(mid_neighbor_id, neighbor_ids, length, ref_stride=10, ref_num=-1):
|
||||
ref_index = []
|
||||
if ref_num == -1:
|
||||
for i in range(0, length, ref_stride):
|
||||
if i not in neighbor_ids:
|
||||
ref_index.append(i)
|
||||
else:
|
||||
start_idx = max(0, mid_neighbor_id - ref_stride * (ref_num // 2))
|
||||
end_idx = min(length, mid_neighbor_id + ref_stride * (ref_num // 2))
|
||||
for i in range(start_idx, end_idx, ref_stride):
|
||||
if i not in neighbor_ids:
|
||||
if len(ref_index) > ref_num:
|
||||
break
|
||||
ref_index.append(i)
|
||||
return ref_index
|
||||
|
||||
|
||||
class VideoInpaint:
|
||||
def __init__(self, sub_video_length=config.MAX_PROCESS_NUM, use_fp16=True):
|
||||
self.device = get_device()
|
||||
self.use_fp16 = use_fp16
|
||||
self.use_half = True if self.use_fp16 else False
|
||||
if self.device == torch.device('cpu'):
|
||||
self.use_half = False
|
||||
# Length of sub-video for long video inference.
|
||||
self.sub_video_length = sub_video_length
|
||||
# Length of local neighboring frames.'
|
||||
self.neighbor_length = 10
|
||||
# Mask dilation for video and flow masking
|
||||
self.mask_dilation = 4
|
||||
# Stride of global reference frames
|
||||
self.ref_stride = 10
|
||||
# Iterations for RAFT inference
|
||||
self.raft_iter = 20
|
||||
# Stride of global reference frames
|
||||
self.ref_stride = 10
|
||||
# 设置raft模型
|
||||
self.fix_raft = self.init_raft_model()
|
||||
# 设置fix_flow模型
|
||||
self.fix_flow_complete = self.init_fix_flow_model()
|
||||
# 设置inpaint模型
|
||||
self.model = self.init_inpaint_model()
|
||||
|
||||
def init_raft_model(self):
|
||||
# set up RAFT and flow competition model
|
||||
return RAFT_bi(os.path.join(config.VIDEO_INPAINT_MODEL_PATH, 'raft-things.pth'), self.device)
|
||||
|
||||
def init_fix_flow_model(self):
|
||||
fix_flow_complete_model = RecurrentFlowCompleteNet(
|
||||
os.path.join(config.VIDEO_INPAINT_MODEL_PATH, 'recurrent_flow_completion.pth'))
|
||||
for p in fix_flow_complete_model.parameters():
|
||||
p.requires_grad = False
|
||||
fix_flow_complete_model.to(self.device)
|
||||
fix_flow_complete_model.eval()
|
||||
return fix_flow_complete_model
|
||||
|
||||
def init_inpaint_model(self):
|
||||
# set up ProPainter model
|
||||
return InpaintGenerator(model_path=os.path.join(config.VIDEO_INPAINT_MODEL_PATH, 'ProPainter.pth')).to(
|
||||
self.device).eval()
|
||||
|
||||
def inpaint(self, frames, mask):
|
||||
if isinstance(frames[0], np.ndarray):
|
||||
frames = [Image.fromarray(cv2.cvtColor(f, cv2.COLOR_BGR2RGB)) for f in frames]
|
||||
size = frames[0].size
|
||||
frames_len = len(frames)
|
||||
flow_masks, masks_dilated = read_mask(mask, frames_len, size,
|
||||
flow_mask_dilates=self.mask_dilation,
|
||||
mask_dilates=self.mask_dilation)
|
||||
w, h = size
|
||||
# for saving the masked frames or video
|
||||
masked_frame_for_save = []
|
||||
for i in range(len(frames)):
|
||||
mask_ = np.expand_dims(np.array(masks_dilated[i]), 2).repeat(3, axis=2) / 255.
|
||||
img = np.array(frames[i])
|
||||
green = np.zeros([h, w, 3])
|
||||
green[:, :, 1] = 255
|
||||
alpha = 0.6
|
||||
# alpha = 1.0
|
||||
fuse_img = (1 - alpha) * img + alpha * green
|
||||
fuse_img = mask_ * fuse_img + (1 - mask_) * img
|
||||
masked_frame_for_save.append(fuse_img.astype(np.uint8))
|
||||
|
||||
frames_inp = [np.array(f).astype(np.uint8) for f in frames]
|
||||
frames = to_tensors()(frames).unsqueeze(0) * 2 - 1
|
||||
flow_masks = to_tensors()(flow_masks).unsqueeze(0)
|
||||
masks_dilated = to_tensors()(masks_dilated).unsqueeze(0)
|
||||
frames, flow_masks, masks_dilated = frames.to(self.device), flow_masks.to(self.device), masks_dilated.to(
|
||||
self.device)
|
||||
video_length = frames.size(1)
|
||||
with torch.no_grad():
|
||||
# ---- compute flow ----
|
||||
if frames.size(-1) <= 640:
|
||||
short_clip_len = 12
|
||||
elif frames.size(-1) <= 720:
|
||||
short_clip_len = 8
|
||||
elif frames.size(-1) <= 1280:
|
||||
short_clip_len = 4
|
||||
else:
|
||||
short_clip_len = 2
|
||||
|
||||
# use fp32 for RAFT
|
||||
if frames.size(1) > short_clip_len:
|
||||
gt_flows_f_list, gt_flows_b_list = [], []
|
||||
for f in range(0, video_length, short_clip_len):
|
||||
end_f = min(video_length, f + short_clip_len)
|
||||
if f == 0:
|
||||
flows_f, flows_b = self.fix_raft(frames[:, f:end_f], iters=self.raft_iter)
|
||||
else:
|
||||
flows_f, flows_b = self.fix_raft(frames[:, f - 1:end_f], iters=self.raft_iter)
|
||||
gt_flows_f_list.append(flows_f)
|
||||
gt_flows_b_list.append(flows_b)
|
||||
torch.cuda.empty_cache()
|
||||
gt_flows_f = torch.cat(gt_flows_f_list, dim=1)
|
||||
gt_flows_b = torch.cat(gt_flows_b_list, dim=1)
|
||||
gt_flows_bi = (gt_flows_f, gt_flows_b)
|
||||
else:
|
||||
gt_flows_bi = self.fix_raft(frames, iters=self.raft_iter)
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
if self.use_half:
|
||||
frames, flow_masks, masks_dilated = frames.half(), flow_masks.half(), masks_dilated.half()
|
||||
gt_flows_bi = (gt_flows_bi[0].half(), gt_flows_bi[1].half())
|
||||
fix_flow_complete = self.fix_flow_complete.half()
|
||||
self.model = self.model.half()
|
||||
|
||||
# ---- complete flow ----
|
||||
flow_length = gt_flows_bi[0].size(1)
|
||||
if flow_length > self.sub_video_length:
|
||||
pred_flows_f, pred_flows_b = [], []
|
||||
pad_len = 5
|
||||
for f in range(0, flow_length, self.sub_video_length):
|
||||
s_f = max(0, f - pad_len)
|
||||
e_f = min(flow_length, f + self.sub_video_length + pad_len)
|
||||
pad_len_s = max(0, f) - s_f
|
||||
pad_len_e = e_f - min(flow_length, f + self.sub_video_length)
|
||||
pred_flows_bi_sub, _ = fix_flow_complete.forward_bidirect_flow(
|
||||
(gt_flows_bi[0][:, s_f:e_f], gt_flows_bi[1][:, s_f:e_f]),
|
||||
flow_masks[:, s_f:e_f + 1])
|
||||
pred_flows_bi_sub = fix_flow_complete.combine_flow(
|
||||
(gt_flows_bi[0][:, s_f:e_f], gt_flows_bi[1][:, s_f:e_f]),
|
||||
pred_flows_bi_sub,
|
||||
flow_masks[:, s_f:e_f + 1])
|
||||
|
||||
pred_flows_f.append(pred_flows_bi_sub[0][:, pad_len_s:e_f - s_f - pad_len_e])
|
||||
pred_flows_b.append(pred_flows_bi_sub[1][:, pad_len_s:e_f - s_f - pad_len_e])
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
pred_flows_f = torch.cat(pred_flows_f, dim=1)
|
||||
pred_flows_b = torch.cat(pred_flows_b, dim=1)
|
||||
pred_flows_bi = (pred_flows_f, pred_flows_b)
|
||||
else:
|
||||
pred_flows_bi, _ = fix_flow_complete.forward_bidirect_flow(gt_flows_bi, flow_masks)
|
||||
pred_flows_bi = fix_flow_complete.combine_flow(gt_flows_bi, pred_flows_bi, flow_masks)
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
# ---- image propagation ----
|
||||
masked_frames = frames * (1 - masks_dilated)
|
||||
# ensure a minimum of 100 frames for image propagation
|
||||
subvideo_length_img_prop = min(100, self.sub_video_length)
|
||||
if video_length > subvideo_length_img_prop:
|
||||
updated_frames, updated_masks = [], []
|
||||
pad_len = 10
|
||||
for f in range(0, video_length, subvideo_length_img_prop):
|
||||
s_f = max(0, f - pad_len)
|
||||
e_f = min(video_length, f + subvideo_length_img_prop + pad_len)
|
||||
pad_len_s = max(0, f) - s_f
|
||||
pad_len_e = e_f - min(video_length, f + subvideo_length_img_prop)
|
||||
|
||||
b, t, _, _, _ = masks_dilated[:, s_f:e_f].size()
|
||||
pred_flows_bi_sub = (pred_flows_bi[0][:, s_f:e_f - 1], pred_flows_bi[1][:, s_f:e_f - 1])
|
||||
prop_imgs_sub, updated_local_masks_sub = self.model.img_propagation(masked_frames[:, s_f:e_f],
|
||||
pred_flows_bi_sub,
|
||||
masks_dilated[:, s_f:e_f],
|
||||
'nearest')
|
||||
updated_frames_sub = frames[:, s_f:e_f] * (1 - masks_dilated[:, s_f:e_f]) + prop_imgs_sub.view(b, t, 3, h, w) * masks_dilated[:, s_f:e_f]
|
||||
updated_masks_sub = updated_local_masks_sub.view(b, t, 1, h, w)
|
||||
updated_frames.append(updated_frames_sub[:, pad_len_s:e_f - s_f - pad_len_e])
|
||||
updated_masks.append(updated_masks_sub[:, pad_len_s:e_f - s_f - pad_len_e])
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
updated_frames = torch.cat(updated_frames, dim=1)
|
||||
updated_masks = torch.cat(updated_masks, dim=1)
|
||||
else:
|
||||
b, t, _, _, _ = masks_dilated.size()
|
||||
prop_imgs, updated_local_masks = self.model.img_propagation(masked_frames, pred_flows_bi, masks_dilated,
|
||||
'nearest')
|
||||
updated_frames = frames * (1 - masks_dilated) + prop_imgs.view(b, t, 3, h, w) * masks_dilated
|
||||
updated_masks = updated_local_masks.view(b, t, 1, h, w)
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
ori_frames = frames_inp
|
||||
comp_frames = [None] * video_length
|
||||
|
||||
neighbor_stride = self.neighbor_length // 2
|
||||
if video_length > self.sub_video_length:
|
||||
ref_num = self.sub_video_length // self.ref_stride
|
||||
else:
|
||||
ref_num = -1
|
||||
|
||||
# ---- feature propagation + transformer ----
|
||||
for f in range(0, video_length, neighbor_stride):
|
||||
neighbor_ids = [
|
||||
i for i in range(max(0, f - neighbor_stride),
|
||||
min(video_length, f + neighbor_stride + 1))
|
||||
]
|
||||
ref_ids = get_ref_index(f, neighbor_ids, video_length, self.ref_stride, ref_num)
|
||||
selected_imgs = updated_frames[:, neighbor_ids + ref_ids, :, :, :]
|
||||
selected_masks = masks_dilated[:, neighbor_ids + ref_ids, :, :, :]
|
||||
selected_update_masks = updated_masks[:, neighbor_ids + ref_ids, :, :, :]
|
||||
selected_pred_flows_bi = (
|
||||
pred_flows_bi[0][:, neighbor_ids[:-1], :, :, :], pred_flows_bi[1][:, neighbor_ids[:-1], :, :, :])
|
||||
|
||||
with torch.no_grad():
|
||||
# 1.0 indicates mask
|
||||
l_t = len(neighbor_ids)
|
||||
pred_img = self.model(selected_imgs, selected_pred_flows_bi, selected_masks, selected_update_masks, l_t)
|
||||
pred_img = pred_img.view(-1, 3, h, w)
|
||||
pred_img = (pred_img + 1) / 2
|
||||
pred_img = pred_img.cpu().permute(0, 2, 3, 1).numpy() * 255
|
||||
binary_masks = masks_dilated[0, neighbor_ids, :, :, :].cpu().permute(
|
||||
0, 2, 3, 1).numpy().astype(np.uint8)
|
||||
for i in range(len(neighbor_ids)):
|
||||
idx = neighbor_ids[i]
|
||||
img = np.array(pred_img[i]).astype(np.uint8) * binary_masks[i] \
|
||||
+ ori_frames[idx] * (1 - binary_masks[i])
|
||||
if comp_frames[idx] is None:
|
||||
comp_frames[idx] = img
|
||||
else:
|
||||
comp_frames[idx] = comp_frames[idx].astype(np.float32) * 0.5 + img.astype(np.float32) * 0.5
|
||||
comp_frames[idx] = comp_frames[idx].astype(np.uint8)
|
||||
torch.cuda.empty_cache()
|
||||
# save videos frame
|
||||
comp_frames = [cv2.cvtColor(i, cv2.COLOR_RGB2BGR) for i in comp_frames]
|
||||
return comp_frames
|
||||
|
||||
|
||||
def read_frames(v_path):
|
||||
video_cap = cv2.VideoCapture(v_path)
|
||||
video_frames = []
|
||||
while True:
|
||||
ret, frame = video_cap.read()
|
||||
if not ret:
|
||||
break
|
||||
video_frames.append(frame)
|
||||
video_frames = [Image.fromarray(f) for f in video_frames]
|
||||
return video_frames
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
# VideoInpaint
|
||||
video_inpaint = VideoInpaint(sub_video_length=80)
|
||||
frames = read_frames('/home/yao/Documents/Project/video-subtitle-remover/local_test/test1.mp4')
|
||||
mask = cv2.imread('/home/yao/Documents/Project/video-subtitle-remover/local_test/test1_mask.png')
|
||||
inpainted_frames = video_inpaint.inpaint(frames, mask)
|
||||
save_root = '/home/yao/Documents/Project/video-subtitle-remover/local_test/'
|
||||
video_out_path = os.path.join(save_root, 'inpaint_out.mp4')
|
||||
print("size: ", inpainted_frames[0].shape)
|
||||
video_writer = cv2.VideoWriter(video_out_path, cv2.VideoWriter_fourcc(*'mp4v'), 24, (640, 360))
|
||||
for comp_frame in inpainted_frames:
|
||||
video_writer.write(comp_frame)
|
||||
video_writer.release()
|
||||
print(f'\nAll results are saved in {save_root}')
|
||||
|
||||
+684
@@ -0,0 +1,684 @@
|
||||
import shutil
|
||||
import subprocess
|
||||
import os
|
||||
from pathlib import Path
|
||||
import threading
|
||||
import cv2
|
||||
import sys
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
from backend.inpaint.lama_inpaint import LamaInpaint
|
||||
from backend.inpaint.video_inpaint import VideoInpaint
|
||||
from backend.tools.inpaint_tools import create_mask, batch_generator
|
||||
import config
|
||||
import importlib
|
||||
import platform
|
||||
import tempfile
|
||||
import torch
|
||||
import multiprocessing
|
||||
from shapely.geometry import Polygon
|
||||
import time
|
||||
from tqdm import tqdm
|
||||
from tools.infer import utility
|
||||
from tools.infer.predict_det import TextDetector
|
||||
|
||||
|
||||
class SubtitleDetect:
|
||||
"""
|
||||
文本框检测类，用于检测视频帧中是否存在文本框
|
||||
"""
|
||||
|
||||
def __init__(self, video_path, sub_area=None):
|
||||
# 获取参数对象
|
||||
importlib.reload(config)
|
||||
args = utility.parse_args()
|
||||
args.det_algorithm = 'DB'
|
||||
args.det_model_dir = config.DET_MODEL_PATH
|
||||
self.text_detector = TextDetector(args)
|
||||
self.video_path = video_path
|
||||
self.sub_area = sub_area
|
||||
|
||||
def detect_subtitle(self, img):
|
||||
dt_boxes, elapse = self.text_detector(img)
|
||||
return dt_boxes, elapse
|
||||
|
||||
@staticmethod
|
||||
def get_coordinates(dt_box):
|
||||
"""
|
||||
从返回的检测框中获取坐标
|
||||
:param dt_box 检测框返回结果
|
||||
:return list 坐标点列表
|
||||
"""
|
||||
coordinate_list = list()
|
||||
if isinstance(dt_box, list):
|
||||
for i in dt_box:
|
||||
i = list(i)
|
||||
(x1, y1) = int(i[0][0]), int(i[0][1])
|
||||
(x2, y2) = int(i[1][0]), int(i[1][1])
|
||||
(x3, y3) = int(i[2][0]), int(i[2][1])
|
||||
(x4, y4) = int(i[3][0]), int(i[3][1])
|
||||
xmin = max(x1, x4)
|
||||
xmax = min(x2, x3)
|
||||
ymin = max(y1, y2)
|
||||
ymax = min(y3, y4)
|
||||
coordinate_list.append((xmin, xmax, ymin, ymax))
|
||||
return coordinate_list
|
||||
|
||||
def find_subtitle_frame_no(self, sub_remover=None):
|
||||
video_cap = cv2.VideoCapture(self.video_path)
|
||||
frame_count = video_cap.get(cv2.CAP_PROP_FRAME_COUNT)
|
||||
tbar = tqdm(total=int(frame_count), unit='frame', position=0, file=sys.__stdout__, desc='Subtitle Finding')
|
||||
current_frame_no = 0
|
||||
subtitle_frame_no_box_dict = {}
|
||||
print('[Processing] start finding subtitles...')
|
||||
while video_cap.isOpened():
|
||||
ret, frame = video_cap.read()
|
||||
# 如果读取视频帧失败（视频读到最后一帧）
|
||||
if not ret:
|
||||
break
|
||||
# 读取视频帧成功
|
||||
current_frame_no += 1
|
||||
dt_boxes, elapse = self.detect_subtitle(frame)
|
||||
coordinate_list = self.get_coordinates(dt_boxes.tolist())
|
||||
if coordinate_list:
|
||||
temp_list = []
|
||||
for coordinate in coordinate_list:
|
||||
xmin, xmax, ymin, ymax = coordinate
|
||||
if self.sub_area is not None:
|
||||
s_ymin, s_ymax, s_xmin, s_xmax = self.sub_area
|
||||
if (s_xmin <= xmin and xmax <= s_xmax
|
||||
and s_ymin <= ymin
|
||||
and ymax <= s_ymax):
|
||||
temp_list.append((xmin, xmax, ymin, ymax))
|
||||
else:
|
||||
temp_list.append((xmin, xmax, ymin, ymax))
|
||||
if len(temp_list) > 0:
|
||||
subtitle_frame_no_box_dict[current_frame_no] = temp_list
|
||||
tbar.update(1)
|
||||
if sub_remover:
|
||||
sub_remover.progress_total = (100 * float(current_frame_no) / float(frame_count)) // 2
|
||||
subtitle_frame_no_box_dict = self.unify_regions(subtitle_frame_no_box_dict)
|
||||
# if config.UNITE_COORDINATES:
|
||||
# subtitle_frame_no_box_dict = self.get_subtitle_frame_no_box_dict_with_united_coordinates(subtitle_frame_no_box_dict)
|
||||
# if sub_remover is not None:
|
||||
# try:
|
||||
# # 当帧数大于1时，说明并非图片或单帧
|
||||
# if sub_remover.frame_count > 1:
|
||||
# subtitle_frame_no_box_dict = self.filter_mistake_sub_area(subtitle_frame_no_box_dict,
|
||||
# sub_remover.fps)
|
||||
# except Exception:
|
||||
# pass
|
||||
# subtitle_frame_no_box_dict = self.prevent_missed_detection(subtitle_frame_no_box_dict)
|
||||
print('[Finished] Finished finding subtitles...')
|
||||
new_subtitle_frame_no_box_dict = dict()
|
||||
for key in subtitle_frame_no_box_dict.keys():
|
||||
if len(subtitle_frame_no_box_dict[key]) > 0:
|
||||
new_subtitle_frame_no_box_dict[key] = subtitle_frame_no_box_dict[key]
|
||||
return new_subtitle_frame_no_box_dict
|
||||
|
||||
@staticmethod
|
||||
def are_similar(region1, region2):
|
||||
"""判断两个区域是否相似。"""
|
||||
xmin1, xmax1, ymin1, ymax1 = region1
|
||||
xmin2, xmax2, ymin2, ymax2 = region2
|
||||
|
||||
return abs(xmin1 - xmin2) <= config.PIXEL_TOLERANCE_X and abs(xmax1 - xmax2) <= config.PIXEL_TOLERANCE_X and \
|
||||
abs(ymin1 - ymin2) <= config.PIXEL_TOLERANCE_Y and abs(ymax1 - ymax2) <= config.PIXEL_TOLERANCE_Y
|
||||
|
||||
def unify_regions(self, raw_regions):
|
||||
"""将连续相似的区域统一，保持列表结构。"""
|
||||
keys = sorted(raw_regions.keys()) # 对键进行排序以确保它们是连续的
|
||||
unified_regions = {}
|
||||
|
||||
# 初始化
|
||||
last_key = keys[0]
|
||||
unify_value_map = {last_key: raw_regions[last_key]}
|
||||
|
||||
for key in keys[1:]:
|
||||
current_regions = raw_regions[key]
|
||||
|
||||
# 新增一个列表来存放匹配过的标准区间
|
||||
new_unify_values = []
|
||||
|
||||
for idx, region in enumerate(current_regions):
|
||||
last_standard_region = unify_value_map[last_key][idx] if idx < len(unify_value_map[last_key]) else None
|
||||
|
||||
# 如果当前的区间与前一个键的对应区间相似，我们统一它们
|
||||
if last_standard_region and self.are_similar(region, last_standard_region):
|
||||
new_unify_values.append(last_standard_region)
|
||||
else:
|
||||
new_unify_values.append(region)
|
||||
|
||||
# 更新unify_value_map为最新的区间值
|
||||
unify_value_map[key] = new_unify_values
|
||||
last_key = key
|
||||
|
||||
# 将最终统一后的结果传递给unified_regions
|
||||
for key in keys:
|
||||
unified_regions[key] = unify_value_map[key]
|
||||
return unified_regions
|
||||
|
||||
@staticmethod
|
||||
def find_continuous_ranges(subtitle_frame_no_box_dict):
|
||||
"""
|
||||
获取字幕出现的起始帧号与结束帧号
|
||||
"""
|
||||
numbers = sorted(list(subtitle_frame_no_box_dict.keys()))
|
||||
ranges = []
|
||||
start = numbers[0] # 初始区间开始值
|
||||
|
||||
for i in range(1, len(numbers)):
|
||||
# 如果当前数字与前一个数字间隔超过1，
|
||||
# 则上一个区间结束，记录当前区间的开始与结束
|
||||
if numbers[i] - numbers[i - 1] != 1:
|
||||
end = numbers[i - 1] # 则该数字是当前连续区间的终点
|
||||
ranges.append((start, end))
|
||||
start = numbers[i] # 开始下一个连续区间
|
||||
# 添加最后一个区间
|
||||
ranges.append((start, numbers[-1]))
|
||||
return ranges
|
||||
|
||||
@staticmethod
|
||||
def find_continuous_ranges_with_same_mask(subtitle_frame_no_box_dict):
|
||||
numbers = sorted(list(subtitle_frame_no_box_dict.keys()))
|
||||
ranges = []
|
||||
start = numbers[0] # 初始区间开始值
|
||||
for i in range(1, len(numbers)):
|
||||
# 如果当前帧号与前一个帧号间隔超过1，
|
||||
# 则上一个区间结束，记录当前区间的开始与结束
|
||||
if numbers[i] - numbers[i - 1] != 1:
|
||||
end = numbers[i - 1] # 则该数字是当前连续区间的终点
|
||||
ranges.append((start, end))
|
||||
start = numbers[i] # 开始下一个连续区间
|
||||
# 如果当前帧号与前一个帧号间隔为1，且当前帧号对应的坐标点与上一帧号对应的坐标点不一致
|
||||
# 记录当前区间的开始与结束
|
||||
if numbers[i] - numbers[i - 1] == 1:
|
||||
if subtitle_frame_no_box_dict[numbers[i]] != subtitle_frame_no_box_dict[numbers[i - 1]]:
|
||||
end = numbers[i - 1] # 则该数字是当前连续区间的终点
|
||||
ranges.append((start, end))
|
||||
start = numbers[i] # 开始下一个连续区间
|
||||
# 添加最后一个区间
|
||||
ranges.append((start, numbers[-1]))
|
||||
return ranges
|
||||
|
||||
@staticmethod
|
||||
def sub_area_to_polygon(sub_area):
|
||||
"""
|
||||
xmin, xmax, ymin, ymax = sub_area
|
||||
"""
|
||||
s_xmin = sub_area[0]
|
||||
s_xmax = sub_area[1]
|
||||
s_ymin = sub_area[2]
|
||||
s_ymax = sub_area[3]
|
||||
return Polygon([[s_xmin, s_ymin], [s_xmax, s_ymin], [s_xmax, s_ymax], [s_xmin, s_ymax]])
|
||||
|
||||
def compute_iou(self, box1, box2):
|
||||
box1_polygon = self.sub_area_to_polygon(box1)
|
||||
box2_polygon = self.sub_area_to_polygon(box2)
|
||||
intersection = box1_polygon.intersection(box2_polygon)
|
||||
if intersection.is_empty:
|
||||
return -1
|
||||
else:
|
||||
union_area = (box1_polygon.area + box2_polygon.area - intersection.area)
|
||||
if union_area > 0:
|
||||
intersection_area_rate = intersection.area / union_area
|
||||
else:
|
||||
intersection_area_rate = 0
|
||||
return intersection_area_rate
|
||||
|
||||
def get_area_max_box_dict(self, sub_frame_no_list_continuous, subtitle_frame_no_box_dict):
|
||||
_area_max_box_dict = dict()
|
||||
for start_no, end_no in sub_frame_no_list_continuous:
|
||||
# 寻找面积最大文本框
|
||||
current_no = start_no
|
||||
# 查找当前区间矩形框最大面积
|
||||
area_max_box_list = []
|
||||
while current_no <= end_no:
|
||||
for coord in subtitle_frame_no_box_dict[current_no]:
|
||||
# 取出每一个文本框坐标
|
||||
xmin, xmax, ymin, ymax = coord
|
||||
# 计算当前文本框坐标面积
|
||||
current_area = abs(xmax - xmin) * abs(ymax - ymin)
|
||||
# 如果区间最大框列表为空，则当前面积为区间最大面积
|
||||
if len(area_max_box_list) < 1:
|
||||
area_max_box_list.append({
|
||||
'area': current_area,
|
||||
'xmin': xmin,
|
||||
'xmax': xmax,
|
||||
'ymin': ymin,
|
||||
'ymax': ymax
|
||||
})
|
||||
# 如果列表非空，判断当前文本框是与区间最大文本框在同一区域
|
||||
else:
|
||||
has_same_position = False
|
||||
# 遍历每个区间最大文本框，判断当前文本框位置是否与区间最大文本框列表的某个文本框位于同一行且交叉
|
||||
for area_max_box in area_max_box_list:
|
||||
if (area_max_box['ymin'] - config.TOLERANCE_Y <= ymin
|
||||
and ymax <= area_max_box['ymax'] + config.TOLERANCE_Y):
|
||||
if self.compute_iou((xmin, xmax, ymin, ymax), (
|
||||
area_max_box['xmin'], area_max_box['xmax'], area_max_box['ymin'],
|
||||
area_max_box['ymax'])) != -1:
|
||||
# 如果高度差异不一样
|
||||
if abs(abs(area_max_box['ymax'] - area_max_box['ymin']) - abs(
|
||||
ymax - ymin)) < config.THRESHOLD_HEIGHT_DIFFERENCE:
|
||||
has_same_position = True
|
||||
# 如果在同一行，则计算当前面积是不是最大
|
||||
# 判断面积大小，若当前面积更大，则将当前行的最大区域坐标点更新
|
||||
if has_same_position and current_area > area_max_box['area']:
|
||||
area_max_box['area'] = current_area
|
||||
area_max_box['xmin'] = xmin
|
||||
area_max_box['xmax'] = xmax
|
||||
area_max_box['ymin'] = ymin
|
||||
area_max_box['ymax'] = ymax
|
||||
# 如果遍历了所有的区间最大文本框列表，发现是新的一行，则直接添加
|
||||
if not has_same_position:
|
||||
new_large_area = {
|
||||
'area': current_area,
|
||||
'xmin': xmin,
|
||||
'xmax': xmax,
|
||||
'ymin': ymin,
|
||||
'ymax': ymax
|
||||
}
|
||||
if new_large_area not in area_max_box_list:
|
||||
area_max_box_list.append(new_large_area)
|
||||
break
|
||||
current_no += 1
|
||||
_area_max_box_list = list()
|
||||
for area_max_box in area_max_box_list:
|
||||
if area_max_box not in _area_max_box_list:
|
||||
_area_max_box_list.append(area_max_box)
|
||||
_area_max_box_dict[f'{start_no}->{end_no}'] = _area_max_box_list
|
||||
return _area_max_box_dict
|
||||
|
||||
def get_subtitle_frame_no_box_dict_with_united_coordinates(self, subtitle_frame_no_box_dict):
|
||||
"""
|
||||
将多个视频帧的文本区域坐标统一
|
||||
"""
|
||||
subtitle_frame_no_box_dict_with_united_coordinates = dict()
|
||||
frame_no_list = self.find_continuous_ranges_with_same_mask(subtitle_frame_no_box_dict)
|
||||
area_max_box_dict = self.get_area_max_box_dict(frame_no_list, subtitle_frame_no_box_dict)
|
||||
for start_no, end_no in frame_no_list:
|
||||
current_no = start_no
|
||||
while True:
|
||||
area_max_box_list = area_max_box_dict[f'{start_no}->{end_no}']
|
||||
current_boxes = subtitle_frame_no_box_dict[current_no]
|
||||
new_subtitle_frame_no_box_list = []
|
||||
for current_box in current_boxes:
|
||||
current_xmin, current_xmax, current_ymin, current_ymax = current_box
|
||||
for max_box in area_max_box_list:
|
||||
large_xmin = max_box['xmin']
|
||||
large_xmax = max_box['xmax']
|
||||
large_ymin = max_box['ymin']
|
||||
large_ymax = max_box['ymax']
|
||||
box1 = (current_xmin, current_xmax, current_ymin, current_ymax)
|
||||
box2 = (large_xmin, large_xmax, large_ymin, large_ymax)
|
||||
res = self.compute_iou(box1, box2)
|
||||
if res != -1:
|
||||
new_subtitle_frame_no_box = (large_xmin, large_xmax, large_ymin, large_ymax)
|
||||
if new_subtitle_frame_no_box not in new_subtitle_frame_no_box_list:
|
||||
new_subtitle_frame_no_box_list.append(new_subtitle_frame_no_box)
|
||||
subtitle_frame_no_box_dict_with_united_coordinates[current_no] = new_subtitle_frame_no_box_list
|
||||
current_no += 1
|
||||
if current_no > end_no:
|
||||
break
|
||||
return subtitle_frame_no_box_dict_with_united_coordinates
|
||||
|
||||
def prevent_missed_detection(self, subtitle_frame_no_box_dict):
|
||||
"""
|
||||
添加额外的文本框，防止漏检
|
||||
"""
|
||||
frame_no_list = self.find_continuous_ranges_with_same_mask(subtitle_frame_no_box_dict)
|
||||
for start_no, end_no in frame_no_list:
|
||||
current_no = start_no
|
||||
while True:
|
||||
current_box_list = subtitle_frame_no_box_dict[current_no]
|
||||
if current_no + 1 != end_no and (current_no + 1) in subtitle_frame_no_box_dict.keys():
|
||||
next_box_list = subtitle_frame_no_box_dict[current_no + 1]
|
||||
if set(current_box_list).issubset(set(next_box_list)):
|
||||
subtitle_frame_no_box_dict[current_no] = subtitle_frame_no_box_dict[current_no + 1]
|
||||
current_no += 1
|
||||
if current_no > end_no:
|
||||
break
|
||||
return subtitle_frame_no_box_dict
|
||||
|
||||
@staticmethod
|
||||
def get_frequency_in_range(sub_frame_no_list_continuous, subtitle_frame_no_box_dict):
|
||||
sub_area_with_frequency = {}
|
||||
for start_no, end_no in sub_frame_no_list_continuous:
|
||||
current_no = start_no
|
||||
while True:
|
||||
current_box_list = subtitle_frame_no_box_dict[current_no]
|
||||
for current_box in current_box_list:
|
||||
if str(current_box) not in sub_area_with_frequency.keys():
|
||||
sub_area_with_frequency[f'{current_box}'] = 1
|
||||
else:
|
||||
sub_area_with_frequency[f'{current_box}'] += 1
|
||||
current_no += 1
|
||||
if current_no > end_no:
|
||||
break
|
||||
return sub_area_with_frequency
|
||||
|
||||
def filter_mistake_sub_area(self, subtitle_frame_no_box_dict, fps):
|
||||
"""
|
||||
过滤错误的字幕区域
|
||||
"""
|
||||
sub_frame_no_list_continuous = self.find_continuous_ranges_with_same_mask(subtitle_frame_no_box_dict)
|
||||
sub_area_with_frequency = self.get_frequency_in_range(sub_frame_no_list_continuous, subtitle_frame_no_box_dict)
|
||||
correct_sub_area = []
|
||||
for sub_area in sub_area_with_frequency.keys():
|
||||
if sub_area_with_frequency[sub_area] >= (fps // 2):
|
||||
correct_sub_area.append(sub_area)
|
||||
else:
|
||||
print(f'drop {sub_area}')
|
||||
correct_subtitle_frame_no_box_dict = dict()
|
||||
for frame_no in subtitle_frame_no_box_dict.keys():
|
||||
current_box_list = subtitle_frame_no_box_dict[frame_no]
|
||||
new_box_list = []
|
||||
for current_box in current_box_list:
|
||||
if str(current_box) in correct_sub_area and current_box not in new_box_list:
|
||||
new_box_list.append(current_box)
|
||||
correct_subtitle_frame_no_box_dict[frame_no] = new_box_list
|
||||
return correct_subtitle_frame_no_box_dict
|
||||
|
||||
|
||||
class SubtitleRemover:
|
||||
def __init__(self, vd_path, sub_area=None):
|
||||
importlib.reload(config)
|
||||
# 线程锁
|
||||
self.lock = threading.RLock()
|
||||
# 用户指定的字幕区域位置
|
||||
self.sub_area = sub_area
|
||||
# 判断是否为图片
|
||||
self.is_picture = False
|
||||
if str(vd_path).endswith(('.bmp', '.dib', '.png', '.jpg', '.jpeg', '.pbm', '.pgm', '.ppm', '.tif', '.tiff')):
|
||||
self.sub_area = None
|
||||
self.is_picture = True
|
||||
# 视频路径
|
||||
self.video_path = vd_path
|
||||
self.video_cap = cv2.VideoCapture(vd_path)
|
||||
# 通过视频路径获取视频名称
|
||||
self.vd_name = Path(self.video_path).stem
|
||||
# 视频帧总数
|
||||
self.frame_count = self.video_cap.get(cv2.CAP_PROP_FRAME_COUNT)
|
||||
# 视频帧率
|
||||
self.fps = self.video_cap.get(cv2.CAP_PROP_FPS)
|
||||
# 视频尺寸
|
||||
self.size = (int(self.video_cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(self.video_cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
|
||||
self.mask_size = (int(self.video_cap.get(cv2.CAP_PROP_FRAME_HEIGHT)), int(self.video_cap.get(cv2.CAP_PROP_FRAME_WIDTH)))
|
||||
self.frame_height = int(self.video_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||||
self.frame_width = int(self.video_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||||
# 创建字幕检测对象
|
||||
self.sub_detector = SubtitleDetect(self.video_path, self.sub_area)
|
||||
# 创建视频临时对象，windows下delete=True会有permission denied的报错
|
||||
self.video_temp_file = tempfile.NamedTemporaryFile(suffix='.mp4', delete=False)
|
||||
# 创建视频写对象
|
||||
self.video_writer = cv2.VideoWriter(self.video_temp_file.name, cv2.VideoWriter_fourcc(*'mp4v'), self.fps,
|
||||
self.size)
|
||||
self.video_out_name = os.path.join(os.path.dirname(self.video_path), f'{self.vd_name}_no_sub.mp4')
|
||||
self.video_inpaint = None
|
||||
self.lama_inpaint = None
|
||||
self.ext = os.path.splitext(vd_path)[-1]
|
||||
if self.is_picture:
|
||||
pic_dir = os.path.join(os.path.dirname(self.video_path), 'no_sub')
|
||||
if not os.path.exists(pic_dir):
|
||||
os.makedirs(pic_dir)
|
||||
self.video_out_name = os.path.join(pic_dir, f'{self.vd_name}{self.ext}')
|
||||
if torch.cuda.is_available():
|
||||
print('use GPU for acceleration')
|
||||
# 总处理进度
|
||||
self.progress_total = 0
|
||||
self.progress_remover = 0
|
||||
self.isFinished = False
|
||||
# 预览帧
|
||||
self.preview_frame = None
|
||||
# 是否将原音频嵌入到去除字幕后的视频
|
||||
self.is_successful_merged = False
|
||||
|
||||
@staticmethod
|
||||
def get_coordinates(dt_box):
|
||||
"""
|
||||
从返回的检测框中获取坐标
|
||||
:param dt_box 检测框返回结果
|
||||
:return list 坐标点列表
|
||||
"""
|
||||
coordinate_list = list()
|
||||
if isinstance(dt_box, list):
|
||||
for i in dt_box:
|
||||
i = list(i)
|
||||
(x1, y1) = int(i[0][0]), int(i[0][1])
|
||||
(x2, y2) = int(i[1][0]), int(i[1][1])
|
||||
(x3, y3) = int(i[2][0]), int(i[2][1])
|
||||
(x4, y4) = int(i[3][0]), int(i[3][1])
|
||||
xmin = max(x1, x4)
|
||||
xmax = min(x2, x3)
|
||||
ymin = max(y1, y2)
|
||||
ymax = min(y3, y4)
|
||||
coordinate_list.append((xmin, xmax, ymin, ymax))
|
||||
return coordinate_list
|
||||
|
||||
@staticmethod
|
||||
def is_current_frame_no_start(frame_no, continuous_frame_no_list):
|
||||
"""
|
||||
判断给定的帧号是否为开头，是的话返回结束帧号，不是的话返回-1
|
||||
"""
|
||||
for start_no, end_no in continuous_frame_no_list:
|
||||
if start_no == frame_no:
|
||||
return True
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def find_frame_no_end(frame_no, continuous_frame_no_list):
|
||||
"""
|
||||
判断给定的帧号是否为开头，是的话返回结束帧号，不是的话返回-1
|
||||
"""
|
||||
for start_no, end_no in continuous_frame_no_list:
|
||||
if start_no <= frame_no <= end_no:
|
||||
return end_no
|
||||
return -1
|
||||
|
||||
def update_progress(self, tbar, increment):
|
||||
tbar.update(increment)
|
||||
current_percentage = (tbar.n / tbar.total) * 100
|
||||
self.progress_remover = int(current_percentage) // 2
|
||||
self.progress_total = 50 + self.progress_remover
|
||||
|
||||
def run(self):
|
||||
|
||||
# 记录开始时间
|
||||
start_time = time.time()
|
||||
# 寻找字幕帧
|
||||
self.progress_total = 0
|
||||
sub_list = self.sub_detector.find_subtitle_frame_no(sub_remover=self)
|
||||
# 测试代码
|
||||
# from test1_dict_raw import test1_raw
|
||||
# sub_list = self.sub_detector.unify_regions(test1_raw)
|
||||
continuous_frame_no_list = self.sub_detector.find_continuous_ranges_with_same_mask(sub_list)
|
||||
tbar = tqdm(total=int(self.frame_count), unit='frame', position=0, file=sys.__stdout__,
|
||||
desc='Subtitle Removing')
|
||||
print('[Processing] start removing subtitles...')
|
||||
|
||||
if self.is_picture:
|
||||
self.lama_inpaint = LamaInpaint()
|
||||
original_frame = cv2.imread(self.video_path)
|
||||
mask = create_mask(original_frame.shape[0:2], sub_list[1])
|
||||
frame = self.lama_inpaint(original_frame, mask)
|
||||
cv2.imencode(self.ext, frame)[1].tofile(self.video_out_name)
|
||||
self.preview_frame = cv2.hconcat([original_frame, frame])
|
||||
tbar.update(1)
|
||||
self.progress_total = 100
|
||||
else:
|
||||
if config.ACCURATE_MODE:
|
||||
# *********************** 批推理方案 start ***********************
|
||||
print('use accurate mode')
|
||||
self.video_inpaint = VideoInpaint(config.MAX_PROCESS_NUM)
|
||||
index = 0
|
||||
while True:
|
||||
ret, frame = self.video_cap.read()
|
||||
if not ret:
|
||||
break
|
||||
index += 1
|
||||
# 如果当前帧没有水印/文本则直接写
|
||||
if index not in sub_list.keys():
|
||||
self.video_writer.write(frame)
|
||||
print(f'write frame: {index}')
|
||||
self.update_progress(tbar, increment=1)
|
||||
continue
|
||||
# 如果有水印，判断该帧是不是开头帧
|
||||
else:
|
||||
# 如果是开头帧，则批推理到尾帧
|
||||
if self.is_current_frame_no_start(index, continuous_frame_no_list):
|
||||
# print(f'No 1 Current index: {index}')
|
||||
start_frame_no = index
|
||||
print(f'find start: {start_frame_no}')
|
||||
# 找到结束帧
|
||||
end_frame_no = self.find_frame_no_end(index, continuous_frame_no_list)
|
||||
# 判断当前帧号是不是字幕起始位置
|
||||
# 如果获取的结束帧号不为-1则说明
|
||||
if end_frame_no != -1:
|
||||
print(f'find end: {end_frame_no}')
|
||||
# ************ 读取该区间所有帧 start ************
|
||||
temp_frames = list()
|
||||
# 将头帧加入处理列表
|
||||
temp_frames.append(frame)
|
||||
inner_index = 0
|
||||
# 一直读取到尾帧
|
||||
while index < end_frame_no:
|
||||
ret, frame = self.video_cap.read()
|
||||
if not ret:
|
||||
break
|
||||
index += 1
|
||||
temp_frames.append(frame)
|
||||
# ************ 读取该区间所有帧 end ************
|
||||
if len(temp_frames) < 1:
|
||||
# 没有待处理，直接跳过
|
||||
continue
|
||||
elif len(temp_frames) == 1:
|
||||
inner_index += 1
|
||||
single_mask = create_mask(self.mask_size, sub_list[index])
|
||||
if self.lama_inpaint is None:
|
||||
self.lama_inpaint = LamaInpaint()
|
||||
inpainted_frame = self.lama_inpaint(frame, single_mask)
|
||||
self.video_writer.write(inpainted_frame)
|
||||
print(f'write frame: {start_frame_no + inner_index} with mask {sub_list[start_frame_no]}')
|
||||
self.update_progress(tbar, increment=1)
|
||||
continue
|
||||
else:
|
||||
# 将读取的视频帧分批处理
|
||||
# 1. 获取当前批次使用的mask
|
||||
mask = create_mask(self.mask_size, sub_list[start_frame_no])
|
||||
for batch in batch_generator(temp_frames, config.MAX_LOAD_NUM):
|
||||
# 2. 调用批推理
|
||||
if len(batch) == 1:
|
||||
single_mask = create_mask(self.mask_size, sub_list[start_frame_no])
|
||||
if self.lama_inpaint is None:
|
||||
self.lama_inpaint = LamaInpaint()
|
||||
inpainted_frame = self.lama_inpaint(frame, single_mask)
|
||||
self.video_writer.write(inpainted_frame)
|
||||
print(f'write frame: {start_frame_no + inner_index} with mask {sub_list[start_frame_no]}')
|
||||
inner_index += 1
|
||||
self.update_progress(tbar, increment=1)
|
||||
elif len(batch) > 1:
|
||||
inpainted_frames = self.video_inpaint.inpaint(batch, mask)
|
||||
for i, inpainted_frame in enumerate(inpainted_frames):
|
||||
self.video_writer.write(inpainted_frame)
|
||||
print(f'write frame: {start_frame_no + inner_index} with mask {sub_list[index]}')
|
||||
inner_index += 1
|
||||
self.preview_frame = cv2.hconcat([batch[i], inpainted_frame])
|
||||
self.update_progress(tbar, increment=len(batch))
|
||||
# *********************** 批推理方案 end ***********************
|
||||
else:
|
||||
# *********************** 单线程方案 start ***********************
|
||||
print('use normal mode')
|
||||
if self.lama_inpaint is None:
|
||||
self.lama_inpaint = LamaInpaint()
|
||||
index = 0
|
||||
while True:
|
||||
ret, frame = self.video_cap.read()
|
||||
if not ret:
|
||||
break
|
||||
original_frame = frame
|
||||
index += 1
|
||||
if index in sub_list.keys():
|
||||
mask = create_mask(self.mask_size, sub_list[index])
|
||||
if config.FAST_MODE:
|
||||
frame = cv2.inpaint(frame, mask, 3, cv2.INPAINT_TELEA)
|
||||
else:
|
||||
frame = self.lama_inpaint(frame, mask)
|
||||
self.preview_frame = cv2.hconcat([original_frame, frame])
|
||||
if self.is_picture:
|
||||
cv2.imencode(self.ext, frame)[1].tofile(self.video_out_name)
|
||||
else:
|
||||
self.video_writer.write(frame)
|
||||
tbar.update(1)
|
||||
self.progress_remover = 100 * float(index) / float(self.frame_count) // 2
|
||||
self.progress_total = 50 + self.progress_remover
|
||||
# *********************** 单线程方案 end ***********************
|
||||
self.video_cap.release()
|
||||
self.video_writer.release()
|
||||
if not self.is_picture:
|
||||
# 将原音频合并到新生成的视频文件中
|
||||
self.merge_audio_to_video()
|
||||
print(f"[Finished]Subtitle successfully removed, video generated at：{self.video_out_name}")
|
||||
else:
|
||||
print(f"[Finished]Subtitle successfully removed, picture generated at：{self.video_out_name}")
|
||||
print(f'time cost: {round(time.time() - start_time, 2)}s')
|
||||
self.isFinished = True
|
||||
self.progress_total = 100
|
||||
if os.path.exists(self.video_temp_file.name):
|
||||
try:
|
||||
os.remove(self.video_temp_file.name)
|
||||
except Exception:
|
||||
if platform.system() in ['Windows']:
|
||||
pass
|
||||
else:
|
||||
print(f'failed to delete temp file {self.video_temp_file.name}')
|
||||
|
||||
def merge_audio_to_video(self):
|
||||
# 创建音频临时对象，windows下delete=True会有permission denied的报错
|
||||
temp = tempfile.NamedTemporaryFile(suffix='.aac', delete=False)
|
||||
audio_extract_command = [config.FFMPEG_PATH,
|
||||
"-y", "-i", self.video_path,
|
||||
"-acodec", "copy",
|
||||
"-vn", "-loglevel", "error", temp.name]
|
||||
use_shell = True if os.name == "nt" else False
|
||||
try:
|
||||
subprocess.check_output(audio_extract_command, stdin=open(os.devnull), shell=use_shell)
|
||||
except Exception:
|
||||
print('fail to extract audio')
|
||||
return
|
||||
else:
|
||||
if os.path.exists(self.video_temp_file.name):
|
||||
audio_merge_command = [config.FFMPEG_PATH,
|
||||
"-y", "-i", self.video_temp_file.name,
|
||||
"-i", temp.name,
|
||||
"-vcodec", "copy",
|
||||
"-acodec", "copy",
|
||||
"-loglevel", "error", self.video_out_name]
|
||||
try:
|
||||
subprocess.check_output(audio_merge_command, stdin=open(os.devnull), shell=use_shell)
|
||||
except Exception:
|
||||
print('fail to merge audio')
|
||||
return
|
||||
if os.path.exists(temp.name):
|
||||
try:
|
||||
os.remove(temp.name)
|
||||
except Exception:
|
||||
print(f'failed to delete temp file {temp.name}')
|
||||
self.is_successful_merged = True
|
||||
finally:
|
||||
temp.close()
|
||||
if not self.is_successful_merged:
|
||||
try:
|
||||
shutil.copy2(self.video_temp_file.name, self.video_out_name)
|
||||
except IOError as e:
|
||||
print("Unable to copy file. %s" % e)
|
||||
self.video_temp_file.close()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
multiprocessing.set_start_method("spawn")
|
||||
# 提示用户输入视频路径
|
||||
video_path = input(f"Please input video file path: ").strip()
|
||||
# 新建字幕提取对象
|
||||
sd = SubtitleRemover(video_path)
|
||||
sd.run()
|
||||
@@ -0,0 +1,4 @@
|
||||
filename,filesize,encoding,header
|
||||
inference_1.pdiparams,50000000,,
|
||||
inference_2.pdiparams,50000000,,
|
||||
inference_3.pdiparams,13295054,,
|
||||
|
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,6 @@
|
||||
filename,filesize,encoding,header
|
||||
big-lama_1.pt,50000000,,
|
||||
big-lama_2.pt,50000000,,
|
||||
big-lama_3.pt,50000000,,
|
||||
big-lama_4.pt,50000000,,
|
||||
big-lama_5.pt,5803670,,
|
||||
|
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,5 @@
|
||||
filename,filesize,encoding,header
|
||||
ProPainter_1.pth,50000000,,
|
||||
ProPainter_2.pth,50000000,,
|
||||
ProPainter_3.pth,50000000,,
|
||||
ProPainter_4.pth,7780510,,
|
||||
|
Binary file not shown.
Binary file not shown.
Executable
+16
@@ -0,0 +1,16 @@
|
||||
# Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
import warnings
|
||||
warnings.filterwarnings("ignore", category=Warning)
|
||||
warnings.filterwarnings("ignore", category=DeprecationWarning)
|
||||
@@ -0,0 +1,109 @@
|
||||
# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
from __future__ import unicode_literals
|
||||
|
||||
import os
|
||||
import sys
|
||||
import numpy as np
|
||||
import skimage
|
||||
import paddle
|
||||
import signal
|
||||
import random
|
||||
|
||||
__dir__ = os.path.dirname(os.path.abspath(__file__))
|
||||
sys.path.append(os.path.abspath(os.path.join(__dir__, '../..')))
|
||||
|
||||
import copy
|
||||
from paddle.io import Dataset, DataLoader, BatchSampler, DistributedBatchSampler
|
||||
import paddle.distributed as dist
|
||||
|
||||
from ppocr.data.imaug import transform, create_operators
|
||||
from ppocr.data.simple_dataset import SimpleDataSet
|
||||
from ppocr.data.lmdb_dataset import LMDBDataSet
|
||||
from ppocr.data.pgnet_dataset import PGDataSet
|
||||
from ppocr.data.pubtab_dataset import PubTabDataSet
|
||||
|
||||
__all__ = ['build_dataloader', 'transform', 'create_operators']
|
||||
|
||||
|
||||
def term_mp(sig_num, frame):
|
||||
""" kill all child processes
|
||||
"""
|
||||
pid = os.getpid()
|
||||
pgid = os.getpgid(os.getpid())
|
||||
print("main proc {} exit, kill process group " "{}".format(pid, pgid))
|
||||
os.killpg(pgid, signal.SIGKILL)
|
||||
|
||||
|
||||
def build_dataloader(config, mode, device, logger, seed=None):
|
||||
config = copy.deepcopy(config)
|
||||
|
||||
support_dict = [
|
||||
'SimpleDataSet', 'LMDBDataSet', 'PGDataSet', 'PubTabDataSet'
|
||||
]
|
||||
module_name = config[mode]['dataset']['name']
|
||||
assert module_name in support_dict, Exception(
|
||||
'DataSet only support {}'.format(support_dict))
|
||||
assert mode in ['Train', 'Eval', 'Test'
|
||||
], "Mode should be Train, Eval or Test."
|
||||
|
||||
dataset = eval(module_name)(config, mode, logger, seed)
|
||||
loader_config = config[mode]['loader']
|
||||
batch_size = loader_config['batch_size_per_card']
|
||||
drop_last = loader_config['drop_last']
|
||||
shuffle = loader_config['shuffle']
|
||||
num_workers = loader_config['num_workers']
|
||||
if 'use_shared_memory' in loader_config.keys():
|
||||
use_shared_memory = loader_config['use_shared_memory']
|
||||
else:
|
||||
use_shared_memory = True
|
||||
|
||||
if mode == "Train":
|
||||
# Distribute data to multiple cards
|
||||
batch_sampler = DistributedBatchSampler(
|
||||
dataset=dataset,
|
||||
batch_size=batch_size,
|
||||
shuffle=shuffle,
|
||||
drop_last=drop_last)
|
||||
else:
|
||||
# Distribute data to single card
|
||||
batch_sampler = BatchSampler(
|
||||
dataset=dataset,
|
||||
batch_size=batch_size,
|
||||
shuffle=shuffle,
|
||||
drop_last=drop_last)
|
||||
|
||||
if 'collate_fn' in loader_config:
|
||||
from . import collate_fn
|
||||
collate_fn = getattr(collate_fn, loader_config['collate_fn'])()
|
||||
else:
|
||||
collate_fn = None
|
||||
data_loader = DataLoader(
|
||||
dataset=dataset,
|
||||
batch_sampler=batch_sampler,
|
||||
places=device,
|
||||
num_workers=num_workers,
|
||||
return_list=True,
|
||||
use_shared_memory=use_shared_memory,
|
||||
collate_fn=collate_fn)
|
||||
|
||||
# support exit using ctrl+c
|
||||
signal.signal(signal.SIGINT, term_mp)
|
||||
signal.signal(signal.SIGTERM, term_mp)
|
||||
|
||||
return data_loader
|
||||
@@ -0,0 +1,72 @@
|
||||
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import paddle
|
||||
import numbers
|
||||
import numpy as np
|
||||
from collections import defaultdict
|
||||
|
||||
|
||||
class DictCollator(object):
|
||||
"""
|
||||
data batch
|
||||
"""
|
||||
|
||||
def __call__(self, batch):
|
||||
# todo：support batch operators
|
||||
data_dict = defaultdict(list)
|
||||
to_tensor_keys = []
|
||||
for sample in batch:
|
||||
for k, v in sample.items():
|
||||
if isinstance(v, (np.ndarray, paddle.Tensor, numbers.Number)):
|
||||
if k not in to_tensor_keys:
|
||||
to_tensor_keys.append(k)
|
||||
data_dict[k].append(v)
|
||||
for k in to_tensor_keys:
|
||||
data_dict[k] = paddle.to_tensor(data_dict[k])
|
||||
return data_dict
|
||||
|
||||
|
||||
class ListCollator(object):
|
||||
"""
|
||||
data batch
|
||||
"""
|
||||
|
||||
def __call__(self, batch):
|
||||
# todo：support batch operators
|
||||
data_dict = defaultdict(list)
|
||||
to_tensor_idxs = []
|
||||
for sample in batch:
|
||||
for idx, v in enumerate(sample):
|
||||
if isinstance(v, (np.ndarray, paddle.Tensor, numbers.Number)):
|
||||
if idx not in to_tensor_idxs:
|
||||
to_tensor_idxs.append(idx)
|
||||
data_dict[idx].append(v)
|
||||
for idx in to_tensor_idxs:
|
||||
data_dict[idx] = paddle.to_tensor(data_dict[idx])
|
||||
return list(data_dict.values())
|
||||
|
||||
|
||||
class SSLRotateCollate(object):
|
||||
"""
|
||||
bach: [
|
||||
[(4*3xH*W), (4,)]
|
||||
[(4*3xH*W), (4,)]
|
||||
...
|
||||
]
|
||||
"""
|
||||
|
||||
def __call__(self, batch):
|
||||
output = [np.concatenate(d, axis=0) for d in zip(*batch)]
|
||||
return output
|
||||
@@ -0,0 +1,26 @@
|
||||
# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from paddle.vision.transforms import ColorJitter as pp_ColorJitter
|
||||
|
||||
__all__ = ['ColorJitter']
|
||||
|
||||
class ColorJitter(object):
|
||||
def __init__(self, brightness=0, contrast=0, saturation=0, hue=0,**kwargs):
|
||||
self.aug = pp_ColorJitter(brightness, contrast, saturation, hue)
|
||||
|
||||
def __call__(self, data):
|
||||
image = data['image']
|
||||
image = self.aug(image)
|
||||
data['image'] = image
|
||||
return data
|
||||
@@ -0,0 +1,74 @@
|
||||
# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
from __future__ import unicode_literals
|
||||
|
||||
from .iaa_augment import IaaAugment
|
||||
from .make_border_map import MakeBorderMap
|
||||
from .make_shrink_map import MakeShrinkMap
|
||||
from .random_crop_data import EastRandomCropData, RandomCropImgMask
|
||||
from .make_pse_gt import MakePseGt
|
||||
|
||||
from .rec_img_aug import RecAug, RecConAug, RecResizeImg, ClsResizeImg, \
|
||||
SRNRecResizeImg, NRTRRecResizeImg, SARRecResizeImg, PRENResizeImg
|
||||
from .ssl_img_aug import SSLRotateResize
|
||||
from .randaugment import RandAugment
|
||||
from .copy_paste import CopyPaste
|
||||
from .ColorJitter import ColorJitter
|
||||
from .operators import *
|
||||
from .label_ops import *
|
||||
|
||||
from .east_process import *
|
||||
from .sast_process import *
|
||||
from .pg_process import *
|
||||
from .gen_table_mask import *
|
||||
|
||||
from .vqa import *
|
||||
|
||||
from .fce_aug import *
|
||||
from .fce_targets import FCENetTargets
|
||||
|
||||
|
||||
def transform(data, ops=None):
|
||||
""" transform """
|
||||
if ops is None:
|
||||
ops = []
|
||||
for op in ops:
|
||||
data = op(data)
|
||||
if data is None:
|
||||
return None
|
||||
return data
|
||||
|
||||
|
||||
def create_operators(op_param_list, global_config=None):
|
||||
"""
|
||||
create operators based on the config
|
||||
|
||||
Args:
|
||||
params(list): a dict list, used to create some operators
|
||||
"""
|
||||
assert isinstance(op_param_list, list), ('operator config should be a list')
|
||||
ops = []
|
||||
for operator in op_param_list:
|
||||
assert isinstance(operator,
|
||||
dict) and len(operator) == 1, "yaml format error"
|
||||
op_name = list(operator)[0]
|
||||
param = {} if operator[op_name] is None else operator[op_name]
|
||||
if global_config is not None:
|
||||
param.update(global_config)
|
||||
op = eval(op_name)(**param)
|
||||
ops.append(op)
|
||||
return ops
|
||||
@@ -0,0 +1,170 @@
|
||||
# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
import copy
|
||||
import cv2
|
||||
import random
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from shapely.geometry import Polygon
|
||||
|
||||
from ppocr.data.imaug.iaa_augment import IaaAugment
|
||||
from ppocr.data.imaug.random_crop_data import is_poly_outside_rect
|
||||
from tools.infer.utility import get_rotate_crop_image
|
||||
|
||||
|
||||
class CopyPaste(object):
|
||||
def __init__(self, objects_paste_ratio=0.2, limit_paste=True, **kwargs):
|
||||
self.ext_data_num = 1
|
||||
self.objects_paste_ratio = objects_paste_ratio
|
||||
self.limit_paste = limit_paste
|
||||
augmenter_args = [{'type': 'Resize', 'args': {'size': [0.5, 3]}}]
|
||||
self.aug = IaaAugment(augmenter_args)
|
||||
|
||||
def __call__(self, data):
|
||||
point_num = data['polys'].shape[1]
|
||||
src_img = data['image']
|
||||
src_polys = data['polys'].tolist()
|
||||
src_ignores = data['ignore_tags'].tolist()
|
||||
ext_data = data['ext_data'][0]
|
||||
ext_image = ext_data['image']
|
||||
ext_polys = ext_data['polys']
|
||||
ext_ignores = ext_data['ignore_tags']
|
||||
|
||||
indexs = [i for i in range(len(ext_ignores)) if not ext_ignores[i]]
|
||||
select_num = max(
|
||||
1, min(int(self.objects_paste_ratio * len(ext_polys)), 30))
|
||||
|
||||
random.shuffle(indexs)
|
||||
select_idxs = indexs[:select_num]
|
||||
select_polys = ext_polys[select_idxs]
|
||||
select_ignores = ext_ignores[select_idxs]
|
||||
|
||||
src_img = cv2.cvtColor(src_img, cv2.COLOR_BGR2RGB)
|
||||
ext_image = cv2.cvtColor(ext_image, cv2.COLOR_BGR2RGB)
|
||||
src_img = Image.fromarray(src_img).convert('RGBA')
|
||||
for poly, tag in zip(select_polys, select_ignores):
|
||||
box_img = get_rotate_crop_image(ext_image, poly)
|
||||
|
||||
src_img, box = self.paste_img(src_img, box_img, src_polys)
|
||||
if box is not None:
|
||||
box = box.tolist()
|
||||
for _ in range(len(box), point_num):
|
||||
box.append(box[-1])
|
||||
src_polys.append(box)
|
||||
src_ignores.append(tag)
|
||||
src_img = cv2.cvtColor(np.array(src_img), cv2.COLOR_RGB2BGR)
|
||||
h, w = src_img.shape[:2]
|
||||
src_polys = np.array(src_polys)
|
||||
src_polys[:, :, 0] = np.clip(src_polys[:, :, 0], 0, w)
|
||||
src_polys[:, :, 1] = np.clip(src_polys[:, :, 1], 0, h)
|
||||
data['image'] = src_img
|
||||
data['polys'] = src_polys
|
||||
data['ignore_tags'] = np.array(src_ignores)
|
||||
return data
|
||||
|
||||
def paste_img(self, src_img, box_img, src_polys):
|
||||
box_img_pil = Image.fromarray(box_img).convert('RGBA')
|
||||
src_w, src_h = src_img.size
|
||||
box_w, box_h = box_img_pil.size
|
||||
|
||||
angle = np.random.randint(0, 360)
|
||||
box = np.array([[[0, 0], [box_w, 0], [box_w, box_h], [0, box_h]]])
|
||||
box = rotate_bbox(box_img, box, angle)[0]
|
||||
box_img_pil = box_img_pil.rotate(angle, expand=1)
|
||||
box_w, box_h = box_img_pil.width, box_img_pil.height
|
||||
if src_w - box_w < 0 or src_h - box_h < 0:
|
||||
return src_img, None
|
||||
|
||||
paste_x, paste_y = self.select_coord(src_polys, box, src_w - box_w,
|
||||
src_h - box_h)
|
||||
if paste_x is None:
|
||||
return src_img, None
|
||||
box[:, 0] += paste_x
|
||||
box[:, 1] += paste_y
|
||||
r, g, b, A = box_img_pil.split()
|
||||
src_img.paste(box_img_pil, (paste_x, paste_y), mask=A)
|
||||
|
||||
return src_img, box
|
||||
|
||||
def select_coord(self, src_polys, box, endx, endy):
|
||||
if self.limit_paste:
|
||||
xmin, ymin, xmax, ymax = box[:, 0].min(), box[:, 1].min(
|
||||
), box[:, 0].max(), box[:, 1].max()
|
||||
for _ in range(50):
|
||||
paste_x = random.randint(0, endx)
|
||||
paste_y = random.randint(0, endy)
|
||||
xmin1 = xmin + paste_x
|
||||
xmax1 = xmax + paste_x
|
||||
ymin1 = ymin + paste_y
|
||||
ymax1 = ymax + paste_y
|
||||
|
||||
num_poly_in_rect = 0
|
||||
for poly in src_polys:
|
||||
if not is_poly_outside_rect(poly, xmin1, ymin1,
|
||||
xmax1 - xmin1, ymax1 - ymin1):
|
||||
num_poly_in_rect += 1
|
||||
break
|
||||
if num_poly_in_rect == 0:
|
||||
return paste_x, paste_y
|
||||
return None, None
|
||||
else:
|
||||
paste_x = random.randint(0, endx)
|
||||
paste_y = random.randint(0, endy)
|
||||
return paste_x, paste_y
|
||||
|
||||
|
||||
def get_union(pD, pG):
|
||||
return Polygon(pD).union(Polygon(pG)).area
|
||||
|
||||
|
||||
def get_intersection_over_union(pD, pG):
|
||||
return get_intersection(pD, pG) / get_union(pD, pG)
|
||||
|
||||
|
||||
def get_intersection(pD, pG):
|
||||
return Polygon(pD).intersection(Polygon(pG)).area
|
||||
|
||||
|
||||
def rotate_bbox(img, text_polys, angle, scale=1):
|
||||
"""
|
||||
from https://github.com/WenmuZhou/DBNet.pytorch/blob/master/data_loader/modules/augment.py
|
||||
Args:
|
||||
img: np.ndarray
|
||||
text_polys: np.ndarray N*4*2
|
||||
angle: int
|
||||
scale: int
|
||||
|
||||
Returns:
|
||||
|
||||
"""
|
||||
w = img.shape[1]
|
||||
h = img.shape[0]
|
||||
|
||||
rangle = np.deg2rad(angle)
|
||||
nw = (abs(np.sin(rangle) * h) + abs(np.cos(rangle) * w))
|
||||
nh = (abs(np.cos(rangle) * h) + abs(np.sin(rangle) * w))
|
||||
rot_mat = cv2.getRotationMatrix2D((nw * 0.5, nh * 0.5), angle, scale)
|
||||
rot_move = np.dot(rot_mat, np.array([(nw - w) * 0.5, (nh - h) * 0.5, 0]))
|
||||
rot_mat[0, 2] += rot_move[0]
|
||||
rot_mat[1, 2] += rot_move[1]
|
||||
|
||||
# ---------------------- rotate box ----------------------
|
||||
rot_text_polys = list()
|
||||
for bbox in text_polys:
|
||||
point1 = np.dot(rot_mat, np.array([bbox[0, 0], bbox[0, 1], 1]))
|
||||
point2 = np.dot(rot_mat, np.array([bbox[1, 0], bbox[1, 1], 1]))
|
||||
point3 = np.dot(rot_mat, np.array([bbox[2, 0], bbox[2, 1], 1]))
|
||||
point4 = np.dot(rot_mat, np.array([bbox[3, 0], bbox[3, 1], 1]))
|
||||
rot_text_polys.append([point1, point2, point3, point4])
|
||||
return np.array(rot_text_polys, dtype=np.float32)
|
||||
@@ -0,0 +1,436 @@
|
||||
#copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
#Licensed under the Apache License, Version 2.0 (the "License");
|
||||
#you may not use this file except in compliance with the License.
|
||||
#You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
#Unless required by applicable law or agreed to in writing, software
|
||||
#distributed under the License is distributed on an "AS IS" BASIS,
|
||||
#WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
#See the License for the specific language governing permissions and
|
||||
#limitations under the License.
|
||||
"""
|
||||
This code is refered from:
|
||||
https://github.com/songdejia/EAST/blob/master/data_utils.py
|
||||
"""
|
||||
import math
|
||||
import cv2
|
||||
import numpy as np
|
||||
import json
|
||||
import sys
|
||||
import os
|
||||
|
||||
__all__ = ['EASTProcessTrain']
|
||||
|
||||
|
||||
class EASTProcessTrain(object):
|
||||
def __init__(self,
|
||||
image_shape=[512, 512],
|
||||
background_ratio=0.125,
|
||||
min_crop_side_ratio=0.1,
|
||||
min_text_size=10,
|
||||
**kwargs):
|
||||
self.input_size = image_shape[1]
|
||||
self.random_scale = np.array([0.5, 1, 2.0, 3.0])
|
||||
self.background_ratio = background_ratio
|
||||
self.min_crop_side_ratio = min_crop_side_ratio
|
||||
self.min_text_size = min_text_size
|
||||
|
||||
def preprocess(self, im):
|
||||
input_size = self.input_size
|
||||
im_shape = im.shape
|
||||
im_size_min = np.min(im_shape[0:2])
|
||||
im_size_max = np.max(im_shape[0:2])
|
||||
im_scale = float(input_size) / float(im_size_max)
|
||||
im = cv2.resize(im, None, None, fx=im_scale, fy=im_scale)
|
||||
img_mean = [0.485, 0.456, 0.406]
|
||||
img_std = [0.229, 0.224, 0.225]
|
||||
# im = im[:, :, ::-1].astype(np.float32)
|
||||
im = im / 255
|
||||
im -= img_mean
|
||||
im /= img_std
|
||||
new_h, new_w, _ = im.shape
|
||||
im_padded = np.zeros((input_size, input_size, 3), dtype=np.float32)
|
||||
im_padded[:new_h, :new_w, :] = im
|
||||
im_padded = im_padded.transpose((2, 0, 1))
|
||||
im_padded = im_padded[np.newaxis, :]
|
||||
return im_padded, im_scale
|
||||
|
||||
def rotate_im_poly(self, im, text_polys):
|
||||
"""
|
||||
rotate image with 90 / 180 / 270 degre
|
||||
"""
|
||||
im_w, im_h = im.shape[1], im.shape[0]
|
||||
dst_im = im.copy()
|
||||
dst_polys = []
|
||||
rand_degree_ratio = np.random.rand()
|
||||
rand_degree_cnt = 1
|
||||
if 0.333 < rand_degree_ratio < 0.666:
|
||||
rand_degree_cnt = 2
|
||||
elif rand_degree_ratio > 0.666:
|
||||
rand_degree_cnt = 3
|
||||
for i in range(rand_degree_cnt):
|
||||
dst_im = np.rot90(dst_im)
|
||||
rot_degree = -90 * rand_degree_cnt
|
||||
rot_angle = rot_degree * math.pi / 180.0
|
||||
n_poly = text_polys.shape[0]
|
||||
cx, cy = 0.5 * im_w, 0.5 * im_h
|
||||
ncx, ncy = 0.5 * dst_im.shape[1], 0.5 * dst_im.shape[0]
|
||||
for i in range(n_poly):
|
||||
wordBB = text_polys[i]
|
||||
poly = []
|
||||
for j in range(4):
|
||||
sx, sy = wordBB[j][0], wordBB[j][1]
|
||||
dx = math.cos(rot_angle) * (sx - cx)\
|
||||
- math.sin(rot_angle) * (sy - cy) + ncx
|
||||
dy = math.sin(rot_angle) * (sx - cx)\
|
||||
+ math.cos(rot_angle) * (sy - cy) + ncy
|
||||
poly.append([dx, dy])
|
||||
dst_polys.append(poly)
|
||||
dst_polys = np.array(dst_polys, dtype=np.float32)
|
||||
return dst_im, dst_polys
|
||||
|
||||
def polygon_area(self, poly):
|
||||
"""
|
||||
compute area of a polygon
|
||||
:param poly:
|
||||
:return:
|
||||
"""
|
||||
edge = [(poly[1][0] - poly[0][0]) * (poly[1][1] + poly[0][1]),
|
||||
(poly[2][0] - poly[1][0]) * (poly[2][1] + poly[1][1]),
|
||||
(poly[3][0] - poly[2][0]) * (poly[3][1] + poly[2][1]),
|
||||
(poly[0][0] - poly[3][0]) * (poly[0][1] + poly[3][1])]
|
||||
return np.sum(edge) / 2.
|
||||
|
||||
def check_and_validate_polys(self, polys, tags, img_height, img_width):
|
||||
"""
|
||||
check so that the text poly is in the same direction,
|
||||
and also filter some invalid polygons
|
||||
:param polys:
|
||||
:param tags:
|
||||
:return:
|
||||
"""
|
||||
h, w = img_height, img_width
|
||||
if polys.shape[0] == 0:
|
||||
return polys
|
||||
polys[:, :, 0] = np.clip(polys[:, :, 0], 0, w - 1)
|
||||
polys[:, :, 1] = np.clip(polys[:, :, 1], 0, h - 1)
|
||||
|
||||
validated_polys = []
|
||||
validated_tags = []
|
||||
for poly, tag in zip(polys, tags):
|
||||
p_area = self.polygon_area(poly)
|
||||
#invalid poly
|
||||
if abs(p_area) < 1:
|
||||
continue
|
||||
if p_area > 0:
|
||||
#'poly in wrong direction'
|
||||
if not tag:
|
||||
tag = True #reversed cases should be ignore
|
||||
poly = poly[(0, 3, 2, 1), :]
|
||||
validated_polys.append(poly)
|
||||
validated_tags.append(tag)
|
||||
return np.array(validated_polys), np.array(validated_tags)
|
||||
|
||||
def draw_img_polys(self, img, polys):
|
||||
if len(img.shape) == 4:
|
||||
img = np.squeeze(img, axis=0)
|
||||
if img.shape[0] == 3:
|
||||
img = img.transpose((1, 2, 0))
|
||||
img[:, :, 2] += 123.68
|
||||
img[:, :, 1] += 116.78
|
||||
img[:, :, 0] += 103.94
|
||||
cv2.imwrite("tmp.jpg", img)
|
||||
img = cv2.imread("tmp.jpg")
|
||||
for box in polys:
|
||||
box = box.astype(np.int32).reshape((-1, 1, 2))
|
||||
cv2.polylines(img, [box], True, color=(255, 255, 0), thickness=2)
|
||||
import random
|
||||
ino = random.randint(0, 100)
|
||||
cv2.imwrite("tmp_%d.jpg" % ino, img)
|
||||
return
|
||||
|
||||
def shrink_poly(self, poly, r):
|
||||
"""
|
||||
fit a poly inside the origin poly, maybe bugs here...
|
||||
used for generate the score map
|
||||
:param poly: the text poly
|
||||
:param r: r in the paper
|
||||
:return: the shrinked poly
|
||||
"""
|
||||
# shrink ratio
|
||||
R = 0.3
|
||||
# find the longer pair
|
||||
dist0 = np.linalg.norm(poly[0] - poly[1])
|
||||
dist1 = np.linalg.norm(poly[2] - poly[3])
|
||||
dist2 = np.linalg.norm(poly[0] - poly[3])
|
||||
dist3 = np.linalg.norm(poly[1] - poly[2])
|
||||
if dist0 + dist1 > dist2 + dist3:
|
||||
# first move (p0, p1), (p2, p3), then (p0, p3), (p1, p2)
|
||||
## p0, p1
|
||||
theta = np.arctan2((poly[1][1] - poly[0][1]),
|
||||
(poly[1][0] - poly[0][0]))
|
||||
poly[0][0] += R * r[0] * np.cos(theta)
|
||||
poly[0][1] += R * r[0] * np.sin(theta)
|
||||
poly[1][0] -= R * r[1] * np.cos(theta)
|
||||
poly[1][1] -= R * r[1] * np.sin(theta)
|
||||
## p2, p3
|
||||
theta = np.arctan2((poly[2][1] - poly[3][1]),
|
||||
(poly[2][0] - poly[3][0]))
|
||||
poly[3][0] += R * r[3] * np.cos(theta)
|
||||
poly[3][1] += R * r[3] * np.sin(theta)
|
||||
poly[2][0] -= R * r[2] * np.cos(theta)
|
||||
poly[2][1] -= R * r[2] * np.sin(theta)
|
||||
## p0, p3
|
||||
theta = np.arctan2((poly[3][0] - poly[0][0]),
|
||||
(poly[3][1] - poly[0][1]))
|
||||
poly[0][0] += R * r[0] * np.sin(theta)
|
||||
poly[0][1] += R * r[0] * np.cos(theta)
|
||||
poly[3][0] -= R * r[3] * np.sin(theta)
|
||||
poly[3][1] -= R * r[3] * np.cos(theta)
|
||||
## p1, p2
|
||||
theta = np.arctan2((poly[2][0] - poly[1][0]),
|
||||
(poly[2][1] - poly[1][1]))
|
||||
poly[1][0] += R * r[1] * np.sin(theta)
|
||||
poly[1][1] += R * r[1] * np.cos(theta)
|
||||
poly[2][0] -= R * r[2] * np.sin(theta)
|
||||
poly[2][1] -= R * r[2] * np.cos(theta)
|
||||
else:
|
||||
## p0, p3
|
||||
# print poly
|
||||
theta = np.arctan2((poly[3][0] - poly[0][0]),
|
||||
(poly[3][1] - poly[0][1]))
|
||||
poly[0][0] += R * r[0] * np.sin(theta)
|
||||
poly[0][1] += R * r[0] * np.cos(theta)
|
||||
poly[3][0] -= R * r[3] * np.sin(theta)
|
||||
poly[3][1] -= R * r[3] * np.cos(theta)
|
||||
## p1, p2
|
||||
theta = np.arctan2((poly[2][0] - poly[1][0]),
|
||||
(poly[2][1] - poly[1][1]))
|
||||
poly[1][0] += R * r[1] * np.sin(theta)
|
||||
poly[1][1] += R * r[1] * np.cos(theta)
|
||||
poly[2][0] -= R * r[2] * np.sin(theta)
|
||||
poly[2][1] -= R * r[2] * np.cos(theta)
|
||||
## p0, p1
|
||||
theta = np.arctan2((poly[1][1] - poly[0][1]),
|
||||
(poly[1][0] - poly[0][0]))
|
||||
poly[0][0] += R * r[0] * np.cos(theta)
|
||||
poly[0][1] += R * r[0] * np.sin(theta)
|
||||
poly[1][0] -= R * r[1] * np.cos(theta)
|
||||
poly[1][1] -= R * r[1] * np.sin(theta)
|
||||
## p2, p3
|
||||
theta = np.arctan2((poly[2][1] - poly[3][1]),
|
||||
(poly[2][0] - poly[3][0]))
|
||||
poly[3][0] += R * r[3] * np.cos(theta)
|
||||
poly[3][1] += R * r[3] * np.sin(theta)
|
||||
poly[2][0] -= R * r[2] * np.cos(theta)
|
||||
poly[2][1] -= R * r[2] * np.sin(theta)
|
||||
return poly
|
||||
|
||||
def generate_quad(self, im_size, polys, tags):
|
||||
"""
|
||||
Generate quadrangle.
|
||||
"""
|
||||
h, w = im_size
|
||||
poly_mask = np.zeros((h, w), dtype=np.uint8)
|
||||
score_map = np.zeros((h, w), dtype=np.uint8)
|
||||
# (x1, y1, ..., x4, y4, short_edge_norm)
|
||||
geo_map = np.zeros((h, w, 9), dtype=np.float32)
|
||||
# mask used during traning, to ignore some hard areas
|
||||
training_mask = np.ones((h, w), dtype=np.uint8)
|
||||
for poly_idx, poly_tag in enumerate(zip(polys, tags)):
|
||||
poly = poly_tag[0]
|
||||
tag = poly_tag[1]
|
||||
|
||||
r = [None, None, None, None]
|
||||
for i in range(4):
|
||||
dist1 = np.linalg.norm(poly[i] - poly[(i + 1) % 4])
|
||||
dist2 = np.linalg.norm(poly[i] - poly[(i - 1) % 4])
|
||||
r[i] = min(dist1, dist2)
|
||||
# score map
|
||||
shrinked_poly = self.shrink_poly(
|
||||
poly.copy(), r).astype(np.int32)[np.newaxis, :, :]
|
||||
cv2.fillPoly(score_map, shrinked_poly, 1)
|
||||
cv2.fillPoly(poly_mask, shrinked_poly, poly_idx + 1)
|
||||
# if the poly is too small, then ignore it during training
|
||||
poly_h = min(
|
||||
np.linalg.norm(poly[0] - poly[3]),
|
||||
np.linalg.norm(poly[1] - poly[2]))
|
||||
poly_w = min(
|
||||
np.linalg.norm(poly[0] - poly[1]),
|
||||
np.linalg.norm(poly[2] - poly[3]))
|
||||
if min(poly_h, poly_w) < self.min_text_size:
|
||||
cv2.fillPoly(training_mask,
|
||||
poly.astype(np.int32)[np.newaxis, :, :], 0)
|
||||
|
||||
if tag:
|
||||
cv2.fillPoly(training_mask,
|
||||
poly.astype(np.int32)[np.newaxis, :, :], 0)
|
||||
|
||||
xy_in_poly = np.argwhere(poly_mask == (poly_idx + 1))
|
||||
# geo map.
|
||||
y_in_poly = xy_in_poly[:, 0]
|
||||
x_in_poly = xy_in_poly[:, 1]
|
||||
poly[:, 0] = np.minimum(np.maximum(poly[:, 0], 0), w)
|
||||
poly[:, 1] = np.minimum(np.maximum(poly[:, 1], 0), h)
|
||||
for pno in range(4):
|
||||
geo_channel_beg = pno * 2
|
||||
geo_map[y_in_poly, x_in_poly, geo_channel_beg] =\
|
||||
x_in_poly - poly[pno, 0]
|
||||
geo_map[y_in_poly, x_in_poly, geo_channel_beg+1] =\
|
||||
y_in_poly - poly[pno, 1]
|
||||
geo_map[y_in_poly, x_in_poly, 8] = \
|
||||
1.0 / max(min(poly_h, poly_w), 1.0)
|
||||
return score_map, geo_map, training_mask
|
||||
|
||||
def crop_area(self, im, polys, tags, crop_background=False, max_tries=50):
|
||||
"""
|
||||
make random crop from the input image
|
||||
:param im:
|
||||
:param polys:
|
||||
:param tags:
|
||||
:param crop_background:
|
||||
:param max_tries:
|
||||
:return:
|
||||
"""
|
||||
h, w, _ = im.shape
|
||||
pad_h = h // 10
|
||||
pad_w = w // 10
|
||||
h_array = np.zeros((h + pad_h * 2), dtype=np.int32)
|
||||
w_array = np.zeros((w + pad_w * 2), dtype=np.int32)
|
||||
for poly in polys:
|
||||
poly = np.round(poly, decimals=0).astype(np.int32)
|
||||
minx = np.min(poly[:, 0])
|
||||
maxx = np.max(poly[:, 0])
|
||||
w_array[minx + pad_w:maxx + pad_w] = 1
|
||||
miny = np.min(poly[:, 1])
|
||||
maxy = np.max(poly[:, 1])
|
||||
h_array[miny + pad_h:maxy + pad_h] = 1
|
||||
# ensure the cropped area not across a text
|
||||
h_axis = np.where(h_array == 0)[0]
|
||||
w_axis = np.where(w_array == 0)[0]
|
||||
if len(h_axis) == 0 or len(w_axis) == 0:
|
||||
return im, polys, tags
|
||||
|
||||
for i in range(max_tries):
|
||||
xx = np.random.choice(w_axis, size=2)
|
||||
xmin = np.min(xx) - pad_w
|
||||
xmax = np.max(xx) - pad_w
|
||||
xmin = np.clip(xmin, 0, w - 1)
|
||||
xmax = np.clip(xmax, 0, w - 1)
|
||||
yy = np.random.choice(h_axis, size=2)
|
||||
ymin = np.min(yy) - pad_h
|
||||
ymax = np.max(yy) - pad_h
|
||||
ymin = np.clip(ymin, 0, h - 1)
|
||||
ymax = np.clip(ymax, 0, h - 1)
|
||||
if xmax - xmin < self.min_crop_side_ratio * w or \
|
||||
ymax - ymin < self.min_crop_side_ratio * h:
|
||||
# area too small
|
||||
continue
|
||||
if polys.shape[0] != 0:
|
||||
poly_axis_in_area = (polys[:, :, 0] >= xmin)\
|
||||
& (polys[:, :, 0] <= xmax)\
|
||||
& (polys[:, :, 1] >= ymin)\
|
||||
& (polys[:, :, 1] <= ymax)
|
||||
selected_polys = np.where(
|
||||
np.sum(poly_axis_in_area, axis=1) == 4)[0]
|
||||
else:
|
||||
selected_polys = []
|
||||
|
||||
if len(selected_polys) == 0:
|
||||
# no text in this area
|
||||
if crop_background:
|
||||
im = im[ymin:ymax + 1, xmin:xmax + 1, :]
|
||||
polys = []
|
||||
tags = []
|
||||
return im, polys, tags
|
||||
else:
|
||||
continue
|
||||
|
||||
im = im[ymin:ymax + 1, xmin:xmax + 1, :]
|
||||
polys = polys[selected_polys]
|
||||
tags = tags[selected_polys]
|
||||
polys[:, :, 0] -= xmin
|
||||
polys[:, :, 1] -= ymin
|
||||
return im, polys, tags
|
||||
return im, polys, tags
|
||||
|
||||
def crop_background_infor(self, im, text_polys, text_tags):
|
||||
im, text_polys, text_tags = self.crop_area(
|
||||
im, text_polys, text_tags, crop_background=True)
|
||||
|
||||
if len(text_polys) > 0:
|
||||
return None
|
||||
# pad and resize image
|
||||
input_size = self.input_size
|
||||
im, ratio = self.preprocess(im)
|
||||
score_map = np.zeros((input_size, input_size), dtype=np.float32)
|
||||
geo_map = np.zeros((input_size, input_size, 9), dtype=np.float32)
|
||||
training_mask = np.ones((input_size, input_size), dtype=np.float32)
|
||||
return im, score_map, geo_map, training_mask
|
||||
|
||||
def crop_foreground_infor(self, im, text_polys, text_tags):
|
||||
im, text_polys, text_tags = self.crop_area(
|
||||
im, text_polys, text_tags, crop_background=False)
|
||||
|
||||
if text_polys.shape[0] == 0:
|
||||
return None
|
||||
#continue for all ignore case
|
||||
if np.sum((text_tags * 1.0)) >= text_tags.size:
|
||||
return None
|
||||
# pad and resize image
|
||||
input_size = self.input_size
|
||||
im, ratio = self.preprocess(im)
|
||||
text_polys[:, :, 0] *= ratio
|
||||
text_polys[:, :, 1] *= ratio
|
||||
_, _, new_h, new_w = im.shape
|
||||
# print(im.shape)
|
||||
# self.draw_img_polys(im, text_polys)
|
||||
score_map, geo_map, training_mask = self.generate_quad(
|
||||
(new_h, new_w), text_polys, text_tags)
|
||||
return im, score_map, geo_map, training_mask
|
||||
|
||||
def __call__(self, data):
|
||||
im = data['image']
|
||||
text_polys = data['polys']
|
||||
text_tags = data['ignore_tags']
|
||||
if im is None:
|
||||
return None
|
||||
if text_polys.shape[0] == 0:
|
||||
return None
|
||||
|
||||
#add rotate cases
|
||||
if np.random.rand() < 0.5:
|
||||
im, text_polys = self.rotate_im_poly(im, text_polys)
|
||||
h, w, _ = im.shape
|
||||
text_polys, text_tags = self.check_and_validate_polys(text_polys,
|
||||
text_tags, h, w)
|
||||
if text_polys.shape[0] == 0:
|
||||
return None
|
||||
|
||||
# random scale this image
|
||||
rd_scale = np.random.choice(self.random_scale)
|
||||
im = cv2.resize(im, dsize=None, fx=rd_scale, fy=rd_scale)
|
||||
text_polys *= rd_scale
|
||||
if np.random.rand() < self.background_ratio:
|
||||
outs = self.crop_background_infor(im, text_polys, text_tags)
|
||||
else:
|
||||
outs = self.crop_foreground_infor(im, text_polys, text_tags)
|
||||
|
||||
if outs is None:
|
||||
return None
|
||||
im, score_map, geo_map, training_mask = outs
|
||||
score_map = score_map[np.newaxis, ::4, ::4].astype(np.float32)
|
||||
geo_map = np.swapaxes(geo_map, 1, 2)
|
||||
geo_map = np.swapaxes(geo_map, 1, 0)
|
||||
geo_map = geo_map[:, ::4, ::4].astype(np.float32)
|
||||
training_mask = training_mask[np.newaxis, ::4, ::4]
|
||||
training_mask = training_mask.astype(np.float32)
|
||||
|
||||
data['image'] = im[0]
|
||||
data['score_map'] = score_map
|
||||
data['geo_map'] = geo_map
|
||||
data['training_mask'] = training_mask
|
||||
return data
|
||||
@@ -0,0 +1,564 @@
|
||||
# copyright (c) 2022 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""
|
||||
This code is refer from:
|
||||
https://github.com/open-mmlab/mmocr/blob/main/mmocr/datasets/pipelines/transforms.py
|
||||
"""
|
||||
import numpy as np
|
||||
from PIL import Image, ImageDraw
|
||||
import cv2
|
||||
from shapely.geometry import Polygon
|
||||
import math
|
||||
from ppocr.utils.poly_nms import poly_intersection
|
||||
|
||||
|
||||
class RandomScaling:
|
||||
def __init__(self, size=800, scale=(3. / 4, 5. / 2), **kwargs):
|
||||
"""Random scale the image while keeping aspect.
|
||||
|
||||
Args:
|
||||
size (int) : Base size before scaling.
|
||||
scale (tuple(float)) : The range of scaling.
|
||||
"""
|
||||
assert isinstance(size, int)
|
||||
assert isinstance(scale, float) or isinstance(scale, tuple)
|
||||
self.size = size
|
||||
self.scale = scale if isinstance(scale, tuple) \
|
||||
else (1 - scale, 1 + scale)
|
||||
|
||||
def __call__(self, data):
|
||||
image = data['image']
|
||||
text_polys = data['polys']
|
||||
h, w, _ = image.shape
|
||||
|
||||
aspect_ratio = np.random.uniform(min(self.scale), max(self.scale))
|
||||
scales = self.size * 1.0 / max(h, w) * aspect_ratio
|
||||
scales = np.array([scales, scales])
|
||||
out_size = (int(h * scales[1]), int(w * scales[0]))
|
||||
image = cv2.resize(image, out_size[::-1])
|
||||
|
||||
data['image'] = image
|
||||
text_polys[:, :, 0::2] = text_polys[:, :, 0::2] * scales[1]
|
||||
text_polys[:, :, 1::2] = text_polys[:, :, 1::2] * scales[0]
|
||||
data['polys'] = text_polys
|
||||
|
||||
return data
|
||||
|
||||
|
||||
class RandomCropFlip:
|
||||
def __init__(self,
|
||||
pad_ratio=0.1,
|
||||
crop_ratio=0.5,
|
||||
iter_num=1,
|
||||
min_area_ratio=0.2,
|
||||
**kwargs):
|
||||
"""Random crop and flip a patch of the image.
|
||||
|
||||
Args:
|
||||
crop_ratio (float): The ratio of cropping.
|
||||
iter_num (int): Number of operations.
|
||||
min_area_ratio (float): Minimal area ratio between cropped patch
|
||||
and original image.
|
||||
"""
|
||||
assert isinstance(crop_ratio, float)
|
||||
assert isinstance(iter_num, int)
|
||||
assert isinstance(min_area_ratio, float)
|
||||
|
||||
self.pad_ratio = pad_ratio
|
||||
self.epsilon = 1e-2
|
||||
self.crop_ratio = crop_ratio
|
||||
self.iter_num = iter_num
|
||||
self.min_area_ratio = min_area_ratio
|
||||
|
||||
def __call__(self, results):
|
||||
for i in range(self.iter_num):
|
||||
results = self.random_crop_flip(results)
|
||||
|
||||
return results
|
||||
|
||||
def random_crop_flip(self, results):
|
||||
image = results['image']
|
||||
polygons = results['polys']
|
||||
ignore_tags = results['ignore_tags']
|
||||
if len(polygons) == 0:
|
||||
return results
|
||||
|
||||
if np.random.random() >= self.crop_ratio:
|
||||
return results
|
||||
|
||||
h, w, _ = image.shape
|
||||
area = h * w
|
||||
pad_h = int(h * self.pad_ratio)
|
||||
pad_w = int(w * self.pad_ratio)
|
||||
h_axis, w_axis = self.generate_crop_target(image, polygons, pad_h,
|
||||
pad_w)
|
||||
if len(h_axis) == 0 or len(w_axis) == 0:
|
||||
return results
|
||||
|
||||
attempt = 0
|
||||
while attempt < 50:
|
||||
attempt += 1
|
||||
polys_keep = []
|
||||
polys_new = []
|
||||
ignore_tags_keep = []
|
||||
ignore_tags_new = []
|
||||
xx = np.random.choice(w_axis, size=2)
|
||||
xmin = np.min(xx) - pad_w
|
||||
xmax = np.max(xx) - pad_w
|
||||
xmin = np.clip(xmin, 0, w - 1)
|
||||
xmax = np.clip(xmax, 0, w - 1)
|
||||
yy = np.random.choice(h_axis, size=2)
|
||||
ymin = np.min(yy) - pad_h
|
||||
ymax = np.max(yy) - pad_h
|
||||
ymin = np.clip(ymin, 0, h - 1)
|
||||
ymax = np.clip(ymax, 0, h - 1)
|
||||
if (xmax - xmin) * (ymax - ymin) < area * self.min_area_ratio:
|
||||
# area too small
|
||||
continue
|
||||
|
||||
pts = np.stack([[xmin, xmax, xmax, xmin],
|
||||
[ymin, ymin, ymax, ymax]]).T.astype(np.int32)
|
||||
pp = Polygon(pts)
|
||||
fail_flag = False
|
||||
for polygon, ignore_tag in zip(polygons, ignore_tags):
|
||||
ppi = Polygon(polygon.reshape(-1, 2))
|
||||
ppiou, _ = poly_intersection(ppi, pp, buffer=0)
|
||||
if np.abs(ppiou - float(ppi.area)) > self.epsilon and \
|
||||
np.abs(ppiou) > self.epsilon:
|
||||
fail_flag = True
|
||||
break
|
||||
elif np.abs(ppiou - float(ppi.area)) < self.epsilon:
|
||||
polys_new.append(polygon)
|
||||
ignore_tags_new.append(ignore_tag)
|
||||
else:
|
||||
polys_keep.append(polygon)
|
||||
ignore_tags_keep.append(ignore_tag)
|
||||
|
||||
if fail_flag:
|
||||
continue
|
||||
else:
|
||||
break
|
||||
|
||||
cropped = image[ymin:ymax, xmin:xmax, :]
|
||||
select_type = np.random.randint(3)
|
||||
if select_type == 0:
|
||||
img = np.ascontiguousarray(cropped[:, ::-1])
|
||||
elif select_type == 1:
|
||||
img = np.ascontiguousarray(cropped[::-1, :])
|
||||
else:
|
||||
img = np.ascontiguousarray(cropped[::-1, ::-1])
|
||||
image[ymin:ymax, xmin:xmax, :] = img
|
||||
results['img'] = image
|
||||
|
||||
if len(polys_new) != 0:
|
||||
height, width, _ = cropped.shape
|
||||
if select_type == 0:
|
||||
for idx, polygon in enumerate(polys_new):
|
||||
poly = polygon.reshape(-1, 2)
|
||||
poly[:, 0] = width - poly[:, 0] + 2 * xmin
|
||||
polys_new[idx] = poly
|
||||
elif select_type == 1:
|
||||
for idx, polygon in enumerate(polys_new):
|
||||
poly = polygon.reshape(-1, 2)
|
||||
poly[:, 1] = height - poly[:, 1] + 2 * ymin
|
||||
polys_new[idx] = poly
|
||||
else:
|
||||
for idx, polygon in enumerate(polys_new):
|
||||
poly = polygon.reshape(-1, 2)
|
||||
poly[:, 0] = width - poly[:, 0] + 2 * xmin
|
||||
poly[:, 1] = height - poly[:, 1] + 2 * ymin
|
||||
polys_new[idx] = poly
|
||||
polygons = polys_keep + polys_new
|
||||
ignore_tags = ignore_tags_keep + ignore_tags_new
|
||||
results['polys'] = np.array(polygons)
|
||||
results['ignore_tags'] = ignore_tags
|
||||
|
||||
return results
|
||||
|
||||
def generate_crop_target(self, image, all_polys, pad_h, pad_w):
|
||||
"""Generate crop target and make sure not to crop the polygon
|
||||
instances.
|
||||
|
||||
Args:
|
||||
image (ndarray): The image waited to be crop.
|
||||
all_polys (list[list[ndarray]]): All polygons including ground
|
||||
truth polygons and ground truth ignored polygons.
|
||||
pad_h (int): Padding length of height.
|
||||
pad_w (int): Padding length of width.
|
||||
Returns:
|
||||
h_axis (ndarray): Vertical cropping range.
|
||||
w_axis (ndarray): Horizontal cropping range.
|
||||
"""
|
||||
h, w, _ = image.shape
|
||||
h_array = np.zeros((h + pad_h * 2), dtype=np.int32)
|
||||
w_array = np.zeros((w + pad_w * 2), dtype=np.int32)
|
||||
|
||||
text_polys = []
|
||||
for polygon in all_polys:
|
||||
rect = cv2.minAreaRect(polygon.astype(np.int32).reshape(-1, 2))
|
||||
box = cv2.boxPoints(rect)
|
||||
box = np.int0(box)
|
||||
text_polys.append([box[0], box[1], box[2], box[3]])
|
||||
|
||||
polys = np.array(text_polys, dtype=np.int32)
|
||||
for poly in polys:
|
||||
poly = np.round(poly, decimals=0).astype(np.int32)
|
||||
minx = np.min(poly[:, 0])
|
||||
maxx = np.max(poly[:, 0])
|
||||
w_array[minx + pad_w:maxx + pad_w] = 1
|
||||
miny = np.min(poly[:, 1])
|
||||
maxy = np.max(poly[:, 1])
|
||||
h_array[miny + pad_h:maxy + pad_h] = 1
|
||||
|
||||
h_axis = np.where(h_array == 0)[0]
|
||||
w_axis = np.where(w_array == 0)[0]
|
||||
return h_axis, w_axis
|
||||
|
||||
|
||||
class RandomCropPolyInstances:
|
||||
"""Randomly crop images and make sure to contain at least one intact
|
||||
instance."""
|
||||
|
||||
def __init__(self, crop_ratio=5.0 / 8.0, min_side_ratio=0.4, **kwargs):
|
||||
super().__init__()
|
||||
self.crop_ratio = crop_ratio
|
||||
self.min_side_ratio = min_side_ratio
|
||||
|
||||
def sample_valid_start_end(self, valid_array, min_len, max_start, min_end):
|
||||
|
||||
assert isinstance(min_len, int)
|
||||
assert len(valid_array) > min_len
|
||||
|
||||
start_array = valid_array.copy()
|
||||
max_start = min(len(start_array) - min_len, max_start)
|
||||
start_array[max_start:] = 0
|
||||
start_array[0] = 1
|
||||
diff_array = np.hstack([0, start_array]) - np.hstack([start_array, 0])
|
||||
region_starts = np.where(diff_array < 0)[0]
|
||||
region_ends = np.where(diff_array > 0)[0]
|
||||
region_ind = np.random.randint(0, len(region_starts))
|
||||
start = np.random.randint(region_starts[region_ind],
|
||||
region_ends[region_ind])
|
||||
|
||||
end_array = valid_array.copy()
|
||||
min_end = max(start + min_len, min_end)
|
||||
end_array[:min_end] = 0
|
||||
end_array[-1] = 1
|
||||
diff_array = np.hstack([0, end_array]) - np.hstack([end_array, 0])
|
||||
region_starts = np.where(diff_array < 0)[0]
|
||||
region_ends = np.where(diff_array > 0)[0]
|
||||
region_ind = np.random.randint(0, len(region_starts))
|
||||
end = np.random.randint(region_starts[region_ind],
|
||||
region_ends[region_ind])
|
||||
return start, end
|
||||
|
||||
def sample_crop_box(self, img_size, results):
|
||||
"""Generate crop box and make sure not to crop the polygon instances.
|
||||
|
||||
Args:
|
||||
img_size (tuple(int)): The image size (h, w).
|
||||
results (dict): The results dict.
|
||||
"""
|
||||
|
||||
assert isinstance(img_size, tuple)
|
||||
h, w = img_size[:2]
|
||||
|
||||
key_masks = results['polys']
|
||||
|
||||
x_valid_array = np.ones(w, dtype=np.int32)
|
||||
y_valid_array = np.ones(h, dtype=np.int32)
|
||||
|
||||
selected_mask = key_masks[np.random.randint(0, len(key_masks))]
|
||||
selected_mask = selected_mask.reshape((-1, 2)).astype(np.int32)
|
||||
max_x_start = max(np.min(selected_mask[:, 0]) - 2, 0)
|
||||
min_x_end = min(np.max(selected_mask[:, 0]) + 3, w - 1)
|
||||
max_y_start = max(np.min(selected_mask[:, 1]) - 2, 0)
|
||||
min_y_end = min(np.max(selected_mask[:, 1]) + 3, h - 1)
|
||||
|
||||
for mask in key_masks:
|
||||
mask = mask.reshape((-1, 2)).astype(np.int32)
|
||||
clip_x = np.clip(mask[:, 0], 0, w - 1)
|
||||
clip_y = np.clip(mask[:, 1], 0, h - 1)
|
||||
min_x, max_x = np.min(clip_x), np.max(clip_x)
|
||||
min_y, max_y = np.min(clip_y), np.max(clip_y)
|
||||
|
||||
x_valid_array[min_x - 2:max_x + 3] = 0
|
||||
y_valid_array[min_y - 2:max_y + 3] = 0
|
||||
|
||||
min_w = int(w * self.min_side_ratio)
|
||||
min_h = int(h * self.min_side_ratio)
|
||||
|
||||
x1, x2 = self.sample_valid_start_end(x_valid_array, min_w, max_x_start,
|
||||
min_x_end)
|
||||
y1, y2 = self.sample_valid_start_end(y_valid_array, min_h, max_y_start,
|
||||
min_y_end)
|
||||
|
||||
return np.array([x1, y1, x2, y2])
|
||||
|
||||
def crop_img(self, img, bbox):
|
||||
assert img.ndim == 3
|
||||
h, w, _ = img.shape
|
||||
assert 0 <= bbox[1] < bbox[3] <= h
|
||||
assert 0 <= bbox[0] < bbox[2] <= w
|
||||
return img[bbox[1]:bbox[3], bbox[0]:bbox[2]]
|
||||
|
||||
def __call__(self, results):
|
||||
image = results['image']
|
||||
polygons = results['polys']
|
||||
ignore_tags = results['ignore_tags']
|
||||
if len(polygons) < 1:
|
||||
return results
|
||||
|
||||
if np.random.random_sample() < self.crop_ratio:
|
||||
|
||||
crop_box = self.sample_crop_box(image.shape, results)
|
||||
img = self.crop_img(image, crop_box)
|
||||
results['image'] = img
|
||||
# crop and filter masks
|
||||
x1, y1, x2, y2 = crop_box
|
||||
w = max(x2 - x1, 1)
|
||||
h = max(y2 - y1, 1)
|
||||
polygons[:, :, 0::2] = polygons[:, :, 0::2] - x1
|
||||
polygons[:, :, 1::2] = polygons[:, :, 1::2] - y1
|
||||
|
||||
valid_masks_list = []
|
||||
valid_tags_list = []
|
||||
for ind, polygon in enumerate(polygons):
|
||||
if (polygon[:, ::2] > -4).all() and (
|
||||
polygon[:, ::2] < w + 4).all() and (
|
||||
polygon[:, 1::2] > -4).all() and (
|
||||
polygon[:, 1::2] < h + 4).all():
|
||||
polygon[:, ::2] = np.clip(polygon[:, ::2], 0, w)
|
||||
polygon[:, 1::2] = np.clip(polygon[:, 1::2], 0, h)
|
||||
valid_masks_list.append(polygon)
|
||||
valid_tags_list.append(ignore_tags[ind])
|
||||
|
||||
results['polys'] = np.array(valid_masks_list)
|
||||
results['ignore_tags'] = valid_tags_list
|
||||
|
||||
return results
|
||||
|
||||
def __repr__(self):
|
||||
repr_str = self.__class__.__name__
|
||||
return repr_str
|
||||
|
||||
|
||||
class RandomRotatePolyInstances:
|
||||
def __init__(self,
|
||||
rotate_ratio=0.5,
|
||||
max_angle=10,
|
||||
pad_with_fixed_color=False,
|
||||
pad_value=(0, 0, 0),
|
||||
**kwargs):
|
||||
"""Randomly rotate images and polygon masks.
|
||||
|
||||
Args:
|
||||
rotate_ratio (float): The ratio of samples to operate rotation.
|
||||
max_angle (int): The maximum rotation angle.
|
||||
pad_with_fixed_color (bool): The flag for whether to pad rotated
|
||||
image with fixed value. If set to False, the rotated image will
|
||||
be padded onto cropped image.
|
||||
pad_value (tuple(int)): The color value for padding rotated image.
|
||||
"""
|
||||
self.rotate_ratio = rotate_ratio
|
||||
self.max_angle = max_angle
|
||||
self.pad_with_fixed_color = pad_with_fixed_color
|
||||
self.pad_value = pad_value
|
||||
|
||||
def rotate(self, center, points, theta, center_shift=(0, 0)):
|
||||
# rotate points.
|
||||
(center_x, center_y) = center
|
||||
center_y = -center_y
|
||||
x, y = points[:, ::2], points[:, 1::2]
|
||||
y = -y
|
||||
|
||||
theta = theta / 180 * math.pi
|
||||
cos = math.cos(theta)
|
||||
sin = math.sin(theta)
|
||||
|
||||
x = (x - center_x)
|
||||
y = (y - center_y)
|
||||
|
||||
_x = center_x + x * cos - y * sin + center_shift[0]
|
||||
_y = -(center_y + x * sin + y * cos) + center_shift[1]
|
||||
|
||||
points[:, ::2], points[:, 1::2] = _x, _y
|
||||
return points
|
||||
|
||||
def cal_canvas_size(self, ori_size, degree):
|
||||
assert isinstance(ori_size, tuple)
|
||||
angle = degree * math.pi / 180.0
|
||||
h, w = ori_size[:2]
|
||||
|
||||
cos = math.cos(angle)
|
||||
sin = math.sin(angle)
|
||||
canvas_h = int(w * math.fabs(sin) + h * math.fabs(cos))
|
||||
canvas_w = int(w * math.fabs(cos) + h * math.fabs(sin))
|
||||
|
||||
canvas_size = (canvas_h, canvas_w)
|
||||
return canvas_size
|
||||
|
||||
def sample_angle(self, max_angle):
|
||||
angle = np.random.random_sample() * 2 * max_angle - max_angle
|
||||
return angle
|
||||
|
||||
def rotate_img(self, img, angle, canvas_size):
|
||||
h, w = img.shape[:2]
|
||||
rotation_matrix = cv2.getRotationMatrix2D((w / 2, h / 2), angle, 1)
|
||||
rotation_matrix[0, 2] += int((canvas_size[1] - w) / 2)
|
||||
rotation_matrix[1, 2] += int((canvas_size[0] - h) / 2)
|
||||
|
||||
if self.pad_with_fixed_color:
|
||||
target_img = cv2.warpAffine(
|
||||
img,
|
||||
rotation_matrix, (canvas_size[1], canvas_size[0]),
|
||||
flags=cv2.INTER_NEAREST,
|
||||
borderValue=self.pad_value)
|
||||
else:
|
||||
mask = np.zeros_like(img)
|
||||
(h_ind, w_ind) = (np.random.randint(0, h * 7 // 8),
|
||||
np.random.randint(0, w * 7 // 8))
|
||||
img_cut = img[h_ind:(h_ind + h // 9), w_ind:(w_ind + w // 9)]
|
||||
img_cut = cv2.resize(img_cut, (canvas_size[1], canvas_size[0]))
|
||||
|
||||
mask = cv2.warpAffine(
|
||||
mask,
|
||||
rotation_matrix, (canvas_size[1], canvas_size[0]),
|
||||
borderValue=[1, 1, 1])
|
||||
target_img = cv2.warpAffine(
|
||||
img,
|
||||
rotation_matrix, (canvas_size[1], canvas_size[0]),
|
||||
borderValue=[0, 0, 0])
|
||||
target_img = target_img + img_cut * mask
|
||||
|
||||
return target_img
|
||||
|
||||
def __call__(self, results):
|
||||
if np.random.random_sample() < self.rotate_ratio:
|
||||
image = results['image']
|
||||
polygons = results['polys']
|
||||
h, w = image.shape[:2]
|
||||
|
||||
angle = self.sample_angle(self.max_angle)
|
||||
canvas_size = self.cal_canvas_size((h, w), angle)
|
||||
center_shift = (int((canvas_size[1] - w) / 2), int(
|
||||
(canvas_size[0] - h) / 2))
|
||||
image = self.rotate_img(image, angle, canvas_size)
|
||||
results['image'] = image
|
||||
# rotate polygons
|
||||
rotated_masks = []
|
||||
for mask in polygons:
|
||||
rotated_mask = self.rotate((w / 2, h / 2), mask, angle,
|
||||
center_shift)
|
||||
rotated_masks.append(rotated_mask)
|
||||
results['polys'] = np.array(rotated_masks)
|
||||
|
||||
return results
|
||||
|
||||
def __repr__(self):
|
||||
repr_str = self.__class__.__name__
|
||||
return repr_str
|
||||
|
||||
|
||||
class SquareResizePad:
|
||||
def __init__(self,
|
||||
target_size,
|
||||
pad_ratio=0.6,
|
||||
pad_with_fixed_color=False,
|
||||
pad_value=(0, 0, 0),
|
||||
**kwargs):
|
||||
"""Resize or pad images to be square shape.
|
||||
|
||||
Args:
|
||||
target_size (int): The target size of square shaped image.
|
||||
pad_with_fixed_color (bool): The flag for whether to pad rotated
|
||||
image with fixed value. If set to False, the rescales image will
|
||||
be padded onto cropped image.
|
||||
pad_value (tuple(int)): The color value for padding rotated image.
|
||||
"""
|
||||
assert isinstance(target_size, int)
|
||||
assert isinstance(pad_ratio, float)
|
||||
assert isinstance(pad_with_fixed_color, bool)
|
||||
assert isinstance(pad_value, tuple)
|
||||
|
||||
self.target_size = target_size
|
||||
self.pad_ratio = pad_ratio
|
||||
self.pad_with_fixed_color = pad_with_fixed_color
|
||||
self.pad_value = pad_value
|
||||
|
||||
def resize_img(self, img, keep_ratio=True):
|
||||
h, w, _ = img.shape
|
||||
if keep_ratio:
|
||||
t_h = self.target_size if h >= w else int(h * self.target_size / w)
|
||||
t_w = self.target_size if h <= w else int(w * self.target_size / h)
|
||||
else:
|
||||
t_h = t_w = self.target_size
|
||||
img = cv2.resize(img, (t_w, t_h))
|
||||
return img, (t_h, t_w)
|
||||
|
||||
def square_pad(self, img):
|
||||
h, w = img.shape[:2]
|
||||
if h == w:
|
||||
return img, (0, 0)
|
||||
pad_size = max(h, w)
|
||||
if self.pad_with_fixed_color:
|
||||
expand_img = np.ones((pad_size, pad_size, 3), dtype=np.uint8)
|
||||
expand_img[:] = self.pad_value
|
||||
else:
|
||||
(h_ind, w_ind) = (np.random.randint(0, h * 7 // 8),
|
||||
np.random.randint(0, w * 7 // 8))
|
||||
img_cut = img[h_ind:(h_ind + h // 9), w_ind:(w_ind + w // 9)]
|
||||
expand_img = cv2.resize(img_cut, (pad_size, pad_size))
|
||||
if h > w:
|
||||
y0, x0 = 0, (h - w) // 2
|
||||
else:
|
||||
y0, x0 = (w - h) // 2, 0
|
||||
expand_img[y0:y0 + h, x0:x0 + w] = img
|
||||
offset = (x0, y0)
|
||||
|
||||
return expand_img, offset
|
||||
|
||||
def square_pad_mask(self, points, offset):
|
||||
x0, y0 = offset
|
||||
pad_points = points.copy()
|
||||
pad_points[::2] = pad_points[::2] + x0
|
||||
pad_points[1::2] = pad_points[1::2] + y0
|
||||
return pad_points
|
||||
|
||||
def __call__(self, results):
|
||||
image = results['image']
|
||||
polygons = results['polys']
|
||||
h, w = image.shape[:2]
|
||||
|
||||
if np.random.random_sample() < self.pad_ratio:
|
||||
image, out_size = self.resize_img(image, keep_ratio=True)
|
||||
image, offset = self.square_pad(image)
|
||||
else:
|
||||
image, out_size = self.resize_img(image, keep_ratio=False)
|
||||
offset = (0, 0)
|
||||
results['image'] = image
|
||||
try:
|
||||
polygons[:, :, 0::2] = polygons[:, :, 0::2] * out_size[
|
||||
1] / w + offset[0]
|
||||
polygons[:, :, 1::2] = polygons[:, :, 1::2] * out_size[
|
||||
0] / h + offset[1]
|
||||
except:
|
||||
pass
|
||||
results['polys'] = polygons
|
||||
|
||||
return results
|
||||
|
||||
def __repr__(self):
|
||||
repr_str = self.__class__.__name__
|
||||
return repr_str
|
||||
@@ -0,0 +1,658 @@
|
||||
# copyright (c) 2022 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""
|
||||
This code is refer from:
|
||||
https://github.com/open-mmlab/mmocr/blob/main/mmocr/datasets/pipelines/textdet_targets/fcenet_targets.py
|
||||
"""
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from numpy.fft import fft
|
||||
from numpy.linalg import norm
|
||||
import sys
|
||||
|
||||
|
||||
class FCENetTargets:
|
||||
"""Generate the ground truth targets of FCENet: Fourier Contour Embedding
|
||||
for Arbitrary-Shaped Text Detection.
|
||||
|
||||
[https://arxiv.org/abs/2104.10442]
|
||||
|
||||
Args:
|
||||
fourier_degree (int): The maximum Fourier transform degree k.
|
||||
resample_step (float): The step size for resampling the text center
|
||||
line (TCL). It's better not to exceed half of the minimum width.
|
||||
center_region_shrink_ratio (float): The shrink ratio of text center
|
||||
region.
|
||||
level_size_divisors (tuple(int)): The downsample ratio on each level.
|
||||
level_proportion_range (tuple(tuple(int))): The range of text sizes
|
||||
assigned to each level.
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
fourier_degree=5,
|
||||
resample_step=4.0,
|
||||
center_region_shrink_ratio=0.3,
|
||||
level_size_divisors=(8, 16, 32),
|
||||
level_proportion_range=((0, 0.25), (0.2, 0.65), (0.55, 1.0)),
|
||||
orientation_thr=2.0,
|
||||
**kwargs):
|
||||
|
||||
super().__init__()
|
||||
assert isinstance(level_size_divisors, tuple)
|
||||
assert isinstance(level_proportion_range, tuple)
|
||||
assert len(level_size_divisors) == len(level_proportion_range)
|
||||
self.fourier_degree = fourier_degree
|
||||
self.resample_step = resample_step
|
||||
self.center_region_shrink_ratio = center_region_shrink_ratio
|
||||
self.level_size_divisors = level_size_divisors
|
||||
self.level_proportion_range = level_proportion_range
|
||||
|
||||
self.orientation_thr = orientation_thr
|
||||
|
||||
def vector_angle(self, vec1, vec2):
|
||||
if vec1.ndim > 1:
|
||||
unit_vec1 = vec1 / (norm(vec1, axis=-1) + 1e-8).reshape((-1, 1))
|
||||
else:
|
||||
unit_vec1 = vec1 / (norm(vec1, axis=-1) + 1e-8)
|
||||
if vec2.ndim > 1:
|
||||
unit_vec2 = vec2 / (norm(vec2, axis=-1) + 1e-8).reshape((-1, 1))
|
||||
else:
|
||||
unit_vec2 = vec2 / (norm(vec2, axis=-1) + 1e-8)
|
||||
return np.arccos(
|
||||
np.clip(
|
||||
np.sum(unit_vec1 * unit_vec2, axis=-1), -1.0, 1.0))
|
||||
|
||||
def resample_line(self, line, n):
|
||||
"""Resample n points on a line.
|
||||
|
||||
Args:
|
||||
line (ndarray): The points composing a line.
|
||||
n (int): The resampled points number.
|
||||
|
||||
Returns:
|
||||
resampled_line (ndarray): The points composing the resampled line.
|
||||
"""
|
||||
|
||||
assert line.ndim == 2
|
||||
assert line.shape[0] >= 2
|
||||
assert line.shape[1] == 2
|
||||
assert isinstance(n, int)
|
||||
assert n > 0
|
||||
|
||||
length_list = [
|
||||
norm(line[i + 1] - line[i]) for i in range(len(line) - 1)
|
||||
]
|
||||
total_length = sum(length_list)
|
||||
length_cumsum = np.cumsum([0.0] + length_list)
|
||||
delta_length = total_length / (float(n) + 1e-8)
|
||||
|
||||
current_edge_ind = 0
|
||||
resampled_line = [line[0]]
|
||||
|
||||
for i in range(1, n):
|
||||
current_line_len = i * delta_length
|
||||
|
||||
while current_line_len >= length_cumsum[current_edge_ind + 1]:
|
||||
current_edge_ind += 1
|
||||
current_edge_end_shift = current_line_len - length_cumsum[
|
||||
current_edge_ind]
|
||||
end_shift_ratio = current_edge_end_shift / length_list[
|
||||
current_edge_ind]
|
||||
current_point = line[current_edge_ind] + (line[current_edge_ind + 1]
|
||||
- line[current_edge_ind]
|
||||
) * end_shift_ratio
|
||||
resampled_line.append(current_point)
|
||||
|
||||
resampled_line.append(line[-1])
|
||||
resampled_line = np.array(resampled_line)
|
||||
|
||||
return resampled_line
|
||||
|
||||
def reorder_poly_edge(self, points):
|
||||
"""Get the respective points composing head edge, tail edge, top
|
||||
sideline and bottom sideline.
|
||||
|
||||
Args:
|
||||
points (ndarray): The points composing a text polygon.
|
||||
|
||||
Returns:
|
||||
head_edge (ndarray): The two points composing the head edge of text
|
||||
polygon.
|
||||
tail_edge (ndarray): The two points composing the tail edge of text
|
||||
polygon.
|
||||
top_sideline (ndarray): The points composing top curved sideline of
|
||||
text polygon.
|
||||
bot_sideline (ndarray): The points composing bottom curved sideline
|
||||
of text polygon.
|
||||
"""
|
||||
|
||||
assert points.ndim == 2
|
||||
assert points.shape[0] >= 4
|
||||
assert points.shape[1] == 2
|
||||
|
||||
head_inds, tail_inds = self.find_head_tail(points, self.orientation_thr)
|
||||
head_edge, tail_edge = points[head_inds], points[tail_inds]
|
||||
|
||||
pad_points = np.vstack([points, points])
|
||||
if tail_inds[1] < 1:
|
||||
tail_inds[1] = len(points)
|
||||
sideline1 = pad_points[head_inds[1]:tail_inds[1]]
|
||||
sideline2 = pad_points[tail_inds[1]:(head_inds[1] + len(points))]
|
||||
sideline_mean_shift = np.mean(
|
||||
sideline1, axis=0) - np.mean(
|
||||
sideline2, axis=0)
|
||||
|
||||
if sideline_mean_shift[1] > 0:
|
||||
top_sideline, bot_sideline = sideline2, sideline1
|
||||
else:
|
||||
top_sideline, bot_sideline = sideline1, sideline2
|
||||
|
||||
return head_edge, tail_edge, top_sideline, bot_sideline
|
||||
|
||||
def find_head_tail(self, points, orientation_thr):
|
||||
"""Find the head edge and tail edge of a text polygon.
|
||||
|
||||
Args:
|
||||
points (ndarray): The points composing a text polygon.
|
||||
orientation_thr (float): The threshold for distinguishing between
|
||||
head edge and tail edge among the horizontal and vertical edges
|
||||
of a quadrangle.
|
||||
|
||||
Returns:
|
||||
head_inds (list): The indexes of two points composing head edge.
|
||||
tail_inds (list): The indexes of two points composing tail edge.
|
||||
"""
|
||||
|
||||
assert points.ndim == 2
|
||||
assert points.shape[0] >= 4
|
||||
assert points.shape[1] == 2
|
||||
assert isinstance(orientation_thr, float)
|
||||
|
||||
if len(points) > 4:
|
||||
pad_points = np.vstack([points, points[0]])
|
||||
edge_vec = pad_points[1:] - pad_points[:-1]
|
||||
|
||||
theta_sum = []
|
||||
adjacent_vec_theta = []
|
||||
for i, edge_vec1 in enumerate(edge_vec):
|
||||
adjacent_ind = [x % len(edge_vec) for x in [i - 1, i + 1]]
|
||||
adjacent_edge_vec = edge_vec[adjacent_ind]
|
||||
temp_theta_sum = np.sum(
|
||||
self.vector_angle(edge_vec1, adjacent_edge_vec))
|
||||
temp_adjacent_theta = self.vector_angle(adjacent_edge_vec[0],
|
||||
adjacent_edge_vec[1])
|
||||
theta_sum.append(temp_theta_sum)
|
||||
adjacent_vec_theta.append(temp_adjacent_theta)
|
||||
theta_sum_score = np.array(theta_sum) / np.pi
|
||||
adjacent_theta_score = np.array(adjacent_vec_theta) / np.pi
|
||||
poly_center = np.mean(points, axis=0)
|
||||
edge_dist = np.maximum(
|
||||
norm(
|
||||
pad_points[1:] - poly_center, axis=-1),
|
||||
norm(
|
||||
pad_points[:-1] - poly_center, axis=-1))
|
||||
dist_score = edge_dist / np.max(edge_dist)
|
||||
position_score = np.zeros(len(edge_vec))
|
||||
score = 0.5 * theta_sum_score + 0.15 * adjacent_theta_score
|
||||
score += 0.35 * dist_score
|
||||
if len(points) % 2 == 0:
|
||||
position_score[(len(score) // 2 - 1)] += 1
|
||||
position_score[-1] += 1
|
||||
score += 0.1 * position_score
|
||||
pad_score = np.concatenate([score, score])
|
||||
score_matrix = np.zeros((len(score), len(score) - 3))
|
||||
x = np.arange(len(score) - 3) / float(len(score) - 4)
|
||||
gaussian = 1. / (np.sqrt(2. * np.pi) * 0.5) * np.exp(-np.power(
|
||||
(x - 0.5) / 0.5, 2.) / 2)
|
||||
gaussian = gaussian / np.max(gaussian)
|
||||
for i in range(len(score)):
|
||||
score_matrix[i, :] = score[i] + pad_score[(i + 2):(i + len(
|
||||
score) - 1)] * gaussian * 0.3
|
||||
|
||||
head_start, tail_increment = np.unravel_index(score_matrix.argmax(),
|
||||
score_matrix.shape)
|
||||
tail_start = (head_start + tail_increment + 2) % len(points)
|
||||
head_end = (head_start + 1) % len(points)
|
||||
tail_end = (tail_start + 1) % len(points)
|
||||
|
||||
if head_end > tail_end:
|
||||
head_start, tail_start = tail_start, head_start
|
||||
head_end, tail_end = tail_end, head_end
|
||||
head_inds = [head_start, head_end]
|
||||
tail_inds = [tail_start, tail_end]
|
||||
else:
|
||||
if self.vector_slope(points[1] - points[0]) + self.vector_slope(
|
||||
points[3] - points[2]) < self.vector_slope(points[
|
||||
2] - points[1]) + self.vector_slope(points[0] - points[
|
||||
3]):
|
||||
horizontal_edge_inds = [[0, 1], [2, 3]]
|
||||
vertical_edge_inds = [[3, 0], [1, 2]]
|
||||
else:
|
||||
horizontal_edge_inds = [[3, 0], [1, 2]]
|
||||
vertical_edge_inds = [[0, 1], [2, 3]]
|
||||
|
||||
vertical_len_sum = norm(points[vertical_edge_inds[0][0]] - points[
|
||||
vertical_edge_inds[0][1]]) + norm(points[vertical_edge_inds[1][
|
||||
0]] - points[vertical_edge_inds[1][1]])
|
||||
horizontal_len_sum = norm(points[horizontal_edge_inds[0][
|
||||
0]] - points[horizontal_edge_inds[0][1]]) + norm(points[
|
||||
horizontal_edge_inds[1][0]] - points[horizontal_edge_inds[1]
|
||||
[1]])
|
||||
|
||||
if vertical_len_sum > horizontal_len_sum * orientation_thr:
|
||||
head_inds = horizontal_edge_inds[0]
|
||||
tail_inds = horizontal_edge_inds[1]
|
||||
else:
|
||||
head_inds = vertical_edge_inds[0]
|
||||
tail_inds = vertical_edge_inds[1]
|
||||
|
||||
return head_inds, tail_inds
|
||||
|
||||
def resample_sidelines(self, sideline1, sideline2, resample_step):
|
||||
"""Resample two sidelines to be of the same points number according to
|
||||
step size.
|
||||
|
||||
Args:
|
||||
sideline1 (ndarray): The points composing a sideline of a text
|
||||
polygon.
|
||||
sideline2 (ndarray): The points composing another sideline of a
|
||||
text polygon.
|
||||
resample_step (float): The resampled step size.
|
||||
|
||||
Returns:
|
||||
resampled_line1 (ndarray): The resampled line 1.
|
||||
resampled_line2 (ndarray): The resampled line 2.
|
||||
"""
|
||||
|
||||
assert sideline1.ndim == sideline2.ndim == 2
|
||||
assert sideline1.shape[1] == sideline2.shape[1] == 2
|
||||
assert sideline1.shape[0] >= 2
|
||||
assert sideline2.shape[0] >= 2
|
||||
assert isinstance(resample_step, float)
|
||||
|
||||
length1 = sum([
|
||||
norm(sideline1[i + 1] - sideline1[i])
|
||||
for i in range(len(sideline1) - 1)
|
||||
])
|
||||
length2 = sum([
|
||||
norm(sideline2[i + 1] - sideline2[i])
|
||||
for i in range(len(sideline2) - 1)
|
||||
])
|
||||
|
||||
total_length = (length1 + length2) / 2
|
||||
resample_point_num = max(int(float(total_length) / resample_step), 1)
|
||||
|
||||
resampled_line1 = self.resample_line(sideline1, resample_point_num)
|
||||
resampled_line2 = self.resample_line(sideline2, resample_point_num)
|
||||
|
||||
return resampled_line1, resampled_line2
|
||||
|
||||
def generate_center_region_mask(self, img_size, text_polys):
|
||||
"""Generate text center region mask.
|
||||
|
||||
Args:
|
||||
img_size (tuple): The image size of (height, width).
|
||||
text_polys (list[list[ndarray]]): The list of text polygons.
|
||||
|
||||
Returns:
|
||||
center_region_mask (ndarray): The text center region mask.
|
||||
"""
|
||||
|
||||
assert isinstance(img_size, tuple)
|
||||
# assert check_argument.is_2dlist(text_polys)
|
||||
|
||||
h, w = img_size
|
||||
|
||||
center_region_mask = np.zeros((h, w), np.uint8)
|
||||
|
||||
center_region_boxes = []
|
||||
for poly in text_polys:
|
||||
# assert len(poly) == 1
|
||||
polygon_points = poly.reshape(-1, 2)
|
||||
_, _, top_line, bot_line = self.reorder_poly_edge(polygon_points)
|
||||
resampled_top_line, resampled_bot_line = self.resample_sidelines(
|
||||
top_line, bot_line, self.resample_step)
|
||||
resampled_bot_line = resampled_bot_line[::-1]
|
||||
center_line = (resampled_top_line + resampled_bot_line) / 2
|
||||
|
||||
line_head_shrink_len = norm(resampled_top_line[0] -
|
||||
resampled_bot_line[0]) / 4.0
|
||||
line_tail_shrink_len = norm(resampled_top_line[-1] -
|
||||
resampled_bot_line[-1]) / 4.0
|
||||
head_shrink_num = int(line_head_shrink_len // self.resample_step)
|
||||
tail_shrink_num = int(line_tail_shrink_len // self.resample_step)
|
||||
if len(center_line) > head_shrink_num + tail_shrink_num + 2:
|
||||
center_line = center_line[head_shrink_num:len(center_line) -
|
||||
tail_shrink_num]
|
||||
resampled_top_line = resampled_top_line[head_shrink_num:len(
|
||||
resampled_top_line) - tail_shrink_num]
|
||||
resampled_bot_line = resampled_bot_line[head_shrink_num:len(
|
||||
resampled_bot_line) - tail_shrink_num]
|
||||
|
||||
for i in range(0, len(center_line) - 1):
|
||||
tl = center_line[i] + (resampled_top_line[i] - center_line[i]
|
||||
) * self.center_region_shrink_ratio
|
||||
tr = center_line[i + 1] + (resampled_top_line[i + 1] -
|
||||
center_line[i + 1]
|
||||
) * self.center_region_shrink_ratio
|
||||
br = center_line[i + 1] + (resampled_bot_line[i + 1] -
|
||||
center_line[i + 1]
|
||||
) * self.center_region_shrink_ratio
|
||||
bl = center_line[i] + (resampled_bot_line[i] - center_line[i]
|
||||
) * self.center_region_shrink_ratio
|
||||
current_center_box = np.vstack([tl, tr, br,
|
||||
bl]).astype(np.int32)
|
||||
center_region_boxes.append(current_center_box)
|
||||
|
||||
cv2.fillPoly(center_region_mask, center_region_boxes, 1)
|
||||
return center_region_mask
|
||||
|
||||
def resample_polygon(self, polygon, n=400):
|
||||
"""Resample one polygon with n points on its boundary.
|
||||
|
||||
Args:
|
||||
polygon (list[float]): The input polygon.
|
||||
n (int): The number of resampled points.
|
||||
Returns:
|
||||
resampled_polygon (list[float]): The resampled polygon.
|
||||
"""
|
||||
length = []
|
||||
|
||||
for i in range(len(polygon)):
|
||||
p1 = polygon[i]
|
||||
if i == len(polygon) - 1:
|
||||
p2 = polygon[0]
|
||||
else:
|
||||
p2 = polygon[i + 1]
|
||||
length.append(((p1[0] - p2[0])**2 + (p1[1] - p2[1])**2)**0.5)
|
||||
|
||||
total_length = sum(length)
|
||||
n_on_each_line = (np.array(length) / (total_length + 1e-8)) * n
|
||||
n_on_each_line = n_on_each_line.astype(np.int32)
|
||||
new_polygon = []
|
||||
|
||||
for i in range(len(polygon)):
|
||||
num = n_on_each_line[i]
|
||||
p1 = polygon[i]
|
||||
if i == len(polygon) - 1:
|
||||
p2 = polygon[0]
|
||||
else:
|
||||
p2 = polygon[i + 1]
|
||||
|
||||
if num == 0:
|
||||
continue
|
||||
|
||||
dxdy = (p2 - p1) / num
|
||||
for j in range(num):
|
||||
point = p1 + dxdy * j
|
||||
new_polygon.append(point)
|
||||
|
||||
return np.array(new_polygon)
|
||||
|
||||
def normalize_polygon(self, polygon):
|
||||
"""Normalize one polygon so that its start point is at right most.
|
||||
|
||||
Args:
|
||||
polygon (list[float]): The origin polygon.
|
||||
Returns:
|
||||
new_polygon (lost[float]): The polygon with start point at right.
|
||||
"""
|
||||
temp_polygon = polygon - polygon.mean(axis=0)
|
||||
x = np.abs(temp_polygon[:, 0])
|
||||
y = temp_polygon[:, 1]
|
||||
index_x = np.argsort(x)
|
||||
index_y = np.argmin(y[index_x[:8]])
|
||||
index = index_x[index_y]
|
||||
new_polygon = np.concatenate([polygon[index:], polygon[:index]])
|
||||
return new_polygon
|
||||
|
||||
def poly2fourier(self, polygon, fourier_degree):
|
||||
"""Perform Fourier transformation to generate Fourier coefficients ck
|
||||
from polygon.
|
||||
|
||||
Args:
|
||||
polygon (ndarray): An input polygon.
|
||||
fourier_degree (int): The maximum Fourier degree K.
|
||||
Returns:
|
||||
c (ndarray(complex)): Fourier coefficients.
|
||||
"""
|
||||
points = polygon[:, 0] + polygon[:, 1] * 1j
|
||||
c_fft = fft(points) / len(points)
|
||||
c = np.hstack((c_fft[-fourier_degree:], c_fft[:fourier_degree + 1]))
|
||||
return c
|
||||
|
||||
def clockwise(self, c, fourier_degree):
|
||||
"""Make sure the polygon reconstructed from Fourier coefficients c in
|
||||
the clockwise direction.
|
||||
|
||||
Args:
|
||||
polygon (list[float]): The origin polygon.
|
||||
Returns:
|
||||
new_polygon (lost[float]): The polygon in clockwise point order.
|
||||
"""
|
||||
if np.abs(c[fourier_degree + 1]) > np.abs(c[fourier_degree - 1]):
|
||||
return c
|
||||
elif np.abs(c[fourier_degree + 1]) < np.abs(c[fourier_degree - 1]):
|
||||
return c[::-1]
|
||||
else:
|
||||
if np.abs(c[fourier_degree + 2]) > np.abs(c[fourier_degree - 2]):
|
||||
return c
|
||||
else:
|
||||
return c[::-1]
|
||||
|
||||
def cal_fourier_signature(self, polygon, fourier_degree):
|
||||
"""Calculate Fourier signature from input polygon.
|
||||
|
||||
Args:
|
||||
polygon (ndarray): The input polygon.
|
||||
fourier_degree (int): The maximum Fourier degree K.
|
||||
Returns:
|
||||
fourier_signature (ndarray): An array shaped (2k+1, 2) containing
|
||||
real part and image part of 2k+1 Fourier coefficients.
|
||||
"""
|
||||
resampled_polygon = self.resample_polygon(polygon)
|
||||
resampled_polygon = self.normalize_polygon(resampled_polygon)
|
||||
|
||||
fourier_coeff = self.poly2fourier(resampled_polygon, fourier_degree)
|
||||
fourier_coeff = self.clockwise(fourier_coeff, fourier_degree)
|
||||
|
||||
real_part = np.real(fourier_coeff).reshape((-1, 1))
|
||||
image_part = np.imag(fourier_coeff).reshape((-1, 1))
|
||||
fourier_signature = np.hstack([real_part, image_part])
|
||||
|
||||
return fourier_signature
|
||||
|
||||
def generate_fourier_maps(self, img_size, text_polys):
|
||||
"""Generate Fourier coefficient maps.
|
||||
|
||||
Args:
|
||||
img_size (tuple): The image size of (height, width).
|
||||
text_polys (list[list[ndarray]]): The list of text polygons.
|
||||
|
||||
Returns:
|
||||
fourier_real_map (ndarray): The Fourier coefficient real part maps.
|
||||
fourier_image_map (ndarray): The Fourier coefficient image part
|
||||
maps.
|
||||
"""
|
||||
|
||||
assert isinstance(img_size, tuple)
|
||||
|
||||
h, w = img_size
|
||||
k = self.fourier_degree
|
||||
real_map = np.zeros((k * 2 + 1, h, w), dtype=np.float32)
|
||||
imag_map = np.zeros((k * 2 + 1, h, w), dtype=np.float32)
|
||||
|
||||
for poly in text_polys:
|
||||
mask = np.zeros((h, w), dtype=np.uint8)
|
||||
polygon = np.array(poly).reshape((1, -1, 2))
|
||||
cv2.fillPoly(mask, polygon.astype(np.int32), 1)
|
||||
fourier_coeff = self.cal_fourier_signature(polygon[0], k)
|
||||
for i in range(-k, k + 1):
|
||||
if i != 0:
|
||||
real_map[i + k, :, :] = mask * fourier_coeff[i + k, 0] + (
|
||||
1 - mask) * real_map[i + k, :, :]
|
||||
imag_map[i + k, :, :] = mask * fourier_coeff[i + k, 1] + (
|
||||
1 - mask) * imag_map[i + k, :, :]
|
||||
else:
|
||||
yx = np.argwhere(mask > 0.5)
|
||||
k_ind = np.ones((len(yx)), dtype=np.int64) * k
|
||||
y, x = yx[:, 0], yx[:, 1]
|
||||
real_map[k_ind, y, x] = fourier_coeff[k, 0] - x
|
||||
imag_map[k_ind, y, x] = fourier_coeff[k, 1] - y
|
||||
|
||||
return real_map, imag_map
|
||||
|
||||
def generate_text_region_mask(self, img_size, text_polys):
|
||||
"""Generate text center region mask and geometry attribute maps.
|
||||
|
||||
Args:
|
||||
img_size (tuple): The image size (height, width).
|
||||
text_polys (list[list[ndarray]]): The list of text polygons.
|
||||
|
||||
Returns:
|
||||
text_region_mask (ndarray): The text region mask.
|
||||
"""
|
||||
|
||||
assert isinstance(img_size, tuple)
|
||||
|
||||
h, w = img_size
|
||||
text_region_mask = np.zeros((h, w), dtype=np.uint8)
|
||||
|
||||
for poly in text_polys:
|
||||
polygon = np.array(poly, dtype=np.int32).reshape((1, -1, 2))
|
||||
cv2.fillPoly(text_region_mask, polygon, 1)
|
||||
|
||||
return text_region_mask
|
||||
|
||||
def generate_effective_mask(self, mask_size: tuple, polygons_ignore):
|
||||
"""Generate effective mask by setting the ineffective regions to 0 and
|
||||
effective regions to 1.
|
||||
|
||||
Args:
|
||||
mask_size (tuple): The mask size.
|
||||
polygons_ignore (list[[ndarray]]: The list of ignored text
|
||||
polygons.
|
||||
|
||||
Returns:
|
||||
mask (ndarray): The effective mask of (height, width).
|
||||
"""
|
||||
|
||||
mask = np.ones(mask_size, dtype=np.uint8)
|
||||
|
||||
for poly in polygons_ignore:
|
||||
instance = poly.reshape(-1, 2).astype(np.int32).reshape(1, -1, 2)
|
||||
cv2.fillPoly(mask, instance, 0)
|
||||
|
||||
return mask
|
||||
|
||||
def generate_level_targets(self, img_size, text_polys, ignore_polys):
|
||||
"""Generate ground truth target on each level.
|
||||
|
||||
Args:
|
||||
img_size (list[int]): Shape of input image.
|
||||
text_polys (list[list[ndarray]]): A list of ground truth polygons.
|
||||
ignore_polys (list[list[ndarray]]): A list of ignored polygons.
|
||||
Returns:
|
||||
level_maps (list(ndarray)): A list of ground target on each level.
|
||||
"""
|
||||
h, w = img_size
|
||||
lv_size_divs = self.level_size_divisors
|
||||
lv_proportion_range = self.level_proportion_range
|
||||
lv_text_polys = [[] for i in range(len(lv_size_divs))]
|
||||
lv_ignore_polys = [[] for i in range(len(lv_size_divs))]
|
||||
level_maps = []
|
||||
for poly in text_polys:
|
||||
polygon = np.array(poly, dtype=np.int).reshape((1, -1, 2))
|
||||
_, _, box_w, box_h = cv2.boundingRect(polygon)
|
||||
proportion = max(box_h, box_w) / (h + 1e-8)
|
||||
|
||||
for ind, proportion_range in enumerate(lv_proportion_range):
|
||||
if proportion_range[0] < proportion < proportion_range[1]:
|
||||
lv_text_polys[ind].append(poly / lv_size_divs[ind])
|
||||
|
||||
for ignore_poly in ignore_polys:
|
||||
polygon = np.array(ignore_poly, dtype=np.int).reshape((1, -1, 2))
|
||||
_, _, box_w, box_h = cv2.boundingRect(polygon)
|
||||
proportion = max(box_h, box_w) / (h + 1e-8)
|
||||
|
||||
for ind, proportion_range in enumerate(lv_proportion_range):
|
||||
if proportion_range[0] < proportion < proportion_range[1]:
|
||||
lv_ignore_polys[ind].append(ignore_poly / lv_size_divs[ind])
|
||||
|
||||
for ind, size_divisor in enumerate(lv_size_divs):
|
||||
current_level_maps = []
|
||||
level_img_size = (h // size_divisor, w // size_divisor)
|
||||
|
||||
text_region = self.generate_text_region_mask(
|
||||
level_img_size, lv_text_polys[ind])[None]
|
||||
current_level_maps.append(text_region)
|
||||
|
||||
center_region = self.generate_center_region_mask(
|
||||
level_img_size, lv_text_polys[ind])[None]
|
||||
current_level_maps.append(center_region)
|
||||
|
||||
effective_mask = self.generate_effective_mask(
|
||||
level_img_size, lv_ignore_polys[ind])[None]
|
||||
current_level_maps.append(effective_mask)
|
||||
|
||||
fourier_real_map, fourier_image_maps = self.generate_fourier_maps(
|
||||
level_img_size, lv_text_polys[ind])
|
||||
current_level_maps.append(fourier_real_map)
|
||||
current_level_maps.append(fourier_image_maps)
|
||||
|
||||
level_maps.append(np.concatenate(current_level_maps))
|
||||
|
||||
return level_maps
|
||||
|
||||
def generate_targets(self, results):
|
||||
"""Generate the ground truth targets for FCENet.
|
||||
|
||||
Args:
|
||||
results (dict): The input result dictionary.
|
||||
|
||||
Returns:
|
||||
results (dict): The output result dictionary.
|
||||
"""
|
||||
|
||||
assert isinstance(results, dict)
|
||||
image = results['image']
|
||||
polygons = results['polys']
|
||||
ignore_tags = results['ignore_tags']
|
||||
h, w, _ = image.shape
|
||||
|
||||
polygon_masks = []
|
||||
polygon_masks_ignore = []
|
||||
for tag, polygon in zip(ignore_tags, polygons):
|
||||
if tag is True:
|
||||
polygon_masks_ignore.append(polygon)
|
||||
else:
|
||||
polygon_masks.append(polygon)
|
||||
|
||||
level_maps = self.generate_level_targets((h, w), polygon_masks,
|
||||
polygon_masks_ignore)
|
||||
|
||||
mapping = {
|
||||
'p3_maps': level_maps[0],
|
||||
'p4_maps': level_maps[1],
|
||||
'p5_maps': level_maps[2]
|
||||
}
|
||||
for key, value in mapping.items():
|
||||
results[key] = value
|
||||
|
||||
return results
|
||||
|
||||
def __call__(self, results):
|
||||
results = self.generate_targets(results)
|
||||
return results
|
||||
@@ -0,0 +1,244 @@
|
||||
"""
|
||||
# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
from __future__ import unicode_literals
|
||||
|
||||
import sys
|
||||
import six
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
|
||||
class GenTableMask(object):
|
||||
""" gen table mask """
|
||||
|
||||
def __init__(self, shrink_h_max, shrink_w_max, mask_type=0, **kwargs):
|
||||
self.shrink_h_max = 5
|
||||
self.shrink_w_max = 5
|
||||
self.mask_type = mask_type
|
||||
|
||||
def projection(self, erosion, h, w, spilt_threshold=0):
|
||||
# 水平投影
|
||||
projection_map = np.ones_like(erosion)
|
||||
project_val_array = [0 for _ in range(0, h)]
|
||||
|
||||
for j in range(0, h):
|
||||
for i in range(0, w):
|
||||
if erosion[j, i] == 255:
|
||||
project_val_array[j] += 1
|
||||
# 根据数组，获取切割点
|
||||
start_idx = 0 # 记录进入字符区的索引
|
||||
end_idx = 0 # 记录进入空白区域的索引
|
||||
in_text = False # 是否遍历到了字符区内
|
||||
box_list = []
|
||||
for i in range(len(project_val_array)):
|
||||
if in_text == False and project_val_array[i] > spilt_threshold: # 进入字符区了
|
||||
in_text = True
|
||||
start_idx = i
|
||||
elif project_val_array[i] <= spilt_threshold and in_text == True: # 进入空白区了
|
||||
end_idx = i
|
||||
in_text = False
|
||||
if end_idx - start_idx <= 2:
|
||||
continue
|
||||
box_list.append((start_idx, end_idx + 1))
|
||||
|
||||
if in_text:
|
||||
box_list.append((start_idx, h - 1))
|
||||
# 绘制投影直方图
|
||||
for j in range(0, h):
|
||||
for i in range(0, project_val_array[j]):
|
||||
projection_map[j, i] = 0
|
||||
return box_list, projection_map
|
||||
|
||||
def projection_cx(self, box_img):
|
||||
box_gray_img = cv2.cvtColor(box_img, cv2.COLOR_BGR2GRAY)
|
||||
h, w = box_gray_img.shape
|
||||
# 灰度图片进行二值化处理
|
||||
ret, thresh1 = cv2.threshold(box_gray_img, 200, 255, cv2.THRESH_BINARY_INV)
|
||||
# 纵向腐蚀
|
||||
if h < w:
|
||||
kernel = np.ones((2, 1), np.uint8)
|
||||
erode = cv2.erode(thresh1, kernel, iterations=1)
|
||||
else:
|
||||
erode = thresh1
|
||||
# 水平膨胀
|
||||
kernel = np.ones((1, 5), np.uint8)
|
||||
erosion = cv2.dilate(erode, kernel, iterations=1)
|
||||
# 水平投影
|
||||
projection_map = np.ones_like(erosion)
|
||||
project_val_array = [0 for _ in range(0, h)]
|
||||
|
||||
for j in range(0, h):
|
||||
for i in range(0, w):
|
||||
if erosion[j, i] == 255:
|
||||
project_val_array[j] += 1
|
||||
# 根据数组，获取切割点
|
||||
start_idx = 0 # 记录进入字符区的索引
|
||||
end_idx = 0 # 记录进入空白区域的索引
|
||||
in_text = False # 是否遍历到了字符区内
|
||||
box_list = []
|
||||
spilt_threshold = 0
|
||||
for i in range(len(project_val_array)):
|
||||
if in_text == False and project_val_array[i] > spilt_threshold: # 进入字符区了
|
||||
in_text = True
|
||||
start_idx = i
|
||||
elif project_val_array[i] <= spilt_threshold and in_text == True: # 进入空白区了
|
||||
end_idx = i
|
||||
in_text = False
|
||||
if end_idx - start_idx <= 2:
|
||||
continue
|
||||
box_list.append((start_idx, end_idx + 1))
|
||||
|
||||
if in_text:
|
||||
box_list.append((start_idx, h - 1))
|
||||
# 绘制投影直方图
|
||||
for j in range(0, h):
|
||||
for i in range(0, project_val_array[j]):
|
||||
projection_map[j, i] = 0
|
||||
split_bbox_list = []
|
||||
if len(box_list) > 1:
|
||||
for i, (h_start, h_end) in enumerate(box_list):
|
||||
if i == 0:
|
||||
h_start = 0
|
||||
if i == len(box_list):
|
||||
h_end = h
|
||||
word_img = erosion[h_start:h_end + 1, :]
|
||||
word_h, word_w = word_img.shape
|
||||
w_split_list, w_projection_map = self.projection(word_img.T, word_w, word_h)
|
||||
w_start, w_end = w_split_list[0][0], w_split_list[-1][1]
|
||||
if h_start > 0:
|
||||
h_start -= 1
|
||||
h_end += 1
|
||||
word_img = box_img[h_start:h_end + 1:, w_start:w_end + 1, :]
|
||||
split_bbox_list.append([w_start, h_start, w_end, h_end])
|
||||
else:
|
||||
split_bbox_list.append([0, 0, w, h])
|
||||
return split_bbox_list
|
||||
|
||||
def shrink_bbox(self, bbox):
|
||||
left, top, right, bottom = bbox
|
||||
sh_h = min(max(int((bottom - top) * 0.1), 1), self.shrink_h_max)
|
||||
sh_w = min(max(int((right - left) * 0.1), 1), self.shrink_w_max)
|
||||
left_new = left + sh_w
|
||||
right_new = right - sh_w
|
||||
top_new = top + sh_h
|
||||
bottom_new = bottom - sh_h
|
||||
if left_new >= right_new:
|
||||
left_new = left
|
||||
right_new = right
|
||||
if top_new >= bottom_new:
|
||||
top_new = top
|
||||
bottom_new = bottom
|
||||
return [left_new, top_new, right_new, bottom_new]
|
||||
|
||||
def __call__(self, data):
|
||||
img = data['image']
|
||||
cells = data['cells']
|
||||
height, width = img.shape[0:2]
|
||||
if self.mask_type == 1:
|
||||
mask_img = np.zeros((height, width), dtype=np.float32)
|
||||
else:
|
||||
mask_img = np.zeros((height, width, 3), dtype=np.float32)
|
||||
cell_num = len(cells)
|
||||
for cno in range(cell_num):
|
||||
if "bbox" in cells[cno]:
|
||||
bbox = cells[cno]['bbox']
|
||||
left, top, right, bottom = bbox
|
||||
box_img = img[top:bottom, left:right, :].copy()
|
||||
split_bbox_list = self.projection_cx(box_img)
|
||||
for sno in range(len(split_bbox_list)):
|
||||
split_bbox_list[sno][0] += left
|
||||
split_bbox_list[sno][1] += top
|
||||
split_bbox_list[sno][2] += left
|
||||
split_bbox_list[sno][3] += top
|
||||
|
||||
for sno in range(len(split_bbox_list)):
|
||||
left, top, right, bottom = split_bbox_list[sno]
|
||||
left, top, right, bottom = self.shrink_bbox([left, top, right, bottom])
|
||||
if self.mask_type == 1:
|
||||
mask_img[top:bottom, left:right] = 1.0
|
||||
data['mask_img'] = mask_img
|
||||
else:
|
||||
mask_img[top:bottom, left:right, :] = (255, 255, 255)
|
||||
data['image'] = mask_img
|
||||
return data
|
||||
|
||||
class ResizeTableImage(object):
|
||||
def __init__(self, max_len, **kwargs):
|
||||
super(ResizeTableImage, self).__init__()
|
||||
self.max_len = max_len
|
||||
|
||||
def get_img_bbox(self, cells):
|
||||
bbox_list = []
|
||||
if len(cells) == 0:
|
||||
return bbox_list
|
||||
cell_num = len(cells)
|
||||
for cno in range(cell_num):
|
||||
if "bbox" in cells[cno]:
|
||||
bbox = cells[cno]['bbox']
|
||||
bbox_list.append(bbox)
|
||||
return bbox_list
|
||||
|
||||
def resize_img_table(self, img, bbox_list, max_len):
|
||||
height, width = img.shape[0:2]
|
||||
ratio = max_len / (max(height, width) * 1.0)
|
||||
resize_h = int(height * ratio)
|
||||
resize_w = int(width * ratio)
|
||||
img_new = cv2.resize(img, (resize_w, resize_h))
|
||||
bbox_list_new = []
|
||||
for bno in range(len(bbox_list)):
|
||||
left, top, right, bottom = bbox_list[bno].copy()
|
||||
left = int(left * ratio)
|
||||
top = int(top * ratio)
|
||||
right = int(right * ratio)
|
||||
bottom = int(bottom * ratio)
|
||||
bbox_list_new.append([left, top, right, bottom])
|
||||
return img_new, bbox_list_new
|
||||
|
||||
def __call__(self, data):
|
||||
img = data['image']
|
||||
if 'cells' not in data:
|
||||
cells = []
|
||||
else:
|
||||
cells = data['cells']
|
||||
bbox_list = self.get_img_bbox(cells)
|
||||
img_new, bbox_list_new = self.resize_img_table(img, bbox_list, self.max_len)
|
||||
data['image'] = img_new
|
||||
cell_num = len(cells)
|
||||
bno = 0
|
||||
for cno in range(cell_num):
|
||||
if "bbox" in data['cells'][cno]:
|
||||
data['cells'][cno]['bbox'] = bbox_list_new[bno]
|
||||
bno += 1
|
||||
data['max_len'] = self.max_len
|
||||
return data
|
||||
|
||||
class PaddingTableImage(object):
|
||||
def __init__(self, **kwargs):
|
||||
super(PaddingTableImage, self).__init__()
|
||||
|
||||
def __call__(self, data):
|
||||
img = data['image']
|
||||
max_len = data['max_len']
|
||||
padding_img = np.zeros((max_len, max_len, 3), dtype=np.float32)
|
||||
height, width = img.shape[0:2]
|
||||
padding_img[0:height, 0:width, :] = img.copy()
|
||||
data['image'] = padding_img
|
||||
return data
|
||||
|
||||
@@ -0,0 +1,105 @@
|
||||
# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""
|
||||
This code is refer from:
|
||||
https://github.com/WenmuZhou/DBNet.pytorch/blob/master/data_loader/modules/iaa_augment.py
|
||||
"""
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
from __future__ import unicode_literals
|
||||
|
||||
import numpy as np
|
||||
import imgaug
|
||||
import imgaug.augmenters as iaa
|
||||
|
||||
|
||||
class AugmenterBuilder(object):
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def build(self, args, root=True):
|
||||
if args is None or len(args) == 0:
|
||||
return None
|
||||
elif isinstance(args, list):
|
||||
if root:
|
||||
sequence = [self.build(value, root=False) for value in args]
|
||||
return iaa.Sequential(sequence)
|
||||
else:
|
||||
return getattr(iaa, args[0])(
|
||||
*[self.to_tuple_if_list(a) for a in args[1:]])
|
||||
elif isinstance(args, dict):
|
||||
cls = getattr(iaa, args['type'])
|
||||
return cls(**{
|
||||
k: self.to_tuple_if_list(v)
|
||||
for k, v in args['args'].items()
|
||||
})
|
||||
else:
|
||||
raise RuntimeError('unknown augmenter arg: ' + str(args))
|
||||
|
||||
def to_tuple_if_list(self, obj):
|
||||
if isinstance(obj, list):
|
||||
return tuple(obj)
|
||||
return obj
|
||||
|
||||
|
||||
class IaaAugment():
|
||||
def __init__(self, augmenter_args=None, **kwargs):
|
||||
if augmenter_args is None:
|
||||
augmenter_args = [{
|
||||
'type': 'Fliplr',
|
||||
'args': {
|
||||
'p': 0.5
|
||||
}
|
||||
}, {
|
||||
'type': 'Affine',
|
||||
'args': {
|
||||
'rotate': [-10, 10]
|
||||
}
|
||||
}, {
|
||||
'type': 'Resize',
|
||||
'args': {
|
||||
'size': [0.5, 3]
|
||||
}
|
||||
}]
|
||||
self.augmenter = AugmenterBuilder().build(augmenter_args)
|
||||
|
||||
def __call__(self, data):
|
||||
image = data['image']
|
||||
shape = image.shape
|
||||
|
||||
if self.augmenter:
|
||||
aug = self.augmenter.to_deterministic()
|
||||
data['image'] = aug.augment_image(image)
|
||||
data = self.may_augment_annotation(aug, data, shape)
|
||||
return data
|
||||
|
||||
def may_augment_annotation(self, aug, data, shape):
|
||||
if aug is None:
|
||||
return data
|
||||
|
||||
line_polys = []
|
||||
for poly in data['polys']:
|
||||
new_poly = self.may_augment_poly(aug, shape, poly)
|
||||
line_polys.append(new_poly)
|
||||
data['polys'] = np.array(line_polys)
|
||||
return data
|
||||
|
||||
def may_augment_poly(self, aug, img_shape, poly):
|
||||
keypoints = [imgaug.Keypoint(p[0], p[1]) for p in poly]
|
||||
keypoints = aug.augment_keypoints(
|
||||
[imgaug.KeypointsOnImage(
|
||||
keypoints, shape=img_shape)])[0].keypoints
|
||||
poly = [(p.x, p.y) for p in keypoints]
|
||||
return poly
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,173 @@
|
||||
# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""
|
||||
This code is refer from:
|
||||
https://github.com/WenmuZhou/DBNet.pytorch/blob/master/data_loader/modules/make_border_map.py
|
||||
"""
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
from __future__ import unicode_literals
|
||||
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
np.seterr(divide='ignore', invalid='ignore')
|
||||
import pyclipper
|
||||
from shapely.geometry import Polygon
|
||||
import sys
|
||||
import warnings
|
||||
|
||||
warnings.simplefilter("ignore")
|
||||
|
||||
__all__ = ['MakeBorderMap']
|
||||
|
||||
|
||||
class MakeBorderMap(object):
|
||||
def __init__(self,
|
||||
shrink_ratio=0.4,
|
||||
thresh_min=0.3,
|
||||
thresh_max=0.7,
|
||||
**kwargs):
|
||||
self.shrink_ratio = shrink_ratio
|
||||
self.thresh_min = thresh_min
|
||||
self.thresh_max = thresh_max
|
||||
|
||||
def __call__(self, data):
|
||||
|
||||
img = data['image']
|
||||
text_polys = data['polys']
|
||||
ignore_tags = data['ignore_tags']
|
||||
|
||||
canvas = np.zeros(img.shape[:2], dtype=np.float32)
|
||||
mask = np.zeros(img.shape[:2], dtype=np.float32)
|
||||
|
||||
for i in range(len(text_polys)):
|
||||
if ignore_tags[i]:
|
||||
continue
|
||||
self.draw_border_map(text_polys[i], canvas, mask=mask)
|
||||
canvas = canvas * (self.thresh_max - self.thresh_min) + self.thresh_min
|
||||
|
||||
data['threshold_map'] = canvas
|
||||
data['threshold_mask'] = mask
|
||||
return data
|
||||
|
||||
def draw_border_map(self, polygon, canvas, mask):
|
||||
polygon = np.array(polygon)
|
||||
assert polygon.ndim == 2
|
||||
assert polygon.shape[1] == 2
|
||||
|
||||
polygon_shape = Polygon(polygon)
|
||||
if polygon_shape.area <= 0:
|
||||
return
|
||||
distance = polygon_shape.area * (
|
||||
1 - np.power(self.shrink_ratio, 2)) / polygon_shape.length
|
||||
subject = [tuple(l) for l in polygon]
|
||||
padding = pyclipper.PyclipperOffset()
|
||||
padding.AddPath(subject, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)
|
||||
|
||||
padded_polygon = np.array(padding.Execute(distance)[0])
|
||||
cv2.fillPoly(mask, [padded_polygon.astype(np.int32)], 1.0)
|
||||
|
||||
xmin = padded_polygon[:, 0].min()
|
||||
xmax = padded_polygon[:, 0].max()
|
||||
ymin = padded_polygon[:, 1].min()
|
||||
ymax = padded_polygon[:, 1].max()
|
||||
width = xmax - xmin + 1
|
||||
height = ymax - ymin + 1
|
||||
|
||||
polygon[:, 0] = polygon[:, 0] - xmin
|
||||
polygon[:, 1] = polygon[:, 1] - ymin
|
||||
|
||||
xs = np.broadcast_to(
|
||||
np.linspace(
|
||||
0, width - 1, num=width).reshape(1, width), (height, width))
|
||||
ys = np.broadcast_to(
|
||||
np.linspace(
|
||||
0, height - 1, num=height).reshape(height, 1), (height, width))
|
||||
|
||||
distance_map = np.zeros(
|
||||
(polygon.shape[0], height, width), dtype=np.float32)
|
||||
for i in range(polygon.shape[0]):
|
||||
j = (i + 1) % polygon.shape[0]
|
||||
absolute_distance = self._distance(xs, ys, polygon[i], polygon[j])
|
||||
distance_map[i] = np.clip(absolute_distance / distance, 0, 1)
|
||||
distance_map = distance_map.min(axis=0)
|
||||
|
||||
xmin_valid = min(max(0, xmin), canvas.shape[1] - 1)
|
||||
xmax_valid = min(max(0, xmax), canvas.shape[1] - 1)
|
||||
ymin_valid = min(max(0, ymin), canvas.shape[0] - 1)
|
||||
ymax_valid = min(max(0, ymax), canvas.shape[0] - 1)
|
||||
canvas[ymin_valid:ymax_valid + 1, xmin_valid:xmax_valid + 1] = np.fmax(
|
||||
1 - distance_map[ymin_valid - ymin:ymax_valid - ymax + height,
|
||||
xmin_valid - xmin:xmax_valid - xmax + width],
|
||||
canvas[ymin_valid:ymax_valid + 1, xmin_valid:xmax_valid + 1])
|
||||
|
||||
def _distance(self, xs, ys, point_1, point_2):
|
||||
'''
|
||||
compute the distance from point to a line
|
||||
ys: coordinates in the first axis
|
||||
xs: coordinates in the second axis
|
||||
point_1, point_2: (x, y), the end of the line
|
||||
'''
|
||||
height, width = xs.shape[:2]
|
||||
square_distance_1 = np.square(xs - point_1[0]) + np.square(ys - point_1[
|
||||
1])
|
||||
square_distance_2 = np.square(xs - point_2[0]) + np.square(ys - point_2[
|
||||
1])
|
||||
square_distance = np.square(point_1[0] - point_2[0]) + np.square(
|
||||
point_1[1] - point_2[1])
|
||||
|
||||
cosin = (square_distance - square_distance_1 - square_distance_2) / (
|
||||
2 * np.sqrt(square_distance_1 * square_distance_2))
|
||||
square_sin = 1 - np.square(cosin)
|
||||
square_sin = np.nan_to_num(square_sin)
|
||||
result = np.sqrt(square_distance_1 * square_distance_2 * square_sin /
|
||||
square_distance)
|
||||
|
||||
result[cosin <
|
||||
0] = np.sqrt(np.fmin(square_distance_1, square_distance_2))[cosin
|
||||
< 0]
|
||||
# self.extend_line(point_1, point_2, result)
|
||||
return result
|
||||
|
||||
def extend_line(self, point_1, point_2, result, shrink_ratio):
|
||||
ex_point_1 = (int(
|
||||
round(point_1[0] + (point_1[0] - point_2[0]) * (1 + shrink_ratio))),
|
||||
int(
|
||||
round(point_1[1] + (point_1[1] - point_2[1]) * (
|
||||
1 + shrink_ratio))))
|
||||
cv2.line(
|
||||
result,
|
||||
tuple(ex_point_1),
|
||||
tuple(point_1),
|
||||
4096.0,
|
||||
1,
|
||||
lineType=cv2.LINE_AA,
|
||||
shift=0)
|
||||
ex_point_2 = (int(
|
||||
round(point_2[0] + (point_2[0] - point_1[0]) * (1 + shrink_ratio))),
|
||||
int(
|
||||
round(point_2[1] + (point_2[1] - point_1[1]) * (
|
||||
1 + shrink_ratio))))
|
||||
cv2.line(
|
||||
result,
|
||||
tuple(ex_point_2),
|
||||
tuple(point_2),
|
||||
4096.0,
|
||||
1,
|
||||
lineType=cv2.LINE_AA,
|
||||
shift=0)
|
||||
return ex_point_1, ex_point_2
|
||||
@@ -0,0 +1,106 @@
|
||||
# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
from __future__ import unicode_literals
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import pyclipper
|
||||
from shapely.geometry import Polygon
|
||||
|
||||
__all__ = ['MakePseGt']
|
||||
|
||||
|
||||
class MakePseGt(object):
|
||||
def __init__(self, kernel_num=7, size=640, min_shrink_ratio=0.4, **kwargs):
|
||||
self.kernel_num = kernel_num
|
||||
self.min_shrink_ratio = min_shrink_ratio
|
||||
self.size = size
|
||||
|
||||
def __call__(self, data):
|
||||
|
||||
image = data['image']
|
||||
text_polys = data['polys']
|
||||
ignore_tags = data['ignore_tags']
|
||||
|
||||
h, w, _ = image.shape
|
||||
short_edge = min(h, w)
|
||||
if short_edge < self.size:
|
||||
# keep short_size >= self.size
|
||||
scale = self.size / short_edge
|
||||
image = cv2.resize(image, dsize=None, fx=scale, fy=scale)
|
||||
text_polys *= scale
|
||||
|
||||
gt_kernels = []
|
||||
for i in range(1, self.kernel_num + 1):
|
||||
# s1->sn, from big to small
|
||||
rate = 1.0 - (1.0 - self.min_shrink_ratio) / (self.kernel_num - 1
|
||||
) * i
|
||||
text_kernel, ignore_tags = self.generate_kernel(
|
||||
image.shape[0:2], rate, text_polys, ignore_tags)
|
||||
gt_kernels.append(text_kernel)
|
||||
|
||||
training_mask = np.ones(image.shape[0:2], dtype='uint8')
|
||||
for i in range(text_polys.shape[0]):
|
||||
if ignore_tags[i]:
|
||||
cv2.fillPoly(training_mask,
|
||||
text_polys[i].astype(np.int32)[np.newaxis, :, :],
|
||||
0)
|
||||
|
||||
gt_kernels = np.array(gt_kernels)
|
||||
gt_kernels[gt_kernels > 0] = 1
|
||||
|
||||
data['image'] = image
|
||||
data['polys'] = text_polys
|
||||
data['gt_kernels'] = gt_kernels[0:]
|
||||
data['gt_text'] = gt_kernels[0]
|
||||
data['mask'] = training_mask.astype('float32')
|
||||
return data
|
||||
|
||||
def generate_kernel(self,
|
||||
img_size,
|
||||
shrink_ratio,
|
||||
text_polys,
|
||||
ignore_tags=None):
|
||||
"""
|
||||
Refer to part of the code:
|
||||
https://github.com/open-mmlab/mmocr/blob/main/mmocr/datasets/pipelines/textdet_targets/base_textdet_targets.py
|
||||
"""
|
||||
|
||||
h, w = img_size
|
||||
text_kernel = np.zeros((h, w), dtype=np.float32)
|
||||
for i, poly in enumerate(text_polys):
|
||||
polygon = Polygon(poly)
|
||||
distance = polygon.area * (1 - shrink_ratio * shrink_ratio) / (
|
||||
polygon.length + 1e-6)
|
||||
subject = [tuple(l) for l in poly]
|
||||
pco = pyclipper.PyclipperOffset()
|
||||
pco.AddPath(subject, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)
|
||||
shrinked = np.array(pco.Execute(-distance))
|
||||
|
||||
if len(shrinked) == 0 or shrinked.size == 0:
|
||||
if ignore_tags is not None:
|
||||
ignore_tags[i] = True
|
||||
continue
|
||||
try:
|
||||
shrinked = np.array(shrinked[0]).reshape(-1, 2)
|
||||
except:
|
||||
if ignore_tags is not None:
|
||||
ignore_tags[i] = True
|
||||
continue
|
||||
cv2.fillPoly(text_kernel, [shrinked.astype(np.int32)], i + 1)
|
||||
return text_kernel, ignore_tags
|
||||
@@ -0,0 +1,123 @@
|
||||
# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""
|
||||
This code is refer from:
|
||||
https://github.com/WenmuZhou/DBNet.pytorch/blob/master/data_loader/modules/make_shrink_map.py
|
||||
"""
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
from __future__ import unicode_literals
|
||||
|
||||
import numpy as np
|
||||
import cv2
|
||||
from shapely.geometry import Polygon
|
||||
import pyclipper
|
||||
|
||||
__all__ = ['MakeShrinkMap']
|
||||
|
||||
|
||||
class MakeShrinkMap(object):
|
||||
r'''
|
||||
Making binary mask from detection data with ICDAR format.
|
||||
Typically following the process of class `MakeICDARData`.
|
||||
'''
|
||||
|
||||
def __init__(self, min_text_size=8, shrink_ratio=0.4, **kwargs):
|
||||
self.min_text_size = min_text_size
|
||||
self.shrink_ratio = shrink_ratio
|
||||
|
||||
def __call__(self, data):
|
||||
image = data['image']
|
||||
text_polys = data['polys']
|
||||
ignore_tags = data['ignore_tags']
|
||||
|
||||
h, w = image.shape[:2]
|
||||
text_polys, ignore_tags = self.validate_polygons(text_polys,
|
||||
ignore_tags, h, w)
|
||||
gt = np.zeros((h, w), dtype=np.float32)
|
||||
mask = np.ones((h, w), dtype=np.float32)
|
||||
for i in range(len(text_polys)):
|
||||
polygon = text_polys[i]
|
||||
height = max(polygon[:, 1]) - min(polygon[:, 1])
|
||||
width = max(polygon[:, 0]) - min(polygon[:, 0])
|
||||
if ignore_tags[i] or min(height, width) < self.min_text_size:
|
||||
cv2.fillPoly(mask,
|
||||
polygon.astype(np.int32)[np.newaxis, :, :], 0)
|
||||
ignore_tags[i] = True
|
||||
else:
|
||||
polygon_shape = Polygon(polygon)
|
||||
subject = [tuple(l) for l in polygon]
|
||||
padding = pyclipper.PyclipperOffset()
|
||||
padding.AddPath(subject, pyclipper.JT_ROUND,
|
||||
pyclipper.ET_CLOSEDPOLYGON)
|
||||
shrinked = []
|
||||
|
||||
# Increase the shrink ratio every time we get multiple polygon returned back
|
||||
possible_ratios = np.arange(self.shrink_ratio, 1,
|
||||
self.shrink_ratio)
|
||||
np.append(possible_ratios, 1)
|
||||
# print(possible_ratios)
|
||||
for ratio in possible_ratios:
|
||||
# print(f"Change shrink ratio to {ratio}")
|
||||
distance = polygon_shape.area * (
|
||||
1 - np.power(ratio, 2)) / polygon_shape.length
|
||||
shrinked = padding.Execute(-distance)
|
||||
if len(shrinked) == 1:
|
||||
break
|
||||
|
||||
if shrinked == []:
|
||||
cv2.fillPoly(mask,
|
||||
polygon.astype(np.int32)[np.newaxis, :, :], 0)
|
||||
ignore_tags[i] = True
|
||||
continue
|
||||
|
||||
for each_shirnk in shrinked:
|
||||
shirnk = np.array(each_shirnk).reshape(-1, 2)
|
||||
cv2.fillPoly(gt, [shirnk.astype(np.int32)], 1)
|
||||
|
||||
data['shrink_map'] = gt
|
||||
data['shrink_mask'] = mask
|
||||
return data
|
||||
|
||||
def validate_polygons(self, polygons, ignore_tags, h, w):
|
||||
'''
|
||||
polygons (numpy.array, required): of shape (num_instances, num_points, 2)
|
||||
'''
|
||||
if len(polygons) == 0:
|
||||
return polygons, ignore_tags
|
||||
assert len(polygons) == len(ignore_tags)
|
||||
for polygon in polygons:
|
||||
polygon[:, 0] = np.clip(polygon[:, 0], 0, w - 1)
|
||||
polygon[:, 1] = np.clip(polygon[:, 1], 0, h - 1)
|
||||
|
||||
for i in range(len(polygons)):
|
||||
area = self.polygon_area(polygons[i])
|
||||
if abs(area) < 1:
|
||||
ignore_tags[i] = True
|
||||
if area > 0:
|
||||
polygons[i] = polygons[i][::-1, :]
|
||||
return polygons, ignore_tags
|
||||
|
||||
def polygon_area(self, polygon):
|
||||
"""
|
||||
compute polygon area
|
||||
"""
|
||||
area = 0
|
||||
q = polygon[-1]
|
||||
for p in polygon:
|
||||
area += p[0] * q[1] - p[1] * q[0]
|
||||
q = p
|
||||
return area / 2.0
|
||||
@@ -0,0 +1,468 @@
|
||||
"""
|
||||
# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
from __future__ import unicode_literals
|
||||
|
||||
import sys
|
||||
import six
|
||||
import cv2
|
||||
import numpy as np
|
||||
import math
|
||||
|
||||
|
||||
class DecodeImage(object):
|
||||
""" decode image """
|
||||
|
||||
def __init__(self,
|
||||
img_mode='RGB',
|
||||
channel_first=False,
|
||||
ignore_orientation=False,
|
||||
**kwargs):
|
||||
self.img_mode = img_mode
|
||||
self.channel_first = channel_first
|
||||
self.ignore_orientation = ignore_orientation
|
||||
|
||||
def __call__(self, data):
|
||||
img = data['image']
|
||||
if six.PY2:
|
||||
assert type(img) is str and len(
|
||||
img) > 0, "invalid input 'img' in DecodeImage"
|
||||
else:
|
||||
assert type(img) is bytes and len(
|
||||
img) > 0, "invalid input 'img' in DecodeImage"
|
||||
img = np.frombuffer(img, dtype='uint8')
|
||||
if self.ignore_orientation:
|
||||
img = cv2.imdecode(img, cv2.IMREAD_IGNORE_ORIENTATION |
|
||||
cv2.IMREAD_COLOR)
|
||||
else:
|
||||
img = cv2.imdecode(img, 1)
|
||||
if img is None:
|
||||
return None
|
||||
if self.img_mode == 'GRAY':
|
||||
img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
|
||||
elif self.img_mode == 'RGB':
|
||||
assert img.shape[2] == 3, 'invalid shape of image[%s]' % (img.shape)
|
||||
img = img[:, :, ::-1]
|
||||
|
||||
if self.channel_first:
|
||||
img = img.transpose((2, 0, 1))
|
||||
|
||||
data['image'] = img
|
||||
return data
|
||||
|
||||
|
||||
class NRTRDecodeImage(object):
|
||||
""" decode image """
|
||||
|
||||
def __init__(self, img_mode='RGB', channel_first=False, **kwargs):
|
||||
self.img_mode = img_mode
|
||||
self.channel_first = channel_first
|
||||
|
||||
def __call__(self, data):
|
||||
img = data['image']
|
||||
if six.PY2:
|
||||
assert type(img) is str and len(
|
||||
img) > 0, "invalid input 'img' in DecodeImage"
|
||||
else:
|
||||
assert type(img) is bytes and len(
|
||||
img) > 0, "invalid input 'img' in DecodeImage"
|
||||
img = np.frombuffer(img, dtype='uint8')
|
||||
|
||||
img = cv2.imdecode(img, 1)
|
||||
|
||||
if img is None:
|
||||
return None
|
||||
if self.img_mode == 'GRAY':
|
||||
img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
|
||||
elif self.img_mode == 'RGB':
|
||||
assert img.shape[2] == 3, 'invalid shape of image[%s]' % (img.shape)
|
||||
img = img[:, :, ::-1]
|
||||
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
||||
if self.channel_first:
|
||||
img = img.transpose((2, 0, 1))
|
||||
data['image'] = img
|
||||
return data
|
||||
|
||||
|
||||
class NormalizeImage(object):
|
||||
""" normalize image such as substract mean, divide std
|
||||
"""
|
||||
|
||||
def __init__(self, scale=None, mean=None, std=None, order='chw', **kwargs):
|
||||
if isinstance(scale, str):
|
||||
scale = eval(scale)
|
||||
self.scale = np.float32(scale if scale is not None else 1.0 / 255.0)
|
||||
mean = mean if mean is not None else [0.485, 0.456, 0.406]
|
||||
std = std if std is not None else [0.229, 0.224, 0.225]
|
||||
|
||||
shape = (3, 1, 1) if order == 'chw' else (1, 1, 3)
|
||||
self.mean = np.array(mean).reshape(shape).astype('float32')
|
||||
self.std = np.array(std).reshape(shape).astype('float32')
|
||||
|
||||
def __call__(self, data):
|
||||
img = data['image']
|
||||
from PIL import Image
|
||||
if isinstance(img, Image.Image):
|
||||
img = np.array(img)
|
||||
assert isinstance(img,
|
||||
np.ndarray), "invalid input 'img' in NormalizeImage"
|
||||
data['image'] = (
|
||||
img.astype('float32') * self.scale - self.mean) / self.std
|
||||
return data
|
||||
|
||||
|
||||
class ToCHWImage(object):
|
||||
""" convert hwc image to chw image
|
||||
"""
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
pass
|
||||
|
||||
def __call__(self, data):
|
||||
img = data['image']
|
||||
from PIL import Image
|
||||
if isinstance(img, Image.Image):
|
||||
img = np.array(img)
|
||||
data['image'] = img.transpose((2, 0, 1))
|
||||
return data
|
||||
|
||||
|
||||
class Fasttext(object):
|
||||
def __init__(self, path="None", **kwargs):
|
||||
import fasttext
|
||||
self.fast_model = fasttext.load_model(path)
|
||||
|
||||
def __call__(self, data):
|
||||
label = data['label']
|
||||
fast_label = self.fast_model[label]
|
||||
data['fast_label'] = fast_label
|
||||
return data
|
||||
|
||||
|
||||
class KeepKeys(object):
|
||||
def __init__(self, keep_keys, **kwargs):
|
||||
self.keep_keys = keep_keys
|
||||
|
||||
def __call__(self, data):
|
||||
data_list = []
|
||||
for key in self.keep_keys:
|
||||
data_list.append(data[key])
|
||||
return data_list
|
||||
|
||||
|
||||
class Pad(object):
|
||||
def __init__(self, size=None, size_div=32, **kwargs):
|
||||
if size is not None and not isinstance(size, (int, list, tuple)):
|
||||
raise TypeError("Type of target_size is invalid. Now is {}".format(
|
||||
type(size)))
|
||||
if isinstance(size, int):
|
||||
size = [size, size]
|
||||
self.size = size
|
||||
self.size_div = size_div
|
||||
|
||||
def __call__(self, data):
|
||||
|
||||
img = data['image']
|
||||
img_h, img_w = img.shape[0], img.shape[1]
|
||||
if self.size:
|
||||
resize_h2, resize_w2 = self.size
|
||||
assert (
|
||||
img_h < resize_h2 and img_w < resize_w2
|
||||
), '(h, w) of target size should be greater than (img_h, img_w)'
|
||||
else:
|
||||
resize_h2 = max(
|
||||
int(math.ceil(img.shape[0] / self.size_div) * self.size_div),
|
||||
self.size_div)
|
||||
resize_w2 = max(
|
||||
int(math.ceil(img.shape[1] / self.size_div) * self.size_div),
|
||||
self.size_div)
|
||||
img = cv2.copyMakeBorder(
|
||||
img,
|
||||
0,
|
||||
resize_h2 - img_h,
|
||||
0,
|
||||
resize_w2 - img_w,
|
||||
cv2.BORDER_CONSTANT,
|
||||
value=0)
|
||||
data['image'] = img
|
||||
return data
|
||||
|
||||
|
||||
class Resize(object):
|
||||
def __init__(self, size=(640, 640), **kwargs):
|
||||
self.size = size
|
||||
|
||||
def resize_image(self, img):
|
||||
resize_h, resize_w = self.size
|
||||
ori_h, ori_w = img.shape[:2] # (h, w, c)
|
||||
ratio_h = float(resize_h) / ori_h
|
||||
ratio_w = float(resize_w) / ori_w
|
||||
img = cv2.resize(img, (int(resize_w), int(resize_h)))
|
||||
return img, [ratio_h, ratio_w]
|
||||
|
||||
def __call__(self, data):
|
||||
img = data['image']
|
||||
if 'polys' in data:
|
||||
text_polys = data['polys']
|
||||
|
||||
img_resize, [ratio_h, ratio_w] = self.resize_image(img)
|
||||
if 'polys' in data:
|
||||
new_boxes = []
|
||||
for box in text_polys:
|
||||
new_box = []
|
||||
for cord in box:
|
||||
new_box.append([cord[0] * ratio_w, cord[1] * ratio_h])
|
||||
new_boxes.append(new_box)
|
||||
data['polys'] = np.array(new_boxes, dtype=np.float32)
|
||||
data['image'] = img_resize
|
||||
return data
|
||||
|
||||
|
||||
class DetResizeForTest(object):
|
||||
def __init__(self, **kwargs):
|
||||
super(DetResizeForTest, self).__init__()
|
||||
self.resize_type = 0
|
||||
if 'image_shape' in kwargs:
|
||||
self.image_shape = kwargs['image_shape']
|
||||
self.resize_type = 1
|
||||
elif 'limit_side_len' in kwargs:
|
||||
self.limit_side_len = kwargs['limit_side_len']
|
||||
self.limit_type = kwargs.get('limit_type', 'min')
|
||||
elif 'resize_long' in kwargs:
|
||||
self.resize_type = 2
|
||||
self.resize_long = kwargs.get('resize_long', 960)
|
||||
else:
|
||||
self.limit_side_len = 736
|
||||
self.limit_type = 'min'
|
||||
|
||||
def __call__(self, data):
|
||||
img = data['image']
|
||||
src_h, src_w, _ = img.shape
|
||||
|
||||
if self.resize_type == 0:
|
||||
# img, shape = self.resize_image_type0(img)
|
||||
img, [ratio_h, ratio_w] = self.resize_image_type0(img)
|
||||
elif self.resize_type == 2:
|
||||
img, [ratio_h, ratio_w] = self.resize_image_type2(img)
|
||||
else:
|
||||
# img, shape = self.resize_image_type1(img)
|
||||
img, [ratio_h, ratio_w] = self.resize_image_type1(img)
|
||||
data['image'] = img
|
||||
data['shape'] = np.array([src_h, src_w, ratio_h, ratio_w])
|
||||
return data
|
||||
|
||||
def resize_image_type1(self, img):
|
||||
resize_h, resize_w = self.image_shape
|
||||
ori_h, ori_w = img.shape[:2] # (h, w, c)
|
||||
ratio_h = float(resize_h) / ori_h
|
||||
ratio_w = float(resize_w) / ori_w
|
||||
img = cv2.resize(img, (int(resize_w), int(resize_h)))
|
||||
# return img, np.array([ori_h, ori_w])
|
||||
return img, [ratio_h, ratio_w]
|
||||
|
||||
def resize_image_type0(self, img):
|
||||
"""
|
||||
resize image to a size multiple of 32 which is required by the network
|
||||
args:
|
||||
img(array): array with shape [h, w, c]
|
||||
return(tuple):
|
||||
img, (ratio_h, ratio_w)
|
||||
"""
|
||||
limit_side_len = self.limit_side_len
|
||||
h, w, c = img.shape
|
||||
|
||||
# limit the max side
|
||||
if self.limit_type == 'max':
|
||||
if max(h, w) > limit_side_len:
|
||||
if h > w:
|
||||
ratio = float(limit_side_len) / h
|
||||
else:
|
||||
ratio = float(limit_side_len) / w
|
||||
else:
|
||||
ratio = 1.
|
||||
elif self.limit_type == 'min':
|
||||
if min(h, w) < limit_side_len:
|
||||
if h < w:
|
||||
ratio = float(limit_side_len) / h
|
||||
else:
|
||||
ratio = float(limit_side_len) / w
|
||||
else:
|
||||
ratio = 1.
|
||||
elif self.limit_type == 'resize_long':
|
||||
ratio = float(limit_side_len) / max(h, w)
|
||||
else:
|
||||
raise Exception('not support limit type, image ')
|
||||
resize_h = int(h * ratio)
|
||||
resize_w = int(w * ratio)
|
||||
|
||||
resize_h = max(int(round(resize_h / 32) * 32), 32)
|
||||
resize_w = max(int(round(resize_w / 32) * 32), 32)
|
||||
|
||||
try:
|
||||
if int(resize_w) <= 0 or int(resize_h) <= 0:
|
||||
return None, (None, None)
|
||||
img = cv2.resize(img, (int(resize_w), int(resize_h)))
|
||||
except:
|
||||
print(img.shape, resize_w, resize_h)
|
||||
sys.exit(0)
|
||||
ratio_h = resize_h / float(h)
|
||||
ratio_w = resize_w / float(w)
|
||||
return img, [ratio_h, ratio_w]
|
||||
|
||||
def resize_image_type2(self, img):
|
||||
h, w, _ = img.shape
|
||||
|
||||
resize_w = w
|
||||
resize_h = h
|
||||
|
||||
if resize_h > resize_w:
|
||||
ratio = float(self.resize_long) / resize_h
|
||||
else:
|
||||
ratio = float(self.resize_long) / resize_w
|
||||
|
||||
resize_h = int(resize_h * ratio)
|
||||
resize_w = int(resize_w * ratio)
|
||||
|
||||
max_stride = 128
|
||||
resize_h = (resize_h + max_stride - 1) // max_stride * max_stride
|
||||
resize_w = (resize_w + max_stride - 1) // max_stride * max_stride
|
||||
img = cv2.resize(img, (int(resize_w), int(resize_h)))
|
||||
ratio_h = resize_h / float(h)
|
||||
ratio_w = resize_w / float(w)
|
||||
|
||||
return img, [ratio_h, ratio_w]
|
||||
|
||||
|
||||
class E2EResizeForTest(object):
|
||||
def __init__(self, **kwargs):
|
||||
super(E2EResizeForTest, self).__init__()
|
||||
self.max_side_len = kwargs['max_side_len']
|
||||
self.valid_set = kwargs['valid_set']
|
||||
|
||||
def __call__(self, data):
|
||||
img = data['image']
|
||||
src_h, src_w, _ = img.shape
|
||||
if self.valid_set == 'totaltext':
|
||||
im_resized, [ratio_h, ratio_w] = self.resize_image_for_totaltext(
|
||||
img, max_side_len=self.max_side_len)
|
||||
else:
|
||||
im_resized, (ratio_h, ratio_w) = self.resize_image(
|
||||
img, max_side_len=self.max_side_len)
|
||||
data['image'] = im_resized
|
||||
data['shape'] = np.array([src_h, src_w, ratio_h, ratio_w])
|
||||
return data
|
||||
|
||||
def resize_image_for_totaltext(self, im, max_side_len=512):
|
||||
|
||||
h, w, _ = im.shape
|
||||
resize_w = w
|
||||
resize_h = h
|
||||
ratio = 1.25
|
||||
if h * ratio > max_side_len:
|
||||
ratio = float(max_side_len) / resize_h
|
||||
resize_h = int(resize_h * ratio)
|
||||
resize_w = int(resize_w * ratio)
|
||||
|
||||
max_stride = 128
|
||||
resize_h = (resize_h + max_stride - 1) // max_stride * max_stride
|
||||
resize_w = (resize_w + max_stride - 1) // max_stride * max_stride
|
||||
im = cv2.resize(im, (int(resize_w), int(resize_h)))
|
||||
ratio_h = resize_h / float(h)
|
||||
ratio_w = resize_w / float(w)
|
||||
return im, (ratio_h, ratio_w)
|
||||
|
||||
def resize_image(self, im, max_side_len=512):
|
||||
"""
|
||||
resize image to a size multiple of max_stride which is required by the network
|
||||
:param im: the resized image
|
||||
:param max_side_len: limit of max image size to avoid out of memory in gpu
|
||||
:return: the resized image and the resize ratio
|
||||
"""
|
||||
h, w, _ = im.shape
|
||||
|
||||
resize_w = w
|
||||
resize_h = h
|
||||
|
||||
# Fix the longer side
|
||||
if resize_h > resize_w:
|
||||
ratio = float(max_side_len) / resize_h
|
||||
else:
|
||||
ratio = float(max_side_len) / resize_w
|
||||
|
||||
resize_h = int(resize_h * ratio)
|
||||
resize_w = int(resize_w * ratio)
|
||||
|
||||
max_stride = 128
|
||||
resize_h = (resize_h + max_stride - 1) // max_stride * max_stride
|
||||
resize_w = (resize_w + max_stride - 1) // max_stride * max_stride
|
||||
im = cv2.resize(im, (int(resize_w), int(resize_h)))
|
||||
ratio_h = resize_h / float(h)
|
||||
ratio_w = resize_w / float(w)
|
||||
|
||||
return im, (ratio_h, ratio_w)
|
||||
|
||||
|
||||
class KieResize(object):
|
||||
def __init__(self, **kwargs):
|
||||
super(KieResize, self).__init__()
|
||||
self.max_side, self.min_side = kwargs['img_scale'][0], kwargs[
|
||||
'img_scale'][1]
|
||||
|
||||
def __call__(self, data):
|
||||
img = data['image']
|
||||
points = data['points']
|
||||
src_h, src_w, _ = img.shape
|
||||
im_resized, scale_factor, [ratio_h, ratio_w
|
||||
], [new_h, new_w] = self.resize_image(img)
|
||||
resize_points = self.resize_boxes(img, points, scale_factor)
|
||||
data['ori_image'] = img
|
||||
data['ori_boxes'] = points
|
||||
data['points'] = resize_points
|
||||
data['image'] = im_resized
|
||||
data['shape'] = np.array([new_h, new_w])
|
||||
return data
|
||||
|
||||
def resize_image(self, img):
|
||||
norm_img = np.zeros([1024, 1024, 3], dtype='float32')
|
||||
scale = [512, 1024]
|
||||
h, w = img.shape[:2]
|
||||
max_long_edge = max(scale)
|
||||
max_short_edge = min(scale)
|
||||
scale_factor = min(max_long_edge / max(h, w),
|
||||
max_short_edge / min(h, w))
|
||||
resize_w, resize_h = int(w * float(scale_factor) + 0.5), int(h * float(
|
||||
scale_factor) + 0.5)
|
||||
max_stride = 32
|
||||
resize_h = (resize_h + max_stride - 1) // max_stride * max_stride
|
||||
resize_w = (resize_w + max_stride - 1) // max_stride * max_stride
|
||||
im = cv2.resize(img, (resize_w, resize_h))
|
||||
new_h, new_w = im.shape[:2]
|
||||
w_scale = new_w / w
|
||||
h_scale = new_h / h
|
||||
scale_factor = np.array(
|
||||
[w_scale, h_scale, w_scale, h_scale], dtype=np.float32)
|
||||
norm_img[:new_h, :new_w, :] = im
|
||||
return norm_img, scale_factor, [h_scale, w_scale], [new_h, new_w]
|
||||
|
||||
def resize_boxes(self, im, points, scale_factor):
|
||||
points = points * scale_factor
|
||||
img_shape = im.shape[:2]
|
||||
points[:, 0::2] = np.clip(points[:, 0::2], 0, img_shape[1])
|
||||
points[:, 1::2] = np.clip(points[:, 1::2], 0, img_shape[0])
|
||||
return points
|
||||
@@ -0,0 +1,906 @@
|
||||
# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import math
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
__all__ = ['PGProcessTrain']
|
||||
|
||||
|
||||
class PGProcessTrain(object):
|
||||
def __init__(self,
|
||||
character_dict_path,
|
||||
max_text_length,
|
||||
max_text_nums,
|
||||
tcl_len,
|
||||
batch_size=14,
|
||||
min_crop_size=24,
|
||||
min_text_size=4,
|
||||
max_text_size=512,
|
||||
**kwargs):
|
||||
self.tcl_len = tcl_len
|
||||
self.max_text_length = max_text_length
|
||||
self.max_text_nums = max_text_nums
|
||||
self.batch_size = batch_size
|
||||
self.min_crop_size = min_crop_size
|
||||
self.min_text_size = min_text_size
|
||||
self.max_text_size = max_text_size
|
||||
self.Lexicon_Table = self.get_dict(character_dict_path)
|
||||
self.pad_num = len(self.Lexicon_Table)
|
||||
self.img_id = 0
|
||||
|
||||
def get_dict(self, character_dict_path):
|
||||
character_str = ""
|
||||
with open(character_dict_path, "rb") as fin:
|
||||
lines = fin.readlines()
|
||||
for line in lines:
|
||||
line = line.decode('utf-8').strip("\n").strip("\r\n")
|
||||
character_str += line
|
||||
dict_character = list(character_str)
|
||||
return dict_character
|
||||
|
||||
def quad_area(self, poly):
|
||||
"""
|
||||
compute area of a polygon
|
||||
:param poly:
|
||||
:return:
|
||||
"""
|
||||
edge = [(poly[1][0] - poly[0][0]) * (poly[1][1] + poly[0][1]),
|
||||
(poly[2][0] - poly[1][0]) * (poly[2][1] + poly[1][1]),
|
||||
(poly[3][0] - poly[2][0]) * (poly[3][1] + poly[2][1]),
|
||||
(poly[0][0] - poly[3][0]) * (poly[0][1] + poly[3][1])]
|
||||
return np.sum(edge) / 2.
|
||||
|
||||
def gen_quad_from_poly(self, poly):
|
||||
"""
|
||||
Generate min area quad from poly.
|
||||
"""
|
||||
point_num = poly.shape[0]
|
||||
min_area_quad = np.zeros((4, 2), dtype=np.float32)
|
||||
rect = cv2.minAreaRect(poly.astype(
|
||||
np.int32)) # (center (x,y), (width, height), angle of rotation)
|
||||
box = np.array(cv2.boxPoints(rect))
|
||||
|
||||
first_point_idx = 0
|
||||
min_dist = 1e4
|
||||
for i in range(4):
|
||||
dist = np.linalg.norm(box[(i + 0) % 4] - poly[0]) + \
|
||||
np.linalg.norm(box[(i + 1) % 4] - poly[point_num // 2 - 1]) + \
|
||||
np.linalg.norm(box[(i + 2) % 4] - poly[point_num // 2]) + \
|
||||
np.linalg.norm(box[(i + 3) % 4] - poly[-1])
|
||||
if dist < min_dist:
|
||||
min_dist = dist
|
||||
first_point_idx = i
|
||||
for i in range(4):
|
||||
min_area_quad[i] = box[(first_point_idx + i) % 4]
|
||||
|
||||
return min_area_quad
|
||||
|
||||
def check_and_validate_polys(self, polys, tags, im_size):
|
||||
"""
|
||||
check so that the text poly is in the same direction,
|
||||
and also filter some invalid polygons
|
||||
:param polys:
|
||||
:param tags:
|
||||
:return:
|
||||
"""
|
||||
(h, w) = im_size
|
||||
if polys.shape[0] == 0:
|
||||
return polys, np.array([]), np.array([])
|
||||
polys[:, :, 0] = np.clip(polys[:, :, 0], 0, w - 1)
|
||||
polys[:, :, 1] = np.clip(polys[:, :, 1], 0, h - 1)
|
||||
|
||||
validated_polys = []
|
||||
validated_tags = []
|
||||
hv_tags = []
|
||||
for poly, tag in zip(polys, tags):
|
||||
quad = self.gen_quad_from_poly(poly)
|
||||
p_area = self.quad_area(quad)
|
||||
if abs(p_area) < 1:
|
||||
print('invalid poly')
|
||||
continue
|
||||
if p_area > 0:
|
||||
if tag == False:
|
||||
print('poly in wrong direction')
|
||||
tag = True # reversed cases should be ignore
|
||||
poly = poly[(0, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2,
|
||||
1), :]
|
||||
quad = quad[(0, 3, 2, 1), :]
|
||||
|
||||
len_w = np.linalg.norm(quad[0] - quad[1]) + np.linalg.norm(quad[3] -
|
||||
quad[2])
|
||||
len_h = np.linalg.norm(quad[0] - quad[3]) + np.linalg.norm(quad[1] -
|
||||
quad[2])
|
||||
hv_tag = 1
|
||||
|
||||
if len_w * 2.0 < len_h:
|
||||
hv_tag = 0
|
||||
|
||||
validated_polys.append(poly)
|
||||
validated_tags.append(tag)
|
||||
hv_tags.append(hv_tag)
|
||||
return np.array(validated_polys), np.array(validated_tags), np.array(
|
||||
hv_tags)
|
||||
|
||||
def crop_area(self,
|
||||
im,
|
||||
polys,
|
||||
tags,
|
||||
hv_tags,
|
||||
txts,
|
||||
crop_background=False,
|
||||
max_tries=25):
|
||||
"""
|
||||
make random crop from the input image
|
||||
:param im:
|
||||
:param polys: [b,4,2]
|
||||
:param tags:
|
||||
:param crop_background:
|
||||
:param max_tries: 50 -> 25
|
||||
:return:
|
||||
"""
|
||||
h, w, _ = im.shape
|
||||
pad_h = h // 10
|
||||
pad_w = w // 10
|
||||
h_array = np.zeros((h + pad_h * 2), dtype=np.int32)
|
||||
w_array = np.zeros((w + pad_w * 2), dtype=np.int32)
|
||||
for poly in polys:
|
||||
poly = np.round(poly, decimals=0).astype(np.int32)
|
||||
minx = np.min(poly[:, 0])
|
||||
maxx = np.max(poly[:, 0])
|
||||
w_array[minx + pad_w:maxx + pad_w] = 1
|
||||
miny = np.min(poly[:, 1])
|
||||
maxy = np.max(poly[:, 1])
|
||||
h_array[miny + pad_h:maxy + pad_h] = 1
|
||||
# ensure the cropped area not across a text
|
||||
h_axis = np.where(h_array == 0)[0]
|
||||
w_axis = np.where(w_array == 0)[0]
|
||||
if len(h_axis) == 0 or len(w_axis) == 0:
|
||||
return im, polys, tags, hv_tags, txts
|
||||
for i in range(max_tries):
|
||||
xx = np.random.choice(w_axis, size=2)
|
||||
xmin = np.min(xx) - pad_w
|
||||
xmax = np.max(xx) - pad_w
|
||||
xmin = np.clip(xmin, 0, w - 1)
|
||||
xmax = np.clip(xmax, 0, w - 1)
|
||||
yy = np.random.choice(h_axis, size=2)
|
||||
ymin = np.min(yy) - pad_h
|
||||
ymax = np.max(yy) - pad_h
|
||||
ymin = np.clip(ymin, 0, h - 1)
|
||||
ymax = np.clip(ymax, 0, h - 1)
|
||||
if xmax - xmin < self.min_crop_size or \
|
||||
ymax - ymin < self.min_crop_size:
|
||||
continue
|
||||
if polys.shape[0] != 0:
|
||||
poly_axis_in_area = (polys[:, :, 0] >= xmin) & (polys[:, :, 0] <= xmax) \
|
||||
& (polys[:, :, 1] >= ymin) & (polys[:, :, 1] <= ymax)
|
||||
selected_polys = np.where(
|
||||
np.sum(poly_axis_in_area, axis=1) == 4)[0]
|
||||
else:
|
||||
selected_polys = []
|
||||
if len(selected_polys) == 0:
|
||||
# no text in this area
|
||||
if crop_background:
|
||||
txts_tmp = []
|
||||
for selected_poly in selected_polys:
|
||||
txts_tmp.append(txts[selected_poly])
|
||||
txts = txts_tmp
|
||||
return im[ymin: ymax + 1, xmin: xmax + 1, :], \
|
||||
polys[selected_polys], tags[selected_polys], hv_tags[selected_polys], txts
|
||||
else:
|
||||
continue
|
||||
im = im[ymin:ymax + 1, xmin:xmax + 1, :]
|
||||
polys = polys[selected_polys]
|
||||
tags = tags[selected_polys]
|
||||
hv_tags = hv_tags[selected_polys]
|
||||
txts_tmp = []
|
||||
for selected_poly in selected_polys:
|
||||
txts_tmp.append(txts[selected_poly])
|
||||
txts = txts_tmp
|
||||
polys[:, :, 0] -= xmin
|
||||
polys[:, :, 1] -= ymin
|
||||
return im, polys, tags, hv_tags, txts
|
||||
|
||||
return im, polys, tags, hv_tags, txts
|
||||
|
||||
def fit_and_gather_tcl_points_v2(self,
|
||||
min_area_quad,
|
||||
poly,
|
||||
max_h,
|
||||
max_w,
|
||||
fixed_point_num=64,
|
||||
img_id=0,
|
||||
reference_height=3):
|
||||
"""
|
||||
Find the center point of poly as key_points, then fit and gather.
|
||||
"""
|
||||
key_point_xys = []
|
||||
point_num = poly.shape[0]
|
||||
for idx in range(point_num // 2):
|
||||
center_point = (poly[idx] + poly[point_num - 1 - idx]) / 2.0
|
||||
key_point_xys.append(center_point)
|
||||
|
||||
tmp_image = np.zeros(
|
||||
shape=(
|
||||
max_h,
|
||||
max_w, ), dtype='float32')
|
||||
cv2.polylines(tmp_image, [np.array(key_point_xys).astype('int32')],
|
||||
False, 1.0)
|
||||
ys, xs = np.where(tmp_image > 0)
|
||||
xy_text = np.array(list(zip(xs, ys)), dtype='float32')
|
||||
|
||||
left_center_pt = (
|
||||
(min_area_quad[0] - min_area_quad[1]) / 2.0).reshape(1, 2)
|
||||
right_center_pt = (
|
||||
(min_area_quad[1] - min_area_quad[2]) / 2.0).reshape(1, 2)
|
||||
proj_unit_vec = (right_center_pt - left_center_pt) / (
|
||||
np.linalg.norm(right_center_pt - left_center_pt) + 1e-6)
|
||||
proj_unit_vec_tile = np.tile(proj_unit_vec,
|
||||
(xy_text.shape[0], 1)) # (n, 2)
|
||||
left_center_pt_tile = np.tile(left_center_pt,
|
||||
(xy_text.shape[0], 1)) # (n, 2)
|
||||
xy_text_to_left_center = xy_text - left_center_pt_tile
|
||||
proj_value = np.sum(xy_text_to_left_center * proj_unit_vec_tile, axis=1)
|
||||
xy_text = xy_text[np.argsort(proj_value)]
|
||||
|
||||
# convert to np and keep the num of point not greater then fixed_point_num
|
||||
pos_info = np.array(xy_text).reshape(-1, 2)[:, ::-1] # xy-> yx
|
||||
point_num = len(pos_info)
|
||||
if point_num > fixed_point_num:
|
||||
keep_ids = [
|
||||
int((point_num * 1.0 / fixed_point_num) * x)
|
||||
for x in range(fixed_point_num)
|
||||
]
|
||||
pos_info = pos_info[keep_ids, :]
|
||||
|
||||
keep = int(min(len(pos_info), fixed_point_num))
|
||||
if np.random.rand() < 0.2 and reference_height >= 3:
|
||||
dl = (np.random.rand(keep) - 0.5) * reference_height * 0.3
|
||||
random_float = np.array([1, 0]).reshape([1, 2]) * dl.reshape(
|
||||
[keep, 1])
|
||||
pos_info += random_float
|
||||
pos_info[:, 0] = np.clip(pos_info[:, 0], 0, max_h - 1)
|
||||
pos_info[:, 1] = np.clip(pos_info[:, 1], 0, max_w - 1)
|
||||
|
||||
# padding to fixed length
|
||||
pos_l = np.zeros((self.tcl_len, 3), dtype=np.int32)
|
||||
pos_l[:, 0] = np.ones((self.tcl_len, )) * img_id
|
||||
pos_m = np.zeros((self.tcl_len, 1), dtype=np.float32)
|
||||
pos_l[:keep, 1:] = np.round(pos_info).astype(np.int32)
|
||||
pos_m[:keep] = 1.0
|
||||
return pos_l, pos_m
|
||||
|
||||
def generate_direction_map(self, poly_quads, n_char, direction_map):
|
||||
"""
|
||||
"""
|
||||
width_list = []
|
||||
height_list = []
|
||||
for quad in poly_quads:
|
||||
quad_w = (np.linalg.norm(quad[0] - quad[1]) +
|
||||
np.linalg.norm(quad[2] - quad[3])) / 2.0
|
||||
quad_h = (np.linalg.norm(quad[0] - quad[3]) +
|
||||
np.linalg.norm(quad[2] - quad[1])) / 2.0
|
||||
width_list.append(quad_w)
|
||||
height_list.append(quad_h)
|
||||
norm_width = max(sum(width_list) / n_char, 1.0)
|
||||
average_height = max(sum(height_list) / len(height_list), 1.0)
|
||||
k = 1
|
||||
for quad in poly_quads:
|
||||
direct_vector_full = (
|
||||
(quad[1] + quad[2]) - (quad[0] + quad[3])) / 2.0
|
||||
direct_vector = direct_vector_full / (
|
||||
np.linalg.norm(direct_vector_full) + 1e-6) * norm_width
|
||||
direction_label = tuple(
|
||||
map(float,
|
||||
[direct_vector[0], direct_vector[1], 1.0 / average_height]))
|
||||
cv2.fillPoly(direction_map,
|
||||
quad.round().astype(np.int32)[np.newaxis, :, :],
|
||||
direction_label)
|
||||
k += 1
|
||||
return direction_map
|
||||
|
||||
def calculate_average_height(self, poly_quads):
|
||||
"""
|
||||
"""
|
||||
height_list = []
|
||||
for quad in poly_quads:
|
||||
quad_h = (np.linalg.norm(quad[0] - quad[3]) +
|
||||
np.linalg.norm(quad[2] - quad[1])) / 2.0
|
||||
height_list.append(quad_h)
|
||||
average_height = max(sum(height_list) / len(height_list), 1.0)
|
||||
return average_height
|
||||
|
||||
def generate_tcl_ctc_label(self,
|
||||
h,
|
||||
w,
|
||||
polys,
|
||||
tags,
|
||||
text_strs,
|
||||
ds_ratio,
|
||||
tcl_ratio=0.3,
|
||||
shrink_ratio_of_width=0.15):
|
||||
"""
|
||||
Generate polygon.
|
||||
"""
|
||||
score_map_big = np.zeros(
|
||||
(
|
||||
h,
|
||||
w, ), dtype=np.float32)
|
||||
h, w = int(h * ds_ratio), int(w * ds_ratio)
|
||||
polys = polys * ds_ratio
|
||||
|
||||
score_map = np.zeros(
|
||||
(
|
||||
h,
|
||||
w, ), dtype=np.float32)
|
||||
score_label_map = np.zeros(
|
||||
(
|
||||
h,
|
||||
w, ), dtype=np.float32)
|
||||
tbo_map = np.zeros((h, w, 5), dtype=np.float32)
|
||||
training_mask = np.ones(
|
||||
(
|
||||
h,
|
||||
w, ), dtype=np.float32)
|
||||
direction_map = np.ones((h, w, 3)) * np.array([0, 0, 1]).reshape(
|
||||
[1, 1, 3]).astype(np.float32)
|
||||
|
||||
label_idx = 0
|
||||
score_label_map_text_label_list = []
|
||||
pos_list, pos_mask, label_list = [], [], []
|
||||
for poly_idx, poly_tag in enumerate(zip(polys, tags)):
|
||||
poly = poly_tag[0]
|
||||
tag = poly_tag[1]
|
||||
|
||||
# generate min_area_quad
|
||||
min_area_quad, center_point = self.gen_min_area_quad_from_poly(poly)
|
||||
min_area_quad_h = 0.5 * (
|
||||
np.linalg.norm(min_area_quad[0] - min_area_quad[3]) +
|
||||
np.linalg.norm(min_area_quad[1] - min_area_quad[2]))
|
||||
min_area_quad_w = 0.5 * (
|
||||
np.linalg.norm(min_area_quad[0] - min_area_quad[1]) +
|
||||
np.linalg.norm(min_area_quad[2] - min_area_quad[3]))
|
||||
|
||||
if min(min_area_quad_h, min_area_quad_w) < self.min_text_size * ds_ratio \
|
||||
or min(min_area_quad_h, min_area_quad_w) > self.max_text_size * ds_ratio:
|
||||
continue
|
||||
|
||||
if tag:
|
||||
cv2.fillPoly(training_mask,
|
||||
poly.astype(np.int32)[np.newaxis, :, :], 0.15)
|
||||
else:
|
||||
text_label = text_strs[poly_idx]
|
||||
text_label = self.prepare_text_label(text_label,
|
||||
self.Lexicon_Table)
|
||||
|
||||
text_label_index_list = [[self.Lexicon_Table.index(c_)]
|
||||
for c_ in text_label
|
||||
if c_ in self.Lexicon_Table]
|
||||
if len(text_label_index_list) < 1:
|
||||
continue
|
||||
|
||||
tcl_poly = self.poly2tcl(poly, tcl_ratio)
|
||||
tcl_quads = self.poly2quads(tcl_poly)
|
||||
poly_quads = self.poly2quads(poly)
|
||||
|
||||
stcl_quads, quad_index = self.shrink_poly_along_width(
|
||||
tcl_quads,
|
||||
shrink_ratio_of_width=shrink_ratio_of_width,
|
||||
expand_height_ratio=1.0 / tcl_ratio)
|
||||
|
||||
cv2.fillPoly(score_map,
|
||||
np.round(stcl_quads).astype(np.int32), 1.0)
|
||||
cv2.fillPoly(score_map_big,
|
||||
np.round(stcl_quads / ds_ratio).astype(np.int32),
|
||||
1.0)
|
||||
|
||||
for idx, quad in enumerate(stcl_quads):
|
||||
quad_mask = np.zeros((h, w), dtype=np.float32)
|
||||
quad_mask = cv2.fillPoly(
|
||||
quad_mask,
|
||||
np.round(quad[np.newaxis, :, :]).astype(np.int32), 1.0)
|
||||
tbo_map = self.gen_quad_tbo(poly_quads[quad_index[idx]],
|
||||
quad_mask, tbo_map)
|
||||
|
||||
# score label map and score_label_map_text_label_list for refine
|
||||
if label_idx == 0:
|
||||
text_pos_list_ = [[len(self.Lexicon_Table)], ]
|
||||
score_label_map_text_label_list.append(text_pos_list_)
|
||||
|
||||
label_idx += 1
|
||||
cv2.fillPoly(score_label_map,
|
||||
np.round(poly_quads).astype(np.int32), label_idx)
|
||||
score_label_map_text_label_list.append(text_label_index_list)
|
||||
|
||||
# direction info, fix-me
|
||||
n_char = len(text_label_index_list)
|
||||
direction_map = self.generate_direction_map(poly_quads, n_char,
|
||||
direction_map)
|
||||
|
||||
# pos info
|
||||
average_shrink_height = self.calculate_average_height(
|
||||
stcl_quads)
|
||||
pos_l, pos_m = self.fit_and_gather_tcl_points_v2(
|
||||
min_area_quad,
|
||||
poly,
|
||||
max_h=h,
|
||||
max_w=w,
|
||||
fixed_point_num=64,
|
||||
img_id=self.img_id,
|
||||
reference_height=average_shrink_height)
|
||||
|
||||
label_l = text_label_index_list
|
||||
if len(text_label_index_list) < 2:
|
||||
continue
|
||||
|
||||
pos_list.append(pos_l)
|
||||
pos_mask.append(pos_m)
|
||||
label_list.append(label_l)
|
||||
|
||||
# use big score_map for smooth tcl lines
|
||||
score_map_big_resized = cv2.resize(
|
||||
score_map_big, dsize=None, fx=ds_ratio, fy=ds_ratio)
|
||||
score_map = np.array(score_map_big_resized > 1e-3, dtype='float32')
|
||||
|
||||
return score_map, score_label_map, tbo_map, direction_map, training_mask, \
|
||||
pos_list, pos_mask, label_list, score_label_map_text_label_list
|
||||
|
||||
def adjust_point(self, poly):
|
||||
"""
|
||||
adjust point order.
|
||||
"""
|
||||
point_num = poly.shape[0]
|
||||
if point_num == 4:
|
||||
len_1 = np.linalg.norm(poly[0] - poly[1])
|
||||
len_2 = np.linalg.norm(poly[1] - poly[2])
|
||||
len_3 = np.linalg.norm(poly[2] - poly[3])
|
||||
len_4 = np.linalg.norm(poly[3] - poly[0])
|
||||
|
||||
if (len_1 + len_3) * 1.5 < (len_2 + len_4):
|
||||
poly = poly[[1, 2, 3, 0], :]
|
||||
|
||||
elif point_num > 4:
|
||||
vector_1 = poly[0] - poly[1]
|
||||
vector_2 = poly[1] - poly[2]
|
||||
cos_theta = np.dot(vector_1, vector_2) / (
|
||||
np.linalg.norm(vector_1) * np.linalg.norm(vector_2) + 1e-6)
|
||||
theta = np.arccos(np.round(cos_theta, decimals=4))
|
||||
|
||||
if abs(theta) > (70 / 180 * math.pi):
|
||||
index = list(range(1, point_num)) + [0]
|
||||
poly = poly[np.array(index), :]
|
||||
return poly
|
||||
|
||||
def gen_min_area_quad_from_poly(self, poly):
|
||||
"""
|
||||
Generate min area quad from poly.
|
||||
"""
|
||||
point_num = poly.shape[0]
|
||||
min_area_quad = np.zeros((4, 2), dtype=np.float32)
|
||||
if point_num == 4:
|
||||
min_area_quad = poly
|
||||
center_point = np.sum(poly, axis=0) / 4
|
||||
else:
|
||||
rect = cv2.minAreaRect(poly.astype(
|
||||
np.int32)) # (center (x,y), (width, height), angle of rotation)
|
||||
center_point = rect[0]
|
||||
box = np.array(cv2.boxPoints(rect))
|
||||
|
||||
first_point_idx = 0
|
||||
min_dist = 1e4
|
||||
for i in range(4):
|
||||
dist = np.linalg.norm(box[(i + 0) % 4] - poly[0]) + \
|
||||
np.linalg.norm(box[(i + 1) % 4] - poly[point_num // 2 - 1]) + \
|
||||
np.linalg.norm(box[(i + 2) % 4] - poly[point_num // 2]) + \
|
||||
np.linalg.norm(box[(i + 3) % 4] - poly[-1])
|
||||
if dist < min_dist:
|
||||
min_dist = dist
|
||||
first_point_idx = i
|
||||
|
||||
for i in range(4):
|
||||
min_area_quad[i] = box[(first_point_idx + i) % 4]
|
||||
|
||||
return min_area_quad, center_point
|
||||
|
||||
def shrink_quad_along_width(self,
|
||||
quad,
|
||||
begin_width_ratio=0.,
|
||||
end_width_ratio=1.):
|
||||
"""
|
||||
Generate shrink_quad_along_width.
|
||||
"""
|
||||
ratio_pair = np.array(
|
||||
[[begin_width_ratio], [end_width_ratio]], dtype=np.float32)
|
||||
p0_1 = quad[0] + (quad[1] - quad[0]) * ratio_pair
|
||||
p3_2 = quad[3] + (quad[2] - quad[3]) * ratio_pair
|
||||
return np.array([p0_1[0], p0_1[1], p3_2[1], p3_2[0]])
|
||||
|
||||
def shrink_poly_along_width(self,
|
||||
quads,
|
||||
shrink_ratio_of_width,
|
||||
expand_height_ratio=1.0):
|
||||
"""
|
||||
shrink poly with given length.
|
||||
"""
|
||||
upper_edge_list = []
|
||||
|
||||
def get_cut_info(edge_len_list, cut_len):
|
||||
for idx, edge_len in enumerate(edge_len_list):
|
||||
cut_len -= edge_len
|
||||
if cut_len <= 0.000001:
|
||||
ratio = (cut_len + edge_len_list[idx]) / edge_len_list[idx]
|
||||
return idx, ratio
|
||||
|
||||
for quad in quads:
|
||||
upper_edge_len = np.linalg.norm(quad[0] - quad[1])
|
||||
upper_edge_list.append(upper_edge_len)
|
||||
|
||||
# length of left edge and right edge.
|
||||
left_length = np.linalg.norm(quads[0][0] - quads[0][
|
||||
3]) * expand_height_ratio
|
||||
right_length = np.linalg.norm(quads[-1][1] - quads[-1][
|
||||
2]) * expand_height_ratio
|
||||
|
||||
shrink_length = min(left_length, right_length,
|
||||
sum(upper_edge_list)) * shrink_ratio_of_width
|
||||
# shrinking length
|
||||
upper_len_left = shrink_length
|
||||
upper_len_right = sum(upper_edge_list) - shrink_length
|
||||
|
||||
left_idx, left_ratio = get_cut_info(upper_edge_list, upper_len_left)
|
||||
left_quad = self.shrink_quad_along_width(
|
||||
quads[left_idx], begin_width_ratio=left_ratio, end_width_ratio=1)
|
||||
right_idx, right_ratio = get_cut_info(upper_edge_list, upper_len_right)
|
||||
right_quad = self.shrink_quad_along_width(
|
||||
quads[right_idx], begin_width_ratio=0, end_width_ratio=right_ratio)
|
||||
|
||||
out_quad_list = []
|
||||
if left_idx == right_idx:
|
||||
out_quad_list.append(
|
||||
[left_quad[0], right_quad[1], right_quad[2], left_quad[3]])
|
||||
else:
|
||||
out_quad_list.append(left_quad)
|
||||
for idx in range(left_idx + 1, right_idx):
|
||||
out_quad_list.append(quads[idx])
|
||||
out_quad_list.append(right_quad)
|
||||
|
||||
return np.array(out_quad_list), list(range(left_idx, right_idx + 1))
|
||||
|
||||
def prepare_text_label(self, label_str, Lexicon_Table):
|
||||
"""
|
||||
Prepare text lablel by given Lexicon_Table.
|
||||
"""
|
||||
if len(Lexicon_Table) == 36:
|
||||
return label_str.lower()
|
||||
else:
|
||||
return label_str
|
||||
|
||||
def vector_angle(self, A, B):
|
||||
"""
|
||||
Calculate the angle between vector AB and x-axis positive direction.
|
||||
"""
|
||||
AB = np.array([B[1] - A[1], B[0] - A[0]])
|
||||
return np.arctan2(*AB)
|
||||
|
||||
def theta_line_cross_point(self, theta, point):
|
||||
"""
|
||||
Calculate the line through given point and angle in ax + by + c =0 form.
|
||||
"""
|
||||
x, y = point
|
||||
cos = np.cos(theta)
|
||||
sin = np.sin(theta)
|
||||
return [sin, -cos, cos * y - sin * x]
|
||||
|
||||
def line_cross_two_point(self, A, B):
|
||||
"""
|
||||
Calculate the line through given point A and B in ax + by + c =0 form.
|
||||
"""
|
||||
angle = self.vector_angle(A, B)
|
||||
return self.theta_line_cross_point(angle, A)
|
||||
|
||||
def average_angle(self, poly):
|
||||
"""
|
||||
Calculate the average angle between left and right edge in given poly.
|
||||
"""
|
||||
p0, p1, p2, p3 = poly
|
||||
angle30 = self.vector_angle(p3, p0)
|
||||
angle21 = self.vector_angle(p2, p1)
|
||||
return (angle30 + angle21) / 2
|
||||
|
||||
def line_cross_point(self, line1, line2):
|
||||
"""
|
||||
line1 and line2 in 0=ax+by+c form, compute the cross point of line1 and line2
|
||||
"""
|
||||
a1, b1, c1 = line1
|
||||
a2, b2, c2 = line2
|
||||
d = a1 * b2 - a2 * b1
|
||||
|
||||
if d == 0:
|
||||
print('Cross point does not exist')
|
||||
return np.array([0, 0], dtype=np.float32)
|
||||
else:
|
||||
x = (b1 * c2 - b2 * c1) / d
|
||||
y = (a2 * c1 - a1 * c2) / d
|
||||
|
||||
return np.array([x, y], dtype=np.float32)
|
||||
|
||||
def quad2tcl(self, poly, ratio):
|
||||
"""
|
||||
Generate center line by poly clock-wise point. (4, 2)
|
||||
"""
|
||||
ratio_pair = np.array(
|
||||
[[0.5 - ratio / 2], [0.5 + ratio / 2]], dtype=np.float32)
|
||||
p0_3 = poly[0] + (poly[3] - poly[0]) * ratio_pair
|
||||
p1_2 = poly[1] + (poly[2] - poly[1]) * ratio_pair
|
||||
return np.array([p0_3[0], p1_2[0], p1_2[1], p0_3[1]])
|
||||
|
||||
def poly2tcl(self, poly, ratio):
|
||||
"""
|
||||
Generate center line by poly clock-wise point.
|
||||
"""
|
||||
ratio_pair = np.array(
|
||||
[[0.5 - ratio / 2], [0.5 + ratio / 2]], dtype=np.float32)
|
||||
tcl_poly = np.zeros_like(poly)
|
||||
point_num = poly.shape[0]
|
||||
|
||||
for idx in range(point_num // 2):
|
||||
point_pair = poly[idx] + (poly[point_num - 1 - idx] - poly[idx]
|
||||
) * ratio_pair
|
||||
tcl_poly[idx] = point_pair[0]
|
||||
tcl_poly[point_num - 1 - idx] = point_pair[1]
|
||||
return tcl_poly
|
||||
|
||||
def gen_quad_tbo(self, quad, tcl_mask, tbo_map):
|
||||
"""
|
||||
Generate tbo_map for give quad.
|
||||
"""
|
||||
# upper and lower line function: ax + by + c = 0;
|
||||
up_line = self.line_cross_two_point(quad[0], quad[1])
|
||||
lower_line = self.line_cross_two_point(quad[3], quad[2])
|
||||
|
||||
quad_h = 0.5 * (np.linalg.norm(quad[0] - quad[3]) +
|
||||
np.linalg.norm(quad[1] - quad[2]))
|
||||
quad_w = 0.5 * (np.linalg.norm(quad[0] - quad[1]) +
|
||||
np.linalg.norm(quad[2] - quad[3]))
|
||||
|
||||
# average angle of left and right line.
|
||||
angle = self.average_angle(quad)
|
||||
|
||||
xy_in_poly = np.argwhere(tcl_mask == 1)
|
||||
for y, x in xy_in_poly:
|
||||
point = (x, y)
|
||||
line = self.theta_line_cross_point(angle, point)
|
||||
cross_point_upper = self.line_cross_point(up_line, line)
|
||||
cross_point_lower = self.line_cross_point(lower_line, line)
|
||||
##FIX, offset reverse
|
||||
upper_offset_x, upper_offset_y = cross_point_upper - point
|
||||
lower_offset_x, lower_offset_y = cross_point_lower - point
|
||||
tbo_map[y, x, 0] = upper_offset_y
|
||||
tbo_map[y, x, 1] = upper_offset_x
|
||||
tbo_map[y, x, 2] = lower_offset_y
|
||||
tbo_map[y, x, 3] = lower_offset_x
|
||||
tbo_map[y, x, 4] = 1.0 / max(min(quad_h, quad_w), 1.0) * 2
|
||||
return tbo_map
|
||||
|
||||
def poly2quads(self, poly):
|
||||
"""
|
||||
Split poly into quads.
|
||||
"""
|
||||
quad_list = []
|
||||
point_num = poly.shape[0]
|
||||
|
||||
# point pair
|
||||
point_pair_list = []
|
||||
for idx in range(point_num // 2):
|
||||
point_pair = [poly[idx], poly[point_num - 1 - idx]]
|
||||
point_pair_list.append(point_pair)
|
||||
|
||||
quad_num = point_num // 2 - 1
|
||||
for idx in range(quad_num):
|
||||
# reshape and adjust to clock-wise
|
||||
quad_list.append((np.array(point_pair_list)[[idx, idx + 1]]
|
||||
).reshape(4, 2)[[0, 2, 3, 1]])
|
||||
|
||||
return np.array(quad_list)
|
||||
|
||||
def rotate_im_poly(self, im, text_polys):
|
||||
"""
|
||||
rotate image with 90 / 180 / 270 degre
|
||||
"""
|
||||
im_w, im_h = im.shape[1], im.shape[0]
|
||||
dst_im = im.copy()
|
||||
dst_polys = []
|
||||
rand_degree_ratio = np.random.rand()
|
||||
rand_degree_cnt = 1
|
||||
if rand_degree_ratio > 0.5:
|
||||
rand_degree_cnt = 3
|
||||
for i in range(rand_degree_cnt):
|
||||
dst_im = np.rot90(dst_im)
|
||||
rot_degree = -90 * rand_degree_cnt
|
||||
rot_angle = rot_degree * math.pi / 180.0
|
||||
n_poly = text_polys.shape[0]
|
||||
cx, cy = 0.5 * im_w, 0.5 * im_h
|
||||
ncx, ncy = 0.5 * dst_im.shape[1], 0.5 * dst_im.shape[0]
|
||||
for i in range(n_poly):
|
||||
wordBB = text_polys[i]
|
||||
poly = []
|
||||
for j in range(4): # 16->4
|
||||
sx, sy = wordBB[j][0], wordBB[j][1]
|
||||
dx = math.cos(rot_angle) * (sx - cx) - math.sin(rot_angle) * (
|
||||
sy - cy) + ncx
|
||||
dy = math.sin(rot_angle) * (sx - cx) + math.cos(rot_angle) * (
|
||||
sy - cy) + ncy
|
||||
poly.append([dx, dy])
|
||||
dst_polys.append(poly)
|
||||
return dst_im, np.array(dst_polys, dtype=np.float32)
|
||||
|
||||
def __call__(self, data):
|
||||
input_size = 512
|
||||
im = data['image']
|
||||
text_polys = data['polys']
|
||||
text_tags = data['ignore_tags']
|
||||
text_strs = data['texts']
|
||||
h, w, _ = im.shape
|
||||
text_polys, text_tags, hv_tags = self.check_and_validate_polys(
|
||||
text_polys, text_tags, (h, w))
|
||||
if text_polys.shape[0] <= 0:
|
||||
return None
|
||||
# set aspect ratio and keep area fix
|
||||
asp_scales = np.arange(1.0, 1.55, 0.1)
|
||||
asp_scale = np.random.choice(asp_scales)
|
||||
if np.random.rand() < 0.5:
|
||||
asp_scale = 1.0 / asp_scale
|
||||
asp_scale = math.sqrt(asp_scale)
|
||||
|
||||
asp_wx = asp_scale
|
||||
asp_hy = 1.0 / asp_scale
|
||||
im = cv2.resize(im, dsize=None, fx=asp_wx, fy=asp_hy)
|
||||
text_polys[:, :, 0] *= asp_wx
|
||||
text_polys[:, :, 1] *= asp_hy
|
||||
|
||||
h, w, _ = im.shape
|
||||
if max(h, w) > 2048:
|
||||
rd_scale = 2048.0 / max(h, w)
|
||||
im = cv2.resize(im, dsize=None, fx=rd_scale, fy=rd_scale)
|
||||
text_polys *= rd_scale
|
||||
h, w, _ = im.shape
|
||||
if min(h, w) < 16:
|
||||
return None
|
||||
|
||||
# no background
|
||||
im, text_polys, text_tags, hv_tags, text_strs = self.crop_area(
|
||||
im,
|
||||
text_polys,
|
||||
text_tags,
|
||||
hv_tags,
|
||||
text_strs,
|
||||
crop_background=False)
|
||||
|
||||
if text_polys.shape[0] == 0:
|
||||
return None
|
||||
# # continue for all ignore case
|
||||
if np.sum((text_tags * 1.0)) >= text_tags.size:
|
||||
return None
|
||||
new_h, new_w, _ = im.shape
|
||||
if (new_h is None) or (new_w is None):
|
||||
return None
|
||||
# resize image
|
||||
std_ratio = float(input_size) / max(new_w, new_h)
|
||||
rand_scales = np.array(
|
||||
[0.25, 0.375, 0.5, 0.625, 0.75, 0.875, 1.0, 1.0, 1.0, 1.0, 1.0])
|
||||
rz_scale = std_ratio * np.random.choice(rand_scales)
|
||||
im = cv2.resize(im, dsize=None, fx=rz_scale, fy=rz_scale)
|
||||
text_polys[:, :, 0] *= rz_scale
|
||||
text_polys[:, :, 1] *= rz_scale
|
||||
|
||||
# add gaussian blur
|
||||
if np.random.rand() < 0.1 * 0.5:
|
||||
ks = np.random.permutation(5)[0] + 1
|
||||
ks = int(ks / 2) * 2 + 1
|
||||
im = cv2.GaussianBlur(im, ksize=(ks, ks), sigmaX=0, sigmaY=0)
|
||||
# add brighter
|
||||
if np.random.rand() < 0.1 * 0.5:
|
||||
im = im * (1.0 + np.random.rand() * 0.5)
|
||||
im = np.clip(im, 0.0, 255.0)
|
||||
# add darker
|
||||
if np.random.rand() < 0.1 * 0.5:
|
||||
im = im * (1.0 - np.random.rand() * 0.5)
|
||||
im = np.clip(im, 0.0, 255.0)
|
||||
|
||||
# Padding the im to [input_size, input_size]
|
||||
new_h, new_w, _ = im.shape
|
||||
if min(new_w, new_h) < input_size * 0.5:
|
||||
return None
|
||||
im_padded = np.ones((input_size, input_size, 3), dtype=np.float32)
|
||||
im_padded[:, :, 2] = 0.485 * 255
|
||||
im_padded[:, :, 1] = 0.456 * 255
|
||||
im_padded[:, :, 0] = 0.406 * 255
|
||||
|
||||
# Random the start position
|
||||
del_h = input_size - new_h
|
||||
del_w = input_size - new_w
|
||||
sh, sw = 0, 0
|
||||
if del_h > 1:
|
||||
sh = int(np.random.rand() * del_h)
|
||||
if del_w > 1:
|
||||
sw = int(np.random.rand() * del_w)
|
||||
|
||||
# Padding
|
||||
im_padded[sh:sh + new_h, sw:sw + new_w, :] = im.copy()
|
||||
text_polys[:, :, 0] += sw
|
||||
text_polys[:, :, 1] += sh
|
||||
|
||||
score_map, score_label_map, border_map, direction_map, training_mask, \
|
||||
pos_list, pos_mask, label_list, score_label_map_text_label = self.generate_tcl_ctc_label(input_size,
|
||||
input_size,
|
||||
text_polys,
|
||||
text_tags,
|
||||
text_strs, 0.25)
|
||||
if len(label_list) <= 0: # eliminate negative samples
|
||||
return None
|
||||
pos_list_temp = np.zeros([64, 3])
|
||||
pos_mask_temp = np.zeros([64, 1])
|
||||
label_list_temp = np.zeros([self.max_text_length, 1]) + self.pad_num
|
||||
|
||||
for i, label in enumerate(label_list):
|
||||
n = len(label)
|
||||
if n > self.max_text_length:
|
||||
label_list[i] = label[:self.max_text_length]
|
||||
continue
|
||||
while n < self.max_text_length:
|
||||
label.append([self.pad_num])
|
||||
n += 1
|
||||
|
||||
for i in range(len(label_list)):
|
||||
label_list[i] = np.array(label_list[i])
|
||||
|
||||
if len(pos_list) <= 0 or len(pos_list) > self.max_text_nums:
|
||||
return None
|
||||
for __ in range(self.max_text_nums - len(pos_list), 0, -1):
|
||||
pos_list.append(pos_list_temp)
|
||||
pos_mask.append(pos_mask_temp)
|
||||
label_list.append(label_list_temp)
|
||||
|
||||
if self.img_id == self.batch_size - 1:
|
||||
self.img_id = 0
|
||||
else:
|
||||
self.img_id += 1
|
||||
|
||||
im_padded[:, :, 2] -= 0.485 * 255
|
||||
im_padded[:, :, 1] -= 0.456 * 255
|
||||
im_padded[:, :, 0] -= 0.406 * 255
|
||||
im_padded[:, :, 2] /= (255.0 * 0.229)
|
||||
im_padded[:, :, 1] /= (255.0 * 0.224)
|
||||
im_padded[:, :, 0] /= (255.0 * 0.225)
|
||||
im_padded = im_padded.transpose((2, 0, 1))
|
||||
images = im_padded[::-1, :, :]
|
||||
tcl_maps = score_map[np.newaxis, :, :]
|
||||
tcl_label_maps = score_label_map[np.newaxis, :, :]
|
||||
border_maps = border_map.transpose((2, 0, 1))
|
||||
direction_maps = direction_map.transpose((2, 0, 1))
|
||||
training_masks = training_mask[np.newaxis, :, :]
|
||||
pos_list = np.array(pos_list)
|
||||
pos_mask = np.array(pos_mask)
|
||||
label_list = np.array(label_list)
|
||||
data['images'] = images
|
||||
data['tcl_maps'] = tcl_maps
|
||||
data['tcl_label_maps'] = tcl_label_maps
|
||||
data['border_maps'] = border_maps
|
||||
data['direction_maps'] = direction_maps
|
||||
data['training_masks'] = training_masks
|
||||
data['label_list'] = label_list
|
||||
data['pos_list'] = pos_list
|
||||
data['pos_mask'] = pos_mask
|
||||
return data
|
||||
@@ -0,0 +1,143 @@
|
||||
# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
from __future__ import unicode_literals
|
||||
|
||||
from PIL import Image, ImageEnhance, ImageOps
|
||||
import numpy as np
|
||||
import random
|
||||
import six
|
||||
|
||||
|
||||
class RawRandAugment(object):
|
||||
def __init__(self,
|
||||
num_layers=2,
|
||||
magnitude=5,
|
||||
fillcolor=(128, 128, 128),
|
||||
**kwargs):
|
||||
self.num_layers = num_layers
|
||||
self.magnitude = magnitude
|
||||
self.max_level = 10
|
||||
|
||||
abso_level = self.magnitude / self.max_level
|
||||
self.level_map = {
|
||||
"shearX": 0.3 * abso_level,
|
||||
"shearY": 0.3 * abso_level,
|
||||
"translateX": 150.0 / 331 * abso_level,
|
||||
"translateY": 150.0 / 331 * abso_level,
|
||||
"rotate": 30 * abso_level,
|
||||
"color": 0.9 * abso_level,
|
||||
"posterize": int(4.0 * abso_level),
|
||||
"solarize": 256.0 * abso_level,
|
||||
"contrast": 0.9 * abso_level,
|
||||
"sharpness": 0.9 * abso_level,
|
||||
"brightness": 0.9 * abso_level,
|
||||
"autocontrast": 0,
|
||||
"equalize": 0,
|
||||
"invert": 0
|
||||
}
|
||||
|
||||
# from https://stackoverflow.com/questions/5252170/
|
||||
# specify-image-filling-color-when-rotating-in-python-with-pil-and-setting-expand
|
||||
def rotate_with_fill(img, magnitude):
|
||||
rot = img.convert("RGBA").rotate(magnitude)
|
||||
return Image.composite(rot,
|
||||
Image.new("RGBA", rot.size, (128, ) * 4),
|
||||
rot).convert(img.mode)
|
||||
|
||||
rnd_ch_op = random.choice
|
||||
|
||||
self.func = {
|
||||
"shearX": lambda img, magnitude: img.transform(
|
||||
img.size,
|
||||
Image.AFFINE,
|
||||
(1, magnitude * rnd_ch_op([-1, 1]), 0, 0, 1, 0),
|
||||
Image.BICUBIC,
|
||||
fillcolor=fillcolor),
|
||||
"shearY": lambda img, magnitude: img.transform(
|
||||
img.size,
|
||||
Image.AFFINE,
|
||||
(1, 0, 0, magnitude * rnd_ch_op([-1, 1]), 1, 0),
|
||||
Image.BICUBIC,
|
||||
fillcolor=fillcolor),
|
||||
"translateX": lambda img, magnitude: img.transform(
|
||||
img.size,
|
||||
Image.AFFINE,
|
||||
(1, 0, magnitude * img.size[0] * rnd_ch_op([-1, 1]), 0, 1, 0),
|
||||
fillcolor=fillcolor),
|
||||
"translateY": lambda img, magnitude: img.transform(
|
||||
img.size,
|
||||
Image.AFFINE,
|
||||
(1, 0, 0, 0, 1, magnitude * img.size[1] * rnd_ch_op([-1, 1])),
|
||||
fillcolor=fillcolor),
|
||||
"rotate": lambda img, magnitude: rotate_with_fill(img, magnitude),
|
||||
"color": lambda img, magnitude: ImageEnhance.Color(img).enhance(
|
||||
1 + magnitude * rnd_ch_op([-1, 1])),
|
||||
"posterize": lambda img, magnitude:
|
||||
ImageOps.posterize(img, magnitude),
|
||||
"solarize": lambda img, magnitude:
|
||||
ImageOps.solarize(img, magnitude),
|
||||
"contrast": lambda img, magnitude:
|
||||
ImageEnhance.Contrast(img).enhance(
|
||||
1 + magnitude * rnd_ch_op([-1, 1])),
|
||||
"sharpness": lambda img, magnitude:
|
||||
ImageEnhance.Sharpness(img).enhance(
|
||||
1 + magnitude * rnd_ch_op([-1, 1])),
|
||||
"brightness": lambda img, magnitude:
|
||||
ImageEnhance.Brightness(img).enhance(
|
||||
1 + magnitude * rnd_ch_op([-1, 1])),
|
||||
"autocontrast": lambda img, magnitude:
|
||||
ImageOps.autocontrast(img),
|
||||
"equalize": lambda img, magnitude: ImageOps.equalize(img),
|
||||
"invert": lambda img, magnitude: ImageOps.invert(img)
|
||||
}
|
||||
|
||||
def __call__(self, img):
|
||||
avaiable_op_names = list(self.level_map.keys())
|
||||
for layer_num in range(self.num_layers):
|
||||
op_name = np.random.choice(avaiable_op_names)
|
||||
img = self.func[op_name](img, self.level_map[op_name])
|
||||
return img
|
||||
|
||||
|
||||
class RandAugment(RawRandAugment):
|
||||
""" RandAugment wrapper to auto fit different img types """
|
||||
|
||||
def __init__(self, prob=0.5, *args, **kwargs):
|
||||
self.prob = prob
|
||||
if six.PY2:
|
||||
super(RandAugment, self).__init__(*args, **kwargs)
|
||||
else:
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
def __call__(self, data):
|
||||
if np.random.rand() > self.prob:
|
||||
return data
|
||||
img = data['image']
|
||||
if not isinstance(img, Image.Image):
|
||||
img = np.ascontiguousarray(img)
|
||||
img = Image.fromarray(img)
|
||||
|
||||
if six.PY2:
|
||||
img = super(RandAugment, self).__call__(img)
|
||||
else:
|
||||
img = super().__call__(img)
|
||||
|
||||
if isinstance(img, Image.Image):
|
||||
img = np.asarray(img)
|
||||
data['image'] = img
|
||||
return data
|
||||
@@ -0,0 +1,234 @@
|
||||
# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""
|
||||
This code is refer from:
|
||||
https://github.com/WenmuZhou/DBNet.pytorch/blob/master/data_loader/modules/random_crop_data.py
|
||||
"""
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
from __future__ import unicode_literals
|
||||
|
||||
import numpy as np
|
||||
import cv2
|
||||
import random
|
||||
|
||||
|
||||
def is_poly_in_rect(poly, x, y, w, h):
|
||||
poly = np.array(poly)
|
||||
if poly[:, 0].min() < x or poly[:, 0].max() > x + w:
|
||||
return False
|
||||
if poly[:, 1].min() < y or poly[:, 1].max() > y + h:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def is_poly_outside_rect(poly, x, y, w, h):
|
||||
poly = np.array(poly)
|
||||
if poly[:, 0].max() < x or poly[:, 0].min() > x + w:
|
||||
return True
|
||||
if poly[:, 1].max() < y or poly[:, 1].min() > y + h:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def split_regions(axis):
|
||||
regions = []
|
||||
min_axis = 0
|
||||
for i in range(1, axis.shape[0]):
|
||||
if axis[i] != axis[i - 1] + 1:
|
||||
region = axis[min_axis:i]
|
||||
min_axis = i
|
||||
regions.append(region)
|
||||
return regions
|
||||
|
||||
|
||||
def random_select(axis, max_size):
|
||||
xx = np.random.choice(axis, size=2)
|
||||
xmin = np.min(xx)
|
||||
xmax = np.max(xx)
|
||||
xmin = np.clip(xmin, 0, max_size - 1)
|
||||
xmax = np.clip(xmax, 0, max_size - 1)
|
||||
return xmin, xmax
|
||||
|
||||
|
||||
def region_wise_random_select(regions, max_size):
|
||||
selected_index = list(np.random.choice(len(regions), 2))
|
||||
selected_values = []
|
||||
for index in selected_index:
|
||||
axis = regions[index]
|
||||
xx = int(np.random.choice(axis, size=1))
|
||||
selected_values.append(xx)
|
||||
xmin = min(selected_values)
|
||||
xmax = max(selected_values)
|
||||
return xmin, xmax
|
||||
|
||||
|
||||
def crop_area(im, text_polys, min_crop_side_ratio, max_tries):
|
||||
h, w, _ = im.shape
|
||||
h_array = np.zeros(h, dtype=np.int32)
|
||||
w_array = np.zeros(w, dtype=np.int32)
|
||||
for points in text_polys:
|
||||
points = np.round(points, decimals=0).astype(np.int32)
|
||||
minx = np.min(points[:, 0])
|
||||
maxx = np.max(points[:, 0])
|
||||
w_array[minx:maxx] = 1
|
||||
miny = np.min(points[:, 1])
|
||||
maxy = np.max(points[:, 1])
|
||||
h_array[miny:maxy] = 1
|
||||
# ensure the cropped area not across a text
|
||||
h_axis = np.where(h_array == 0)[0]
|
||||
w_axis = np.where(w_array == 0)[0]
|
||||
|
||||
if len(h_axis) == 0 or len(w_axis) == 0:
|
||||
return 0, 0, w, h
|
||||
|
||||
h_regions = split_regions(h_axis)
|
||||
w_regions = split_regions(w_axis)
|
||||
|
||||
for i in range(max_tries):
|
||||
if len(w_regions) > 1:
|
||||
xmin, xmax = region_wise_random_select(w_regions, w)
|
||||
else:
|
||||
xmin, xmax = random_select(w_axis, w)
|
||||
if len(h_regions) > 1:
|
||||
ymin, ymax = region_wise_random_select(h_regions, h)
|
||||
else:
|
||||
ymin, ymax = random_select(h_axis, h)
|
||||
|
||||
if xmax - xmin < min_crop_side_ratio * w or ymax - ymin < min_crop_side_ratio * h:
|
||||
# area too small
|
||||
continue
|
||||
num_poly_in_rect = 0
|
||||
for poly in text_polys:
|
||||
if not is_poly_outside_rect(poly, xmin, ymin, xmax - xmin,
|
||||
ymax - ymin):
|
||||
num_poly_in_rect += 1
|
||||
break
|
||||
|
||||
if num_poly_in_rect > 0:
|
||||
return xmin, ymin, xmax - xmin, ymax - ymin
|
||||
|
||||
return 0, 0, w, h
|
||||
|
||||
|
||||
class EastRandomCropData(object):
|
||||
def __init__(self,
|
||||
size=(640, 640),
|
||||
max_tries=10,
|
||||
min_crop_side_ratio=0.1,
|
||||
keep_ratio=True,
|
||||
**kwargs):
|
||||
self.size = size
|
||||
self.max_tries = max_tries
|
||||
self.min_crop_side_ratio = min_crop_side_ratio
|
||||
self.keep_ratio = keep_ratio
|
||||
|
||||
def __call__(self, data):
|
||||
img = data['image']
|
||||
text_polys = data['polys']
|
||||
ignore_tags = data['ignore_tags']
|
||||
texts = data['texts']
|
||||
all_care_polys = [
|
||||
text_polys[i] for i, tag in enumerate(ignore_tags) if not tag
|
||||
]
|
||||
# 计算crop区域
|
||||
crop_x, crop_y, crop_w, crop_h = crop_area(
|
||||
img, all_care_polys, self.min_crop_side_ratio, self.max_tries)
|
||||
# crop 图片 保持比例填充
|
||||
scale_w = self.size[0] / crop_w
|
||||
scale_h = self.size[1] / crop_h
|
||||
scale = min(scale_w, scale_h)
|
||||
h = int(crop_h * scale)
|
||||
w = int(crop_w * scale)
|
||||
if self.keep_ratio:
|
||||
padimg = np.zeros((self.size[1], self.size[0], img.shape[2]),
|
||||
img.dtype)
|
||||
padimg[:h, :w] = cv2.resize(
|
||||
img[crop_y:crop_y + crop_h, crop_x:crop_x + crop_w], (w, h))
|
||||
img = padimg
|
||||
else:
|
||||
img = cv2.resize(
|
||||
img[crop_y:crop_y + crop_h, crop_x:crop_x + crop_w],
|
||||
tuple(self.size))
|
||||
# crop 文本框
|
||||
text_polys_crop = []
|
||||
ignore_tags_crop = []
|
||||
texts_crop = []
|
||||
for poly, text, tag in zip(text_polys, texts, ignore_tags):
|
||||
poly = ((poly - (crop_x, crop_y)) * scale).tolist()
|
||||
if not is_poly_outside_rect(poly, 0, 0, w, h):
|
||||
text_polys_crop.append(poly)
|
||||
ignore_tags_crop.append(tag)
|
||||
texts_crop.append(text)
|
||||
data['image'] = img
|
||||
data['polys'] = np.array(text_polys_crop)
|
||||
data['ignore_tags'] = ignore_tags_crop
|
||||
data['texts'] = texts_crop
|
||||
return data
|
||||
|
||||
|
||||
class RandomCropImgMask(object):
|
||||
def __init__(self, size, main_key, crop_keys, p=3 / 8, **kwargs):
|
||||
self.size = size
|
||||
self.main_key = main_key
|
||||
self.crop_keys = crop_keys
|
||||
self.p = p
|
||||
|
||||
def __call__(self, data):
|
||||
image = data['image']
|
||||
|
||||
h, w = image.shape[0:2]
|
||||
th, tw = self.size
|
||||
if w == tw and h == th:
|
||||
return data
|
||||
|
||||
mask = data[self.main_key]
|
||||
if np.max(mask) > 0 and random.random() > self.p:
|
||||
# make sure to crop the text region
|
||||
tl = np.min(np.where(mask > 0), axis=1) - (th, tw)
|
||||
tl[tl < 0] = 0
|
||||
br = np.max(np.where(mask > 0), axis=1) - (th, tw)
|
||||
br[br < 0] = 0
|
||||
|
||||
br[0] = min(br[0], h - th)
|
||||
br[1] = min(br[1], w - tw)
|
||||
|
||||
i = random.randint(tl[0], br[0]) if tl[0] < br[0] else 0
|
||||
j = random.randint(tl[1], br[1]) if tl[1] < br[1] else 0
|
||||
else:
|
||||
i = random.randint(0, h - th) if h - th > 0 else 0
|
||||
j = random.randint(0, w - tw) if w - tw > 0 else 0
|
||||
|
||||
# return i, j, th, tw
|
||||
for k in data:
|
||||
if k in self.crop_keys:
|
||||
if len(data[k].shape) == 3:
|
||||
if np.argmin(data[k].shape) == 0:
|
||||
img = data[k][:, i:i + th, j:j + tw]
|
||||
if img.shape[1] != img.shape[2]:
|
||||
a = 1
|
||||
elif np.argmin(data[k].shape) == 2:
|
||||
img = data[k][i:i + th, j:j + tw, :]
|
||||
if img.shape[1] != img.shape[0]:
|
||||
a = 1
|
||||
else:
|
||||
img = data[k]
|
||||
else:
|
||||
img = data[k][i:i + th, j:j + tw]
|
||||
if img.shape[0] != img.shape[1]:
|
||||
a = 1
|
||||
data[k] = img
|
||||
return data
|
||||
@@ -0,0 +1,601 @@
|
||||
# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import math
|
||||
import cv2
|
||||
import numpy as np
|
||||
import random
|
||||
import copy
|
||||
from PIL import Image
|
||||
from .text_image_aug import tia_perspective, tia_stretch, tia_distort
|
||||
|
||||
|
||||
class RecAug(object):
|
||||
def __init__(self, use_tia=True, aug_prob=0.4, **kwargs):
|
||||
self.use_tia = use_tia
|
||||
self.aug_prob = aug_prob
|
||||
|
||||
def __call__(self, data):
|
||||
img = data['image']
|
||||
img = warp(img, 10, self.use_tia, self.aug_prob)
|
||||
data['image'] = img
|
||||
return data
|
||||
|
||||
|
||||
class RecConAug(object):
|
||||
def __init__(self,
|
||||
prob=0.5,
|
||||
image_shape=(32, 320, 3),
|
||||
max_text_length=25,
|
||||
ext_data_num=1,
|
||||
**kwargs):
|
||||
self.ext_data_num = ext_data_num
|
||||
self.prob = prob
|
||||
self.max_text_length = max_text_length
|
||||
self.image_shape = image_shape
|
||||
self.max_wh_ratio = self.image_shape[1] / self.image_shape[0]
|
||||
|
||||
def merge_ext_data(self, data, ext_data):
|
||||
ori_w = round(data['image'].shape[1] / data['image'].shape[0] *
|
||||
self.image_shape[0])
|
||||
ext_w = round(ext_data['image'].shape[1] / ext_data['image'].shape[0] *
|
||||
self.image_shape[0])
|
||||
data['image'] = cv2.resize(data['image'], (ori_w, self.image_shape[0]))
|
||||
ext_data['image'] = cv2.resize(ext_data['image'],
|
||||
(ext_w, self.image_shape[0]))
|
||||
data['image'] = np.concatenate(
|
||||
[data['image'], ext_data['image']], axis=1)
|
||||
data["label"] += ext_data["label"]
|
||||
return data
|
||||
|
||||
def __call__(self, data):
|
||||
rnd_num = random.random()
|
||||
if rnd_num > self.prob:
|
||||
return data
|
||||
for idx, ext_data in enumerate(data["ext_data"]):
|
||||
if len(data["label"]) + len(ext_data[
|
||||
"label"]) > self.max_text_length:
|
||||
break
|
||||
concat_ratio = data['image'].shape[1] / data['image'].shape[
|
||||
0] + ext_data['image'].shape[1] / ext_data['image'].shape[0]
|
||||
if concat_ratio > self.max_wh_ratio:
|
||||
break
|
||||
data = self.merge_ext_data(data, ext_data)
|
||||
data.pop("ext_data")
|
||||
return data
|
||||
|
||||
|
||||
class ClsResizeImg(object):
|
||||
def __init__(self, image_shape, **kwargs):
|
||||
self.image_shape = image_shape
|
||||
|
||||
def __call__(self, data):
|
||||
img = data['image']
|
||||
norm_img, _ = resize_norm_img(img, self.image_shape)
|
||||
data['image'] = norm_img
|
||||
return data
|
||||
|
||||
|
||||
class NRTRRecResizeImg(object):
|
||||
def __init__(self, image_shape, resize_type, padding=False, **kwargs):
|
||||
self.image_shape = image_shape
|
||||
self.resize_type = resize_type
|
||||
self.padding = padding
|
||||
|
||||
def __call__(self, data):
|
||||
img = data['image']
|
||||
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
||||
image_shape = self.image_shape
|
||||
if self.padding:
|
||||
imgC, imgH, imgW = image_shape
|
||||
# todo: change to 0 and modified image shape
|
||||
h = img.shape[0]
|
||||
w = img.shape[1]
|
||||
ratio = w / float(h)
|
||||
if math.ceil(imgH * ratio) > imgW:
|
||||
resized_w = imgW
|
||||
else:
|
||||
resized_w = int(math.ceil(imgH * ratio))
|
||||
resized_image = cv2.resize(img, (resized_w, imgH))
|
||||
norm_img = np.expand_dims(resized_image, -1)
|
||||
norm_img = norm_img.transpose((2, 0, 1))
|
||||
resized_image = norm_img.astype(np.float32) / 128. - 1.
|
||||
padding_im = np.zeros((imgC, imgH, imgW), dtype=np.float32)
|
||||
padding_im[:, :, 0:resized_w] = resized_image
|
||||
data['image'] = padding_im
|
||||
return data
|
||||
if self.resize_type == 'PIL':
|
||||
image_pil = Image.fromarray(np.uint8(img))
|
||||
img = image_pil.resize(self.image_shape, Image.ANTIALIAS)
|
||||
img = np.array(img)
|
||||
if self.resize_type == 'OpenCV':
|
||||
img = cv2.resize(img, self.image_shape)
|
||||
norm_img = np.expand_dims(img, -1)
|
||||
norm_img = norm_img.transpose((2, 0, 1))
|
||||
data['image'] = norm_img.astype(np.float32) / 128. - 1.
|
||||
return data
|
||||
|
||||
|
||||
class RecResizeImg(object):
|
||||
def __init__(self,
|
||||
image_shape,
|
||||
infer_mode=False,
|
||||
character_dict_path='./ppocr/utils/ppocr_keys_v1.txt',
|
||||
padding=True,
|
||||
**kwargs):
|
||||
self.image_shape = image_shape
|
||||
self.infer_mode = infer_mode
|
||||
self.character_dict_path = character_dict_path
|
||||
self.padding = padding
|
||||
|
||||
def __call__(self, data):
|
||||
img = data['image']
|
||||
if self.infer_mode and self.character_dict_path is not None:
|
||||
norm_img, valid_ratio = resize_norm_img_chinese(img,
|
||||
self.image_shape)
|
||||
else:
|
||||
norm_img, valid_ratio = resize_norm_img(img, self.image_shape,
|
||||
self.padding)
|
||||
data['image'] = norm_img
|
||||
data['valid_ratio'] = valid_ratio
|
||||
return data
|
||||
|
||||
|
||||
class SRNRecResizeImg(object):
|
||||
def __init__(self, image_shape, num_heads, max_text_length, **kwargs):
|
||||
self.image_shape = image_shape
|
||||
self.num_heads = num_heads
|
||||
self.max_text_length = max_text_length
|
||||
|
||||
def __call__(self, data):
|
||||
img = data['image']
|
||||
norm_img = resize_norm_img_srn(img, self.image_shape)
|
||||
data['image'] = norm_img
|
||||
[encoder_word_pos, gsrm_word_pos, gsrm_slf_attn_bias1, gsrm_slf_attn_bias2] = \
|
||||
srn_other_inputs(self.image_shape, self.num_heads, self.max_text_length)
|
||||
|
||||
data['encoder_word_pos'] = encoder_word_pos
|
||||
data['gsrm_word_pos'] = gsrm_word_pos
|
||||
data['gsrm_slf_attn_bias1'] = gsrm_slf_attn_bias1
|
||||
data['gsrm_slf_attn_bias2'] = gsrm_slf_attn_bias2
|
||||
return data
|
||||
|
||||
|
||||
class SARRecResizeImg(object):
|
||||
def __init__(self, image_shape, width_downsample_ratio=0.25, **kwargs):
|
||||
self.image_shape = image_shape
|
||||
self.width_downsample_ratio = width_downsample_ratio
|
||||
|
||||
def __call__(self, data):
|
||||
img = data['image']
|
||||
norm_img, resize_shape, pad_shape, valid_ratio = resize_norm_img_sar(
|
||||
img, self.image_shape, self.width_downsample_ratio)
|
||||
data['image'] = norm_img
|
||||
data['resized_shape'] = resize_shape
|
||||
data['pad_shape'] = pad_shape
|
||||
data['valid_ratio'] = valid_ratio
|
||||
return data
|
||||
|
||||
|
||||
class PRENResizeImg(object):
|
||||
def __init__(self, image_shape, **kwargs):
|
||||
"""
|
||||
Accroding to original paper's realization, it's a hard resize method here.
|
||||
So maybe you should optimize it to fit for your task better.
|
||||
"""
|
||||
self.dst_h, self.dst_w = image_shape
|
||||
|
||||
def __call__(self, data):
|
||||
img = data['image']
|
||||
resized_img = cv2.resize(
|
||||
img, (self.dst_w, self.dst_h), interpolation=cv2.INTER_LINEAR)
|
||||
resized_img = resized_img.transpose((2, 0, 1)) / 255
|
||||
resized_img -= 0.5
|
||||
resized_img /= 0.5
|
||||
data['image'] = resized_img.astype(np.float32)
|
||||
return data
|
||||
|
||||
|
||||
def resize_norm_img_sar(img, image_shape, width_downsample_ratio=0.25):
|
||||
imgC, imgH, imgW_min, imgW_max = image_shape
|
||||
h = img.shape[0]
|
||||
w = img.shape[1]
|
||||
valid_ratio = 1.0
|
||||
# make sure new_width is an integral multiple of width_divisor.
|
||||
width_divisor = int(1 / width_downsample_ratio)
|
||||
# resize
|
||||
ratio = w / float(h)
|
||||
resize_w = math.ceil(imgH * ratio)
|
||||
if resize_w % width_divisor != 0:
|
||||
resize_w = round(resize_w / width_divisor) * width_divisor
|
||||
if imgW_min is not None:
|
||||
resize_w = max(imgW_min, resize_w)
|
||||
if imgW_max is not None:
|
||||
valid_ratio = min(1.0, 1.0 * resize_w / imgW_max)
|
||||
resize_w = min(imgW_max, resize_w)
|
||||
resized_image = cv2.resize(img, (resize_w, imgH))
|
||||
resized_image = resized_image.astype('float32')
|
||||
# norm
|
||||
if image_shape[0] == 1:
|
||||
resized_image = resized_image / 255
|
||||
resized_image = resized_image[np.newaxis, :]
|
||||
else:
|
||||
resized_image = resized_image.transpose((2, 0, 1)) / 255
|
||||
resized_image -= 0.5
|
||||
resized_image /= 0.5
|
||||
resize_shape = resized_image.shape
|
||||
padding_im = -1.0 * np.ones((imgC, imgH, imgW_max), dtype=np.float32)
|
||||
padding_im[:, :, 0:resize_w] = resized_image
|
||||
pad_shape = padding_im.shape
|
||||
|
||||
return padding_im, resize_shape, pad_shape, valid_ratio
|
||||
|
||||
|
||||
def resize_norm_img(img, image_shape, padding=True):
|
||||
imgC, imgH, imgW = image_shape
|
||||
h = img.shape[0]
|
||||
w = img.shape[1]
|
||||
if not padding:
|
||||
resized_image = cv2.resize(
|
||||
img, (imgW, imgH), interpolation=cv2.INTER_LINEAR)
|
||||
resized_w = imgW
|
||||
else:
|
||||
ratio = w / float(h)
|
||||
if math.ceil(imgH * ratio) > imgW:
|
||||
resized_w = imgW
|
||||
else:
|
||||
resized_w = int(math.ceil(imgH * ratio))
|
||||
resized_image = cv2.resize(img, (resized_w, imgH))
|
||||
resized_image = resized_image.astype('float32')
|
||||
if image_shape[0] == 1:
|
||||
resized_image = resized_image / 255
|
||||
resized_image = resized_image[np.newaxis, :]
|
||||
else:
|
||||
resized_image = resized_image.transpose((2, 0, 1)) / 255
|
||||
resized_image -= 0.5
|
||||
resized_image /= 0.5
|
||||
padding_im = np.zeros((imgC, imgH, imgW), dtype=np.float32)
|
||||
padding_im[:, :, 0:resized_w] = resized_image
|
||||
valid_ratio = min(1.0, float(resized_w / imgW))
|
||||
return padding_im, valid_ratio
|
||||
|
||||
|
||||
def resize_norm_img_chinese(img, image_shape):
|
||||
imgC, imgH, imgW = image_shape
|
||||
# todo: change to 0 and modified image shape
|
||||
max_wh_ratio = imgW * 1.0 / imgH
|
||||
h, w = img.shape[0], img.shape[1]
|
||||
ratio = w * 1.0 / h
|
||||
max_wh_ratio = max(max_wh_ratio, ratio)
|
||||
imgW = int(imgH * max_wh_ratio)
|
||||
if math.ceil(imgH * ratio) > imgW:
|
||||
resized_w = imgW
|
||||
else:
|
||||
resized_w = int(math.ceil(imgH * ratio))
|
||||
resized_image = cv2.resize(img, (resized_w, imgH))
|
||||
resized_image = resized_image.astype('float32')
|
||||
if image_shape[0] == 1:
|
||||
resized_image = resized_image / 255
|
||||
resized_image = resized_image[np.newaxis, :]
|
||||
else:
|
||||
resized_image = resized_image.transpose((2, 0, 1)) / 255
|
||||
resized_image -= 0.5
|
||||
resized_image /= 0.5
|
||||
padding_im = np.zeros((imgC, imgH, imgW), dtype=np.float32)
|
||||
padding_im[:, :, 0:resized_w] = resized_image
|
||||
valid_ratio = min(1.0, float(resized_w / imgW))
|
||||
return padding_im, valid_ratio
|
||||
|
||||
|
||||
def resize_norm_img_srn(img, image_shape):
|
||||
imgC, imgH, imgW = image_shape
|
||||
|
||||
img_black = np.zeros((imgH, imgW))
|
||||
im_hei = img.shape[0]
|
||||
im_wid = img.shape[1]
|
||||
|
||||
if im_wid <= im_hei * 1:
|
||||
img_new = cv2.resize(img, (imgH * 1, imgH))
|
||||
elif im_wid <= im_hei * 2:
|
||||
img_new = cv2.resize(img, (imgH * 2, imgH))
|
||||
elif im_wid <= im_hei * 3:
|
||||
img_new = cv2.resize(img, (imgH * 3, imgH))
|
||||
else:
|
||||
img_new = cv2.resize(img, (imgW, imgH))
|
||||
|
||||
img_np = np.asarray(img_new)
|
||||
img_np = cv2.cvtColor(img_np, cv2.COLOR_BGR2GRAY)
|
||||
img_black[:, 0:img_np.shape[1]] = img_np
|
||||
img_black = img_black[:, :, np.newaxis]
|
||||
|
||||
row, col, c = img_black.shape
|
||||
c = 1
|
||||
|
||||
return np.reshape(img_black, (c, row, col)).astype(np.float32)
|
||||
|
||||
|
||||
def srn_other_inputs(image_shape, num_heads, max_text_length):
|
||||
|
||||
imgC, imgH, imgW = image_shape
|
||||
feature_dim = int((imgH / 8) * (imgW / 8))
|
||||
|
||||
encoder_word_pos = np.array(range(0, feature_dim)).reshape(
|
||||
(feature_dim, 1)).astype('int64')
|
||||
gsrm_word_pos = np.array(range(0, max_text_length)).reshape(
|
||||
(max_text_length, 1)).astype('int64')
|
||||
|
||||
gsrm_attn_bias_data = np.ones((1, max_text_length, max_text_length))
|
||||
gsrm_slf_attn_bias1 = np.triu(gsrm_attn_bias_data, 1).reshape(
|
||||
[1, max_text_length, max_text_length])
|
||||
gsrm_slf_attn_bias1 = np.tile(gsrm_slf_attn_bias1,
|
||||
[num_heads, 1, 1]) * [-1e9]
|
||||
|
||||
gsrm_slf_attn_bias2 = np.tril(gsrm_attn_bias_data, -1).reshape(
|
||||
[1, max_text_length, max_text_length])
|
||||
gsrm_slf_attn_bias2 = np.tile(gsrm_slf_attn_bias2,
|
||||
[num_heads, 1, 1]) * [-1e9]
|
||||
|
||||
return [
|
||||
encoder_word_pos, gsrm_word_pos, gsrm_slf_attn_bias1,
|
||||
gsrm_slf_attn_bias2
|
||||
]
|
||||
|
||||
|
||||
def flag():
|
||||
"""
|
||||
flag
|
||||
"""
|
||||
return 1 if random.random() > 0.5000001 else -1
|
||||
|
||||
|
||||
def cvtColor(img):
|
||||
"""
|
||||
cvtColor
|
||||
"""
|
||||
hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
|
||||
delta = 0.001 * random.random() * flag()
|
||||
hsv[:, :, 2] = hsv[:, :, 2] * (1 + delta)
|
||||
new_img = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)
|
||||
return new_img
|
||||
|
||||
|
||||
def blur(img):
|
||||
"""
|
||||
blur
|
||||
"""
|
||||
h, w, _ = img.shape
|
||||
if h > 10 and w > 10:
|
||||
return cv2.GaussianBlur(img, (5, 5), 1)
|
||||
else:
|
||||
return img
|
||||
|
||||
|
||||
def jitter(img):
|
||||
"""
|
||||
jitter
|
||||
"""
|
||||
w, h, _ = img.shape
|
||||
if h > 10 and w > 10:
|
||||
thres = min(w, h)
|
||||
s = int(random.random() * thres * 0.01)
|
||||
src_img = img.copy()
|
||||
for i in range(s):
|
||||
img[i:, i:, :] = src_img[:w - i, :h - i, :]
|
||||
return img
|
||||
else:
|
||||
return img
|
||||
|
||||
|
||||
def add_gasuss_noise(image, mean=0, var=0.1):
|
||||
"""
|
||||
Gasuss noise
|
||||
"""
|
||||
|
||||
noise = np.random.normal(mean, var**0.5, image.shape)
|
||||
out = image + 0.5 * noise
|
||||
out = np.clip(out, 0, 255)
|
||||
out = np.uint8(out)
|
||||
return out
|
||||
|
||||
|
||||
def get_crop(image):
|
||||
"""
|
||||
random crop
|
||||
"""
|
||||
h, w, _ = image.shape
|
||||
top_min = 1
|
||||
top_max = 8
|
||||
top_crop = int(random.randint(top_min, top_max))
|
||||
top_crop = min(top_crop, h - 1)
|
||||
crop_img = image.copy()
|
||||
ratio = random.randint(0, 1)
|
||||
if ratio:
|
||||
crop_img = crop_img[top_crop:h, :, :]
|
||||
else:
|
||||
crop_img = crop_img[0:h - top_crop, :, :]
|
||||
return crop_img
|
||||
|
||||
|
||||
class Config:
|
||||
"""
|
||||
Config
|
||||
"""
|
||||
|
||||
def __init__(self, use_tia):
|
||||
self.anglex = random.random() * 30
|
||||
self.angley = random.random() * 15
|
||||
self.anglez = random.random() * 10
|
||||
self.fov = 42
|
||||
self.r = 0
|
||||
self.shearx = random.random() * 0.3
|
||||
self.sheary = random.random() * 0.05
|
||||
self.borderMode = cv2.BORDER_REPLICATE
|
||||
self.use_tia = use_tia
|
||||
|
||||
def make(self, w, h, ang):
|
||||
"""
|
||||
make
|
||||
"""
|
||||
self.anglex = random.random() * 5 * flag()
|
||||
self.angley = random.random() * 5 * flag()
|
||||
self.anglez = -1 * random.random() * int(ang) * flag()
|
||||
self.fov = 42
|
||||
self.r = 0
|
||||
self.shearx = 0
|
||||
self.sheary = 0
|
||||
self.borderMode = cv2.BORDER_REPLICATE
|
||||
self.w = w
|
||||
self.h = h
|
||||
|
||||
self.perspective = self.use_tia
|
||||
self.stretch = self.use_tia
|
||||
self.distort = self.use_tia
|
||||
|
||||
self.crop = True
|
||||
self.affine = False
|
||||
self.reverse = True
|
||||
self.noise = True
|
||||
self.jitter = True
|
||||
self.blur = True
|
||||
self.color = True
|
||||
|
||||
|
||||
def rad(x):
|
||||
"""
|
||||
rad
|
||||
"""
|
||||
return x * np.pi / 180
|
||||
|
||||
|
||||
def get_warpR(config):
|
||||
"""
|
||||
get_warpR
|
||||
"""
|
||||
anglex, angley, anglez, fov, w, h, r = \
|
||||
config.anglex, config.angley, config.anglez, config.fov, config.w, config.h, config.r
|
||||
if w > 69 and w < 112:
|
||||
anglex = anglex * 1.5
|
||||
|
||||
z = np.sqrt(w**2 + h**2) / 2 / np.tan(rad(fov / 2))
|
||||
# Homogeneous coordinate transformation matrix
|
||||
rx = np.array([[1, 0, 0, 0],
|
||||
[0, np.cos(rad(anglex)), -np.sin(rad(anglex)), 0], [
|
||||
0,
|
||||
-np.sin(rad(anglex)),
|
||||
np.cos(rad(anglex)),
|
||||
0,
|
||||
], [0, 0, 0, 1]], np.float32)
|
||||
ry = np.array([[np.cos(rad(angley)), 0, np.sin(rad(angley)), 0],
|
||||
[0, 1, 0, 0], [
|
||||
-np.sin(rad(angley)),
|
||||
0,
|
||||
np.cos(rad(angley)),
|
||||
0,
|
||||
], [0, 0, 0, 1]], np.float32)
|
||||
rz = np.array([[np.cos(rad(anglez)), np.sin(rad(anglez)), 0, 0],
|
||||
[-np.sin(rad(anglez)), np.cos(rad(anglez)), 0, 0],
|
||||
[0, 0, 1, 0], [0, 0, 0, 1]], np.float32)
|
||||
r = rx.dot(ry).dot(rz)
|
||||
# generate 4 points
|
||||
pcenter = np.array([h / 2, w / 2, 0, 0], np.float32)
|
||||
p1 = np.array([0, 0, 0, 0], np.float32) - pcenter
|
||||
p2 = np.array([w, 0, 0, 0], np.float32) - pcenter
|
||||
p3 = np.array([0, h, 0, 0], np.float32) - pcenter
|
||||
p4 = np.array([w, h, 0, 0], np.float32) - pcenter
|
||||
dst1 = r.dot(p1)
|
||||
dst2 = r.dot(p2)
|
||||
dst3 = r.dot(p3)
|
||||
dst4 = r.dot(p4)
|
||||
list_dst = np.array([dst1, dst2, dst3, dst4])
|
||||
org = np.array([[0, 0], [w, 0], [0, h], [w, h]], np.float32)
|
||||
dst = np.zeros((4, 2), np.float32)
|
||||
# Project onto the image plane
|
||||
dst[:, 0] = list_dst[:, 0] * z / (z - list_dst[:, 2]) + pcenter[0]
|
||||
dst[:, 1] = list_dst[:, 1] * z / (z - list_dst[:, 2]) + pcenter[1]
|
||||
|
||||
warpR = cv2.getPerspectiveTransform(org, dst)
|
||||
|
||||
dst1, dst2, dst3, dst4 = dst
|
||||
r1 = int(min(dst1[1], dst2[1]))
|
||||
r2 = int(max(dst3[1], dst4[1]))
|
||||
c1 = int(min(dst1[0], dst3[0]))
|
||||
c2 = int(max(dst2[0], dst4[0]))
|
||||
|
||||
try:
|
||||
ratio = min(1.0 * h / (r2 - r1), 1.0 * w / (c2 - c1))
|
||||
|
||||
dx = -c1
|
||||
dy = -r1
|
||||
T1 = np.float32([[1., 0, dx], [0, 1., dy], [0, 0, 1.0 / ratio]])
|
||||
ret = T1.dot(warpR)
|
||||
except:
|
||||
ratio = 1.0
|
||||
T1 = np.float32([[1., 0, 0], [0, 1., 0], [0, 0, 1.]])
|
||||
ret = T1
|
||||
return ret, (-r1, -c1), ratio, dst
|
||||
|
||||
|
||||
def get_warpAffine(config):
|
||||
"""
|
||||
get_warpAffine
|
||||
"""
|
||||
anglez = config.anglez
|
||||
rz = np.array([[np.cos(rad(anglez)), np.sin(rad(anglez)), 0],
|
||||
[-np.sin(rad(anglez)), np.cos(rad(anglez)), 0]], np.float32)
|
||||
return rz
|
||||
|
||||
|
||||
def warp(img, ang, use_tia=True, prob=0.4):
|
||||
"""
|
||||
warp
|
||||
"""
|
||||
h, w, _ = img.shape
|
||||
config = Config(use_tia=use_tia)
|
||||
config.make(w, h, ang)
|
||||
new_img = img
|
||||
|
||||
if config.distort:
|
||||
img_height, img_width = img.shape[0:2]
|
||||
if random.random() <= prob and img_height >= 20 and img_width >= 20:
|
||||
new_img = tia_distort(new_img, random.randint(3, 6))
|
||||
|
||||
if config.stretch:
|
||||
img_height, img_width = img.shape[0:2]
|
||||
if random.random() <= prob and img_height >= 20 and img_width >= 20:
|
||||
new_img = tia_stretch(new_img, random.randint(3, 6))
|
||||
|
||||
if config.perspective:
|
||||
if random.random() <= prob:
|
||||
new_img = tia_perspective(new_img)
|
||||
|
||||
if config.crop:
|
||||
img_height, img_width = img.shape[0:2]
|
||||
if random.random() <= prob and img_height >= 20 and img_width >= 20:
|
||||
new_img = get_crop(new_img)
|
||||
|
||||
if config.blur:
|
||||
if random.random() <= prob:
|
||||
new_img = blur(new_img)
|
||||
if config.color:
|
||||
if random.random() <= prob:
|
||||
new_img = cvtColor(new_img)
|
||||
if config.jitter:
|
||||
new_img = jitter(new_img)
|
||||
if config.noise:
|
||||
if random.random() <= prob:
|
||||
new_img = add_gasuss_noise(new_img)
|
||||
if config.reverse:
|
||||
if random.random() <= prob:
|
||||
new_img = 255 - new_img
|
||||
return new_img
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user