├── .gitignore ├── LICENSE.txt ├── NOTICE.txt ├── README.md ├── benchmark.py ├── bin ├── benchmark_run.sh ├── docker_run.sh ├── exportCPU.sh ├── exportGPU.sh ├── get_test_images_run.sh ├── image_demo_run.sh ├── inspect_saved_model.sh ├── upgrade-tf-v2.sh ├── video_demo_run.sh └── webcam_demo_run.sh ├── docker ├── Dockerfile └── docker_img_build.sh ├── get_test_images.py ├── image_demo.py ├── posenet ├── __init__.py ├── base_model.py ├── constants.py ├── converter │ ├── common.py │ ├── config.py │ ├── config.yaml │ ├── tfjs2tf.py │ └── tfjsdownload.py ├── decode.py ├── decode_multi.py ├── mobilenet.py ├── posenet.py ├── posenet_factory.py ├── resnet.py └── utils.py ├── requirements.txt ├── video_demo.py └── webcam_demo.py /.gitignore: -------------------------------------------------------------------------------- 1 | images/* 2 | output/* 3 | output_old/* 4 | .idea/* 5 | .idea 6 | _models/* 7 | _tf_models/* 8 | _tfjs_models/* 9 | _posenet_weights/* 10 | docker/requirements.txt 11 | *.mp4 12 | 13 | # Byte-compiled / optimized / DLL files 14 | __pycache__/ 15 | *.py[cod] 16 | *$py.class 17 | 18 | # C extensions 19 | *.so 20 | 21 | # Distribution / packaging 22 | .Python 23 | build/ 24 | develop-eggs/ 25 | dist/ 26 | downloads/ 27 | eggs/ 28 | .eggs/ 29 | lib/ 30 | lib64/ 31 | parts/ 32 | sdist/ 33 | var/ 34 | wheels/ 35 | pip-wheel-metadata/ 36 | share/python-wheels/ 37 | *.egg-info/ 38 | .installed.cfg 39 | *.egg 40 | MANIFEST 41 | 42 | # PyInstaller 43 | # Usually these files are written by a python script from a template 44 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 45 | *.manifest 46 | *.spec 47 | 48 | # Installer logs 49 | pip-log.txt 50 | pip-delete-this-directory.txt 51 | 52 | # Unit test / coverage reports 53 | htmlcov/ 54 | .tox/ 55 | .nox/ 56 | .coverage 57 | .coverage.* 58 | .cache 59 | nosetests.xml 60 | coverage.xml 61 | *.cover 62 | .hypothesis/ 63 | .pytest_cache/ 64 | 65 | # Translations 66 | *.mo 67 | *.pot 68 | 69 | # Django stuff: 70 | *.log 71 | local_settings.py 72 | db.sqlite3 73 | 74 | # Flask stuff: 75 | instance/ 76 | .webassets-cache 77 | 78 | # Scrapy stuff: 79 | .scrapy 80 | 81 | # Sphinx documentation 82 | docs/_build/ 83 | 84 | # PyBuilder 85 | target/ 86 | 87 | # Jupyter Notebook 88 | .ipynb_checkpoints 89 | 90 | # IPython 91 | profile_default/ 92 | ipython_config.py 93 | 94 | # pyenv 95 | .python-version 96 | 97 | # pipenv 98 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 99 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 100 | # having no cross-platform support, pipenv may install dependencies that don’t work, or not 101 | # install all needed dependencies. 102 | #Pipfile.lock 103 | 104 | # celery beat schedule file 105 | celerybeat-schedule 106 | 107 | # SageMath parsed files 108 | *.sage.py 109 | 110 | # Environments 111 | .env 112 | .venv 113 | env/ 114 | venv/ 115 | ENV/ 116 | env.bak/ 117 | venv.bak/ 118 | 119 | # Spyder project settings 120 | .spyderproject 121 | .spyproject 122 | 123 | # Rope project settings 124 | .ropeproject 125 | 126 | # mkdocs documentation 127 | /site 128 | 129 | # mypy 130 | .mypy_cache/ 131 | .dmypy.json 132 | dmypy.json 133 | 134 | # Pyre type checker 135 | .pyre/ 136 | 137 | 138 | # Covers JetBrains IDEs: IntelliJ, RubyMine, PhpStorm, AppCode, PyCharm, CLion, Android Studio and WebStorm 139 | # Reference: https://intellij-support.jetbrains.com/hc/en-us/articles/206544839 140 | 141 | # User-specific stuff 142 | .idea/**/workspace.xml 143 | .idea/**/tasks.xml 144 | .idea/**/usage.statistics.xml 145 | .idea/**/dictionaries 146 | .idea/**/shelf 147 | 148 | # Generated files 149 | .idea/**/contentModel.xml 150 | 151 | # Sensitive or high-churn files 152 | .idea/**/dataSources/ 153 | .idea/**/dataSources.ids 154 | .idea/**/dataSources.local.xml 155 | .idea/**/sqlDataSources.xml 156 | .idea/**/dynamic.xml 157 | .idea/**/uiDesigner.xml 158 | .idea/**/dbnavigator.xml 159 | 160 | # Gradle 161 | .idea/**/gradle.xml 162 | .idea/**/libraries 163 | 164 | # Gradle and Maven with auto-import 165 | # When using Gradle or Maven with auto-import, you should exclude module files, 166 | # since they will be recreated, and may cause churn. Uncomment if using 167 | # auto-import. 168 | .idea/modules.xml 169 | .idea/*.iml 170 | .idea/modules 171 | 172 | # CMake 173 | cmake-build-*/ 174 | 175 | # Mongo Explorer plugin 176 | .idea/**/mongoSettings.xml 177 | 178 | # File-based project format 179 | *.iws 180 | 181 | # IntelliJ 182 | out/ 183 | output.txt 184 | 185 | # mpeltonen/sbt-idea plugin 186 | .idea_modules/ 187 | 188 | # JIRA plugin 189 | atlassian-ide-plugin.xml 190 | 191 | # Cursive Clojure plugin 192 | .idea/replstate.xml 193 | 194 | # Crashlytics plugin (for Android Studio and IntelliJ) 195 | com_crashlytics_export_strings.xml 196 | crashlytics.properties 197 | crashlytics-build.properties 198 | fabric.properties 199 | 200 | # Editor-based Rest Client 201 | .idea/httpRequests 202 | 203 | # Android studio 3.1+ serialized cache file 204 | .idea/caches/build_file_checksums.ser 205 | -------------------------------------------------------------------------------- /LICENSE.txt: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS 177 | 178 | APPENDIX: How to apply the Apache License to your work. 179 | 180 | To apply the Apache License to your work, attach the following 181 | boilerplate notice, with the fields enclosed by brackets "[]" 182 | replaced with your own identifying information. (Don't include 183 | the brackets!) The text should be enclosed in the appropriate 184 | comment syntax for the file format. We also recommend that a 185 | file or class name and description of purpose be included on the 186 | same "printed page" as the copyright notice for easier 187 | identification within third-party archives. 188 | 189 | Copyright 2018 Ross Wightman 190 | 191 | Licensed under the Apache License, Version 2.0 (the "License"); 192 | you may not use this file except in compliance with the License. 193 | You may obtain a copy of the License at 194 | 195 | http://www.apache.org/licenses/LICENSE-2.0 196 | 197 | Unless required by applicable law or agreed to in writing, software 198 | distributed under the License is distributed on an "AS IS" BASIS, 199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 200 | See the License for the specific language governing permissions and 201 | limitations under the License. -------------------------------------------------------------------------------- /NOTICE.txt: -------------------------------------------------------------------------------- 1 | PoseNet Python 2 | Copyright 2018 Ross Wightman 3 | 4 | PoseNet Python numerous refactorings 5 | Copyright 2020 Peter Rigole 6 | 7 | tfjs PoseNet weights and original JS code 8 | Copyright 2018 Google LLC. All Rights Reserved. 9 | (https://github.com/tensorflow/tfjs-models | Apache License 2.0) 10 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | ## PoseNet Python 2 | 3 | This repository originates from [rwightman/posenet-python](https://github.com/rwightman/posenet-python) and has been 4 | heavily refactored to: 5 | * make it run the posenet v2 networks 6 | * get it to work with the latest tfjs graph serialization 7 | * extend it with the ResNet50 network 8 | * make the code run on TF 2.x 9 | * get all code running in docker containers for ease of use and installation (no conda necessary) 10 | 11 | This repository contains a pure Python implementation (multi-pose only) of the Google TensorFlow.js Posenet model. 12 | For a (slightly faster) PyTorch implementation that followed from this, 13 | see (https://github.com/rwightman/posenet-pytorch) 14 | 15 | 16 | ### Install 17 | 18 | A suitable Python 3.x environment with Tensorflow 2.x. For a quick setup, use docker. 19 | 20 | If you want to use the webcam demo, a pip version of opencv (`pip install opencv-python`) is required instead of 21 | the conda version. Anaconda's default opencv does not include ffpmeg/VideoCapture support. Also, you may have to 22 | force install version 3.4.x as 4.x has a broken drawKeypoints binding. 23 | 24 | Have a look at the docker configuration for a quick setup. If you want conda, have a look at the `requirements.txt` 25 | file to see what you should install. Know that we rely on https://github.com/patlevin/tfjs-to-tf for 26 | converting the tensorflow.js serialization to the tensorflow saved model. So you have to install this package: 27 | 28 | ```bash 29 | git clone https://github.com/patlevin/tfjs-to-tf.git 30 | cd tfjs-to-tf 31 | pip install . --no-deps 32 | ``` 33 | 34 | Use the `--no-deps` flag to prevent tfjs-to-tf from installing Tensorflow 1.x as this would uninstall your 35 | Tensorflow 2.x! 36 | 37 | 38 | ### Using Docker 39 | 40 | A convenient way to run this project is by building and running the docker image, because it has all the requirements 41 | built-in. 42 | The GPU version is tested on a Linux machine. You need to install the nvidia host driver and the nvidia-docker toolkit. 43 | Once set up, you can make as many images as you want with different dependencies without touching your host OS 44 | (or fiddling with conda). 45 | 46 | If you just want to test this code, you can run everything on a CPU just as well. You still get 8fps on mobilenet and 47 | 4fps on resnet50. Replace `GPU` below with `CPU` to test on a CPU. 48 | 49 | ```bash 50 | cd docker 51 | ./docker_img_build.sh GPU 52 | cd .. 53 | . ./bin/exportGPU.sh 54 | ./bin/get_test_images_run.sh 55 | ./bin/image_demo_run.sh 56 | ``` 57 | 58 | Some pointers to get you going on the Linux machine setup. Most links are based on Ubuntu, but other distributions 59 | should work fine as well. 60 | * [Install docker](https://docs.docker.com/install/linux/docker-ce/ubuntu/ ) 61 | * [Install the NVIDIA host driver](https://docs.nvidia.com/cuda/cuda-installation-guide-linux/index.html#ubuntu-installation) 62 | * remember to reboot here 63 | * [Install the NVIDIA Container Toolkit](https://github.com/NVIDIA/nvidia-docker) 64 | * check your installation: `docker run --gpus all nvidia/cuda nvidia-smi` 65 | 66 | 67 | ### Usage 68 | 69 | There are three demo apps in the root that utilize the PoseNet model. They are very basic and could definitely be 70 | improved. 71 | 72 | The first time these apps are run (or the library is used) model weights will be downloaded from the TensorFlow.js 73 | version and converted on the fly. 74 | 75 | #### image_demo.py 76 | 77 | Image demo runs inference on an input folder of images and outputs those images with the keypoints and skeleton 78 | overlayed. 79 | 80 | `python image_demo.py --model resnet50 --stride 16 --image_dir ./images --output_dir ./output` 81 | 82 | A folder of suitable test images can be downloaded by first running the `get_test_images.py` script. 83 | 84 | #### benchmark.py 85 | 86 | A minimal performance benchmark based on image_demo. Images in `--image_dir` are pre-loaded and inference is 87 | run `--num_images` times with no drawing and no text output. 88 | 89 | Running the benchmark cycling 1000 times through the example images on a Geforce GTX 1080ti gives these average FPS 90 | using TF 2.0.0: 91 | 92 | ``` 93 | ResNet50 stride 16: 32.41 FPS 94 | ResNet50 stride 32: 38.70 FPS 95 | MobileNet stride 8: 37.90 FPS (this is surprisingly slow for mobilenet, ran this several times, same result) 96 | MobileNet stride 16: 58.64 FPS 97 | ``` 98 | 99 | Faster FPS have been reported by Ross Wightmann on the original codebase in 100 | [rwightman/posenet-python](https://github.com/rwightman/posenet-python), so if anyone has a pull request that 101 | improves the performance of this codebase, feel free to let me know! 102 | 103 | #### webcam_demo.py 104 | 105 | The webcam demo uses OpenCV to capture images from a connected webcam. The result is overlayed with the keypoints and 106 | skeletons and rendered to the screen. The default args for the webcam_demo assume device_id=0 for the camera and 107 | that 1280x720 resolution is possible. 108 | 109 | ### Credits 110 | 111 | The original model, weights, code, etc. was created by Google and can be found at 112 | https://github.com/tensorflow/tfjs-models/tree/master/posenet 113 | 114 | This port is initially created by Ross Wightman and later upgraded by Peter Rigole and is in no way related to Google. 115 | 116 | The Python conversion code that started me on my way was adapted from the CoreML port at 117 | https://github.com/infocom-tpo/PoseNet-CoreML 118 | 119 | ### TODO 120 | * Performance improvements (especially edge loops in 'decode.py') 121 | * OpenGL rendering/drawing 122 | * Comment interfaces, tensor dimensions, etc 123 | -------------------------------------------------------------------------------- /benchmark.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import cv2 3 | import time 4 | import argparse 5 | import os 6 | from posenet.posenet_factory import load_model 7 | 8 | 9 | parser = argparse.ArgumentParser() 10 | parser.add_argument('--model', type=str, default='resnet50') # mobilenet resnet50 11 | parser.add_argument('--stride', type=int, default=16) # 8, 16, 32 (max 16 for mobilenet) 12 | parser.add_argument('--quant_bytes', type=int, default=4) # 4 = float 13 | parser.add_argument('--multiplier', type=float, default=1.0) # only for mobilenet 14 | parser.add_argument('--image_dir', type=str, default='./images') 15 | parser.add_argument('--num_images', type=int, default=1000) 16 | args = parser.parse_args() 17 | 18 | 19 | def main(): 20 | 21 | print('Tensorflow version: %s' % tf.__version__) 22 | assert tf.__version__.startswith('2.'), "Tensorflow version 2.x must be used!" 23 | 24 | model = args.model # mobilenet resnet50 25 | stride = args.stride # 8, 16, 32 (max 16 for mobilenet) 26 | quant_bytes = args.quant_bytes # float 27 | multiplier = args.multiplier # only for mobilenet 28 | 29 | posenet = load_model(model, stride, quant_bytes, multiplier) 30 | 31 | num_images = args.num_images 32 | filenames = [ 33 | f.path for f in os.scandir(args.image_dir) if f.is_file() and f.path.endswith(('.png', '.jpg'))] 34 | if len(filenames) > num_images: 35 | filenames = filenames[:num_images] 36 | 37 | images = {f: cv2.imread(f) for f in filenames} 38 | 39 | start = time.time() 40 | for i in range(num_images): 41 | image = images[filenames[i % len(filenames)]] 42 | posenet.estimate_multiple_poses(image) 43 | 44 | print('Average FPS:', num_images / (time.time() - start)) 45 | 46 | 47 | if __name__ == "__main__": 48 | main() 49 | -------------------------------------------------------------------------------- /bin/benchmark_run.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | ./bin/docker_run.sh python benchmark.py --model mobilenet --stride 16 --image_dir ./images --num_images 1000 4 | -------------------------------------------------------------------------------- /bin/docker_run.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | WORK=$(pwd) 4 | 5 | if [ -z "$POSENET_PYTHON_DEVICE" ]; then 6 | echo "set the environment variable POSENET_PYTHON_DEVICE to CPU or GPU, or enter your choice below:" 7 | read -p "Enter your device (CPU or GPU): " device 8 | if [ "$device" = "GPU" ]; then 9 | source exportGPU.sh 10 | elif [ "$device" = "CPU" ]; then 11 | source exportCPU.sh 12 | else 13 | echo "Device configuration failed..." 14 | exit 1 15 | fi 16 | fi 17 | 18 | echo "device is: $POSENET_PYTHON_DEVICE" 19 | 20 | if [ "$POSENET_PYTHON_DEVICE" = "GPU" ]; then 21 | image="posenet-python-gpu" 22 | gpu_opts="--gpus all" 23 | else 24 | image="posenet-python-cpu" 25 | gpu_opts="" 26 | fi 27 | 28 | docker run $gpu_opts -it --rm -v $WORK:/work "$image" "$@" 29 | -------------------------------------------------------------------------------- /bin/exportCPU.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | # source this file to set your environment on a CPU device 3 | # $ . exportCPU.sh 4 | export POSENET_PYTHON_DEVICE=CPU 5 | -------------------------------------------------------------------------------- /bin/exportGPU.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | # source this file to set your environment on a GPU device 3 | # $ . exportGPU.sh 4 | export POSENET_PYTHON_DEVICE=GPU -------------------------------------------------------------------------------- /bin/get_test_images_run.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | ./bin/docker_run.sh python get_test_images.py 4 | -------------------------------------------------------------------------------- /bin/image_demo_run.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | ./bin/docker_run.sh python image_demo.py --model resnet50 --stride 16 --image_dir ./images --output_dir ./output 4 | -------------------------------------------------------------------------------- /bin/inspect_saved_model.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | FOLDER=$1 4 | 5 | # e.g.: $> ./inspect_saved_model.sh _tf_models/posenet/mobilenet_v1_100/stride16 6 | ./bin/docker_run.sh saved_model_cli show --dir "$FOLDER" --all 7 | -------------------------------------------------------------------------------- /bin/upgrade-tf-v2.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | # run this from the top-level folder of the project 4 | 5 | WORK=$(dirname $(pwd)) 6 | 7 | docker run --gpus all -it -v $WORK:/work posenet-python tf_upgrade_v2 \ 8 | --intree posenet-python/ \ 9 | --outtree posenet-python_v2/ \ 10 | --reportfile posenet-python/report.txt 11 | -------------------------------------------------------------------------------- /bin/video_demo_run.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | #./bin/docker_run.sh python video_demo.py --model resnet50 --stride 16 --input_file "Pexels Videos 3552510.mp4" --output_file "Pexels Videos 3552510-with_pose.mp4" 4 | ./bin/docker_run.sh python video_demo.py --model resnet50 --stride 16 --input_file "exki.mp4" --output_file "exki_with_pose.mp4" 5 | ./bin/docker_run.sh python video_demo.py --model resnet50 --stride 16 --input_file "night-bridge.mp4" --output_file "night-bridge_with_pose.mp4" 6 | ./bin/docker_run.sh python video_demo.py --model resnet50 --stride 16 --input_file "night-colorful.mp4" --output_file "night-colorful_with_pose.mp4" 7 | ./bin/docker_run.sh python video_demo.py --model resnet50 --stride 16 --input_file "night-street.mp4" --output_file "night-street_with_pose.mp4" 8 | ./bin/docker_run.sh python video_demo.py --model resnet50 --stride 16 --input_file "pedestrians.mp4" --output_file "pedestrians_with_pose.mp4" 9 | ./bin/docker_run.sh python video_demo.py --model resnet50 --stride 16 --input_file "sidewalk.mp4" --output_file "sidewalk_with_pose.mp4" 10 | -------------------------------------------------------------------------------- /bin/webcam_demo_run.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | ./bin/docker_run.sh python webcam_demo.py --model resnet50 --stride 16 --image_dir ./images --output_dir ./output 4 | -------------------------------------------------------------------------------- /docker/Dockerfile: -------------------------------------------------------------------------------- 1 | # default image version, override using --build-arg IMAGE_VERSION=otherversion 2 | ARG IMAGE_VERSION=2.1.0-py3-jupyter 3 | FROM tensorflow/tensorflow:$IMAGE_VERSION 4 | # The default version is the CPU version! 5 | # see: https://www.tensorflow.org/install/docker 6 | # see: https://hub.docker.com/r/tensorflow/tensorflow/ 7 | 8 | # Install system packages 9 | RUN apt-get update && apt-get install -y --no-install-recommends \ 10 | bzip2 \ 11 | git \ 12 | wget && \ 13 | pip install --upgrade pip && \ 14 | rm -rf /var/lib/apt/lists/* 15 | 16 | COPY requirements.txt /work/ 17 | 18 | WORKDIR /work 19 | 20 | # run pip install with the '--no-deps' argument, to avoid that tensorflowjs installs an old version of tensorflow! 21 | # It also ensures that we know and controll the transitive dependencies (although the tensorflow docker image comes 22 | # with a lot of packages pre-installed). 23 | RUN pip install -r requirements.txt --no-deps 24 | 25 | RUN git clone https://github.com/patlevin/tfjs-to-tf.git && \ 26 | cd tfjs-to-tf && \ 27 | git checkout v0.3.0 && \ 28 | pip install . --no-deps && \ 29 | cd .. && \ 30 | rm -r tfjs-to-tf 31 | 32 | ENV PYTHONPATH='/work/:$PYTHONPATH' 33 | 34 | CMD ["bash"] 35 | -------------------------------------------------------------------------------- /docker/docker_img_build.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | if [ -z "$1" ]; then 4 | echo "pass CPU or GPU as argument" 5 | echo "Docker image build failed..." 6 | exit 1 7 | fi 8 | 9 | if [ "$1" = "GPU" ]; then 10 | image="posenet-python-gpu" 11 | # version="--build-arg IMAGE_VERSION=2.1.0rc2-gpu-py3-jupyter" 12 | version="--build-arg IMAGE_VERSION=2.0.0-gpu-py3-jupyter" 13 | else 14 | image="posenet-python-cpu" 15 | version="--build-arg IMAGE_VERSION=2.1.0-py3-jupyter" 16 | fi 17 | 18 | cp ../requirements.txt . 19 | 20 | docker rmi -f "$image" 21 | 22 | docker build -t "$image" $version . 23 | -------------------------------------------------------------------------------- /get_test_images.py: -------------------------------------------------------------------------------- 1 | import urllib.request 2 | import os 3 | import argparse 4 | 5 | GOOGLE_CLOUD_IMAGE_BUCKET = 'https://storage.googleapis.com/tfjs-models/assets/posenet/' 6 | 7 | TEST_IMAGES = [ 8 | 'frisbee.jpg', 9 | 'frisbee_2.jpg', 10 | 'backpackman.jpg', 11 | 'boy_doughnut.jpg', 12 | 'soccer.png', 13 | 'with_computer.jpg', 14 | 'snowboard.jpg', 15 | 'person_bench.jpg', 16 | 'skiing.jpg', 17 | 'fire_hydrant.jpg', 18 | 'kyte.jpg', 19 | 'looking_at_computer.jpg', 20 | 'tennis.jpg', 21 | 'tennis_standing.jpg', 22 | 'truck.jpg', 23 | 'on_bus.jpg', 24 | 'tie_with_beer.jpg', 25 | 'baseball.jpg', 26 | 'multi_skiing.jpg', 27 | 'riding_elephant.jpg', 28 | 'skate_park_venice.jpg', 29 | 'skate_park.jpg', 30 | 'tennis_in_crowd.jpg', 31 | 'two_on_bench.jpg', 32 | ] 33 | 34 | parser = argparse.ArgumentParser() 35 | parser.add_argument('--image_dir', type=str, default='./images') 36 | args = parser.parse_args() 37 | 38 | 39 | def main(): 40 | if not os.path.exists(args.image_dir): 41 | os.makedirs(args.image_dir) 42 | 43 | for f in TEST_IMAGES: 44 | url = os.path.join(GOOGLE_CLOUD_IMAGE_BUCKET, f) 45 | print('Downloading %s' % f) 46 | urllib.request.urlretrieve(url, os.path.join(args.image_dir, f)) 47 | 48 | 49 | if __name__ == "__main__": 50 | main() 51 | -------------------------------------------------------------------------------- /image_demo.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import cv2 3 | import time 4 | import argparse 5 | import os 6 | from posenet.posenet_factory import load_model 7 | 8 | parser = argparse.ArgumentParser() 9 | parser.add_argument('--model', type=str, default='resnet50') # mobilenet resnet50 10 | parser.add_argument('--stride', type=int, default=16) # 8, 16, 32 (max 16 for mobilenet) 11 | parser.add_argument('--quant_bytes', type=int, default=4) # 4 = float 12 | parser.add_argument('--multiplier', type=float, default=1.0) # only for mobilenet 13 | parser.add_argument('--notxt', action='store_true') 14 | parser.add_argument('--image_dir', type=str, default='./images') 15 | parser.add_argument('--output_dir', type=str, default='./output') 16 | args = parser.parse_args() 17 | 18 | 19 | def main(): 20 | 21 | print('Tensorflow version: %s' % tf.__version__) 22 | assert tf.__version__.startswith('2.'), "Tensorflow version 2.x must be used!" 23 | 24 | if args.output_dir: 25 | if not os.path.exists(args.output_dir): 26 | os.makedirs(args.output_dir) 27 | 28 | model = args.model # mobilenet resnet50 29 | stride = args.stride # 8, 16, 32 (max 16 for mobilenet, min 16 for resnet50) 30 | quant_bytes = args.quant_bytes # float 31 | multiplier = args.multiplier # only for mobilenet 32 | 33 | posenet = load_model(model, stride, quant_bytes, multiplier) 34 | 35 | filenames = [f.path for f in os.scandir(args.image_dir) if f.is_file() and f.path.endswith(('.png', '.jpg'))] 36 | 37 | start = time.time() 38 | for f in filenames: 39 | img = cv2.imread(f) 40 | pose_scores, keypoint_scores, keypoint_coords = posenet.estimate_multiple_poses(img) 41 | img_poses = posenet.draw_poses(img, pose_scores, keypoint_scores, keypoint_coords) 42 | posenet.print_scores(f, pose_scores, keypoint_scores, keypoint_coords) 43 | cv2.imwrite(os.path.join(args.output_dir, os.path.relpath(f, args.image_dir)), img_poses) 44 | 45 | print('Average FPS:', len(filenames) / (time.time() - start)) 46 | 47 | 48 | if __name__ == "__main__": 49 | main() 50 | -------------------------------------------------------------------------------- /posenet/__init__.py: -------------------------------------------------------------------------------- 1 | from posenet.constants import * 2 | from posenet.decode_multi import decode_multiple_poses 3 | from posenet.utils import * 4 | from posenet import converter 5 | -------------------------------------------------------------------------------- /posenet/base_model.py: -------------------------------------------------------------------------------- 1 | from abc import ABC, abstractmethod 2 | import tensorflow as tf 3 | 4 | 5 | class BaseModel(ABC): 6 | 7 | # keys for the output_tensor_names map 8 | HEATMAP_KEY = "heatmap" 9 | OFFSETS_KEY = "offsets" 10 | DISPLACEMENT_FWD_KEY = "displacement_fwd" 11 | DISPLACEMENT_BWD_KEY = "displacement_bwd" 12 | 13 | def __init__(self, model_function, output_tensor_names, output_stride): 14 | self.output_stride = output_stride 15 | self.output_tensor_names = output_tensor_names 16 | self.model_function = model_function 17 | 18 | def valid_resolution(self, width, height): 19 | # calculate closest smaller width and height that is divisible by the stride after subtracting 1 (for the bias?) 20 | target_width = (int(width) // self.output_stride) * self.output_stride + 1 21 | target_height = (int(height) // self.output_stride) * self.output_stride + 1 22 | return target_width, target_height 23 | 24 | @abstractmethod 25 | def preprocess_input(self, image): 26 | pass 27 | 28 | def predict(self, image): 29 | input_image, image_scale = self.preprocess_input(image) 30 | 31 | input_image = tf.convert_to_tensor(input_image, dtype=tf.float32) 32 | 33 | result = self.model_function(input_image) 34 | 35 | heatmap_result = result[self.output_tensor_names[self.HEATMAP_KEY]] 36 | offsets_result = result[self.output_tensor_names[self.OFFSETS_KEY]] 37 | displacement_fwd_result = result[self.output_tensor_names[self.DISPLACEMENT_FWD_KEY]] 38 | displacement_bwd_result = result[self.output_tensor_names[self.DISPLACEMENT_BWD_KEY]] 39 | 40 | return tf.sigmoid(heatmap_result), offsets_result, displacement_fwd_result, displacement_bwd_result, image_scale 41 | -------------------------------------------------------------------------------- /posenet/constants.py: -------------------------------------------------------------------------------- 1 | 2 | PART_NAMES = [ 3 | "nose", "leftEye", "rightEye", "leftEar", "rightEar", "leftShoulder", 4 | "rightShoulder", "leftElbow", "rightElbow", "leftWrist", "rightWrist", 5 | "leftHip", "rightHip", "leftKnee", "rightKnee", "leftAnkle", "rightAnkle" 6 | ] 7 | 8 | NUM_KEYPOINTS = len(PART_NAMES) 9 | 10 | PART_IDS = {pn: pid for pid, pn in enumerate(PART_NAMES)} 11 | 12 | CONNECTED_PART_NAMES = [ 13 | ("leftHip", "leftShoulder"), ("leftElbow", "leftShoulder"), 14 | ("leftElbow", "leftWrist"), ("leftHip", "leftKnee"), 15 | ("leftKnee", "leftAnkle"), ("rightHip", "rightShoulder"), 16 | ("rightElbow", "rightShoulder"), ("rightElbow", "rightWrist"), 17 | ("rightHip", "rightKnee"), ("rightKnee", "rightAnkle"), 18 | ("leftShoulder", "rightShoulder"), ("leftHip", "rightHip") 19 | ] 20 | 21 | CONNECTED_PART_INDICES = [(PART_IDS[a], PART_IDS[b]) for a, b in CONNECTED_PART_NAMES] 22 | 23 | LOCAL_MAXIMUM_RADIUS = 1 24 | 25 | POSE_CHAIN = [ 26 | ("nose", "leftEye"), ("leftEye", "leftEar"), ("nose", "rightEye"), 27 | ("rightEye", "rightEar"), ("nose", "leftShoulder"), 28 | ("leftShoulder", "leftElbow"), ("leftElbow", "leftWrist"), 29 | ("leftShoulder", "leftHip"), ("leftHip", "leftKnee"), 30 | ("leftKnee", "leftAnkle"), ("nose", "rightShoulder"), 31 | ("rightShoulder", "rightElbow"), ("rightElbow", "rightWrist"), 32 | ("rightShoulder", "rightHip"), ("rightHip", "rightKnee"), 33 | ("rightKnee", "rightAnkle") 34 | ] 35 | 36 | PARENT_CHILD_TUPLES = [(PART_IDS[parent], PART_IDS[child]) for parent, child in POSE_CHAIN] 37 | 38 | PART_CHANNELS = [ 39 | 'left_face', 40 | 'right_face', 41 | 'right_upper_leg_front', 42 | 'right_lower_leg_back', 43 | 'right_upper_leg_back', 44 | 'left_lower_leg_front', 45 | 'left_upper_leg_front', 46 | 'left_upper_leg_back', 47 | 'left_lower_leg_back', 48 | 'right_feet', 49 | 'right_lower_leg_front', 50 | 'left_feet', 51 | 'torso_front', 52 | 'torso_back', 53 | 'right_upper_arm_front', 54 | 'right_upper_arm_back', 55 | 'right_lower_arm_back', 56 | 'left_lower_arm_front', 57 | 'left_upper_arm_front', 58 | 'left_upper_arm_back', 59 | 'left_lower_arm_back', 60 | 'right_hand', 61 | 'right_lower_arm_front', 62 | 'left_hand' 63 | ] -------------------------------------------------------------------------------- /posenet/converter/common.py: -------------------------------------------------------------------------------- 1 | TFJS_OP_KEY = 'op' 2 | TFJS_DILATIONS_KEY = 'dilations' 3 | TFJS_CONV2D_KEY = 'Conv2D' 4 | 5 | TF_NHWC = 'NHWC' 6 | TF_NCHW = 'NCHW' 7 | -------------------------------------------------------------------------------- /posenet/converter/config.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | BASE_DIR = os.path.dirname(__file__) 4 | TFJS_MODEL_DIR = './_tfjs_models' 5 | TF_MODEL_DIR = './_tf_models' 6 | 7 | MOBILENET_BASE_URL = 'https://storage.googleapis.com/tfjs-models/savedmodel/posenet/mobilenet/' 8 | RESNET50_BASE_URL = 'https://storage.googleapis.com/tfjs-models/savedmodel/posenet/resnet50/' 9 | 10 | POSENET_ARCHITECTURE = 'posenet' 11 | 12 | RESNET50_MODEL = 'resnet50' 13 | MOBILENET_MODEL = 'mobilenet' 14 | 15 | 16 | def bodypix_resnet50_config(stride, quant_bytes=4): 17 | 18 | graph_json = 'model-stride' + str(stride) + '.json' 19 | 20 | # quantBytes = 4 corresponding to the non - quantized full - precision checkpoints. 21 | if quant_bytes == 4: 22 | base_url = RESNET50_BASE_URL + 'float' 23 | model_dir = RESNET50_MODEL + '_float' 24 | else: 25 | base_url = RESNET50_BASE_URL + 'quant' + str(quant_bytes) + '/' 26 | model_dir = RESNET50_MODEL + '_quant' + str(quant_bytes) 27 | 28 | stride_dir = 'stride' + str(stride) 29 | 30 | return { 31 | 'base_url': base_url, 32 | 'filename': graph_json, 33 | 'output_stride': stride, 34 | 'data_format': 'NHWC', 35 | 'input_tensors': { 36 | 'image': 'sub_2:0' 37 | }, 38 | 'output_tensors': { 39 | 'heatmap': 'float_heatmaps:0', 40 | 'offsets': 'float_short_offsets:0', 41 | 'displacement_fwd': 'resnet_v1_50/displacement_fwd_2/BiasAdd:0', 42 | 'displacement_bwd': 'resnet_v1_50/displacement_bwd_2/BiasAdd:0' 43 | }, 44 | 'tfjs_dir': os.path.join(TFJS_MODEL_DIR, POSENET_ARCHITECTURE, model_dir, stride_dir), 45 | 'tf_dir': os.path.join(TF_MODEL_DIR, POSENET_ARCHITECTURE, model_dir, stride_dir) 46 | } 47 | 48 | 49 | def bodypix_mobilenet_config(stride, quant_bytes=4, multiplier=1.0): 50 | 51 | graph_json = 'model-stride' + str(stride) + '.json' 52 | 53 | multiplier_map = { 54 | 1.0: "100", 55 | 0.75: "075", 56 | 0.5: "050" 57 | } 58 | 59 | # quantBytes = 4 corresponding to the non - quantized full - precision checkpoints. 60 | if quant_bytes == 4: 61 | base_url = MOBILENET_BASE_URL + 'float/' + multiplier_map[multiplier] + '/' 62 | model_dir = MOBILENET_MODEL + '_float_' + multiplier_map[multiplier] 63 | else: 64 | base_url = MOBILENET_BASE_URL + 'quant' + str(quant_bytes) + '/' + multiplier_map[multiplier] + '/' 65 | model_dir = MOBILENET_MODEL + '_quant' + str(quant_bytes) + '_' + multiplier_map[multiplier] 66 | 67 | stride_dir = 'stride' + str(stride) 68 | 69 | return { 70 | 'base_url': base_url, 71 | 'filename': graph_json, 72 | 'output_stride': stride, 73 | 'data_format': 'NHWC', 74 | 'input_tensors': { 75 | 'image': 'sub_2:0' 76 | }, 77 | 'output_tensors': { 78 | 'heatmap': 'MobilenetV1/heatmap_2/BiasAdd:0', 79 | 'offsets': 'MobilenetV1/offset_2/BiasAdd:0', 80 | 'displacement_fwd': 'MobilenetV1/displacement_fwd_2/BiasAdd:0', 81 | 'displacement_bwd': 'MobilenetV1/displacement_bwd_2/BiasAdd:0' 82 | }, 83 | 'tfjs_dir': os.path.join(TFJS_MODEL_DIR, POSENET_ARCHITECTURE, model_dir, stride_dir), 84 | 'tf_dir': os.path.join(TF_MODEL_DIR, POSENET_ARCHITECTURE, model_dir, stride_dir) 85 | } 86 | -------------------------------------------------------------------------------- /posenet/converter/config.yaml: -------------------------------------------------------------------------------- 1 | # This config file is no longer used. 2 | # It is left in place as a reference for future integration of the bodypix architecture. 3 | models: 4 | bodypix: 5 | resnet50_v1: 6 | base_url: 'https://storage.googleapis.com/tfjs-models/savedmodel/bodypix/resnet50/float' 7 | model_variant: 8 | stride16: 9 | filename: 'model-stride16.json' 10 | output_stride: 16 11 | data_format: 'NHWC' 12 | input_tensors: 13 | image: 'sub_2:0' 14 | output_tensors: 15 | heatmap: 'float_heatmaps:0' 16 | offsets: 'float_short_offsets:0' 17 | displacement_fwd: 'resnet_v1_50/displacement_fwd_2/BiasAdd:0' 18 | displacement_bwd: 'resnet_v1_50/displacement_bwd_2/BiasAdd:0' 19 | part_heatmap: 'float_part_heatmaps:0' 20 | part_offsets: 'float_part_offsets:0' 21 | long_offsets: 'float_long_offsets:0' 22 | segments: 'float_segments:0' 23 | mobilenet_v1_100: 24 | base_url: 'https://storage.googleapis.com/tfjs-models/savedmodel/bodypix/mobilenet/float/100' 25 | model_variant: 26 | stride16: 27 | filename: 'model-stride16.json' 28 | output_stride: 16 29 | data_format: 'NHWC' 30 | input_tensors: 31 | image: 'sub_2:0' 32 | output_tensors: 33 | heatmap: 'float_heatmaps:0' 34 | offsets: 'float_short_offsets:0' 35 | displacement_fwd: 'MobilenetV1/displacement_fwd_2/BiasAdd:0' 36 | displacement_bwd: 'MobilenetV1/displacement_bwd_2/BiasAdd:0' 37 | part_heatmap: 'float_part_heatmaps:0' 38 | part_offsets: 'float_part_offsets:0' 39 | long_offsets: 'float_long_offsets:0' 40 | segments: 'float_segments:0' 41 | posenet: 42 | resnet50_v1: 43 | base_url: 'https://storage.googleapis.com/tfjs-models/savedmodel/posenet/resnet50/float' 44 | model_variant: 45 | stride16: 46 | filename: 'model-stride16.json' 47 | output_stride: 16 48 | data_format: 'NHWC' 49 | stride32: 50 | filename: 'model-stride32.json' 51 | output_stride: 32 52 | data_format: 'NHWC' 53 | input_tensors: 54 | image: 'sub_2:0' 55 | output_tensors: 56 | heatmap: 'float_heatmaps:0' 57 | offsets: 'float_short_offsets:0' 58 | displacement_fwd: 'resnet_v1_50/displacement_fwd_2/BiasAdd:0' 59 | displacement_bwd: 'resnet_v1_50/displacement_bwd_2/BiasAdd:0' 60 | mobilenet_v1_100: 61 | base_url: 'https://storage.googleapis.com/tfjs-models/savedmodel/posenet/mobilenet/float/100' 62 | model_variant: 63 | stride16: 64 | filename: 'model-stride16.json' 65 | output_stride: 16 66 | data_format: 'NHWC' 67 | input_tensors: 68 | image: 'sub_2:0' 69 | output_tensors: 70 | heatmap: 'MobilenetV1/heatmap_2/BiasAdd:0' 71 | offsets: 'MobilenetV1/offset_2/BiasAdd:0' 72 | displacement_fwd: 'MobilenetV1/displacement_fwd_2/BiasAdd:0' 73 | displacement_bwd: 'MobilenetV1/displacement_bwd_2/BiasAdd:0' 74 | -------------------------------------------------------------------------------- /posenet/converter/tfjs2tf.py: -------------------------------------------------------------------------------- 1 | import os 2 | import tensorflow as tf 3 | import tfjs_graph_converter as tfjs 4 | import posenet.converter.config as config 5 | import posenet.converter.tfjsdownload as tfjsdownload 6 | 7 | 8 | def __tensor_info_def(sess, tensor_names): 9 | signatures = {} 10 | for tensor_name in tensor_names: 11 | tensor = sess.graph.get_tensor_by_name(tensor_name) 12 | tensor_info = tf.compat.v1.saved_model.build_tensor_info(tensor) 13 | signatures[tensor_name] = tensor_info 14 | return signatures 15 | 16 | 17 | def convert(model_cfg): 18 | model_file_path = os.path.join(model_cfg['tfjs_dir'], model_cfg['filename']) 19 | if not os.path.exists(model_file_path): 20 | print('Cannot find tfjs model path %s, downloading tfjs model...' % model_file_path) 21 | tfjsdownload.download_tfjs_model(model_cfg) 22 | 23 | # 'graph_model_to_saved_model' doesn't store the signature for the model! 24 | # tfjs.api.graph_model_to_saved_model(model_cfg['tfjs_dir'], model_cfg['tf_dir'], ['serve']) 25 | # So we do it manually below. 26 | # This link was a great help to do this: 27 | # https://www.programcreek.com/python/example/104885/tensorflow.python.saved_model.signature_def_utils.build_signature_def 28 | 29 | graph = tfjs.api.load_graph_model(model_cfg['tfjs_dir']) 30 | builder = tf.compat.v1.saved_model.Builder(model_cfg['tf_dir']) 31 | 32 | with tf.compat.v1.Session(graph=graph) as sess: 33 | input_tensor_names = tfjs.util.get_input_tensors(graph) 34 | output_tensor_names = tfjs.util.get_output_tensors(graph) 35 | 36 | signature_inputs = __tensor_info_def(sess, input_tensor_names) 37 | signature_outputs = __tensor_info_def(sess, output_tensor_names) 38 | 39 | method_name = tf.compat.v1.saved_model.signature_constants.PREDICT_METHOD_NAME 40 | signature_def = tf.compat.v1.saved_model.build_signature_def(inputs=signature_inputs, 41 | outputs=signature_outputs, 42 | method_name=method_name) 43 | signature_map = {tf.compat.v1.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signature_def} 44 | builder.add_meta_graph_and_variables(sess=sess, 45 | tags=['serve'], 46 | signature_def_map=signature_map) 47 | return builder.save() 48 | -------------------------------------------------------------------------------- /posenet/converter/tfjsdownload.py: -------------------------------------------------------------------------------- 1 | import urllib.request 2 | import posixpath 3 | import json 4 | import zlib 5 | import os 6 | import shutil 7 | import posenet.converter.config as config 8 | 9 | 10 | def fix_model_file(model_cfg): 11 | model_file_path = os.path.join(model_cfg['tfjs_dir'], model_cfg['filename']) 12 | 13 | if not model_cfg['filename'] == 'model.json': 14 | # The expected filename for the model json file is 'model.json'. 15 | # See tfjs_common.ARTIFACT_MODEL_JSON_FILE_NAME in the tensorflowjs codebase. 16 | normalized_model_json_file = os.path.join(model_cfg['tfjs_dir'], 'model.json') 17 | shutil.copyfile(model_file_path, normalized_model_json_file) 18 | 19 | with open(model_file_path, 'r') as f: 20 | json_model_def = json.load(f) 21 | 22 | return json_model_def 23 | 24 | 25 | def download_single_file(base_url, filename, save_dir): 26 | output_path = os.path.join(save_dir, filename) 27 | url = posixpath.join(base_url, filename) 28 | req = urllib.request.Request(url) 29 | response = urllib.request.urlopen(req) 30 | if response.info().get('Content-Encoding') == 'gzip': 31 | data = zlib.decompress(response.read(), zlib.MAX_WBITS | 32) 32 | else: 33 | # this path not tested since gzip encoding default on google server 34 | # may need additional encoding/text handling if hit in the future 35 | data = response.read() 36 | with open(output_path, 'wb') as f: 37 | f.write(data) 38 | 39 | 40 | def download_tfjs_model(model_cfg): 41 | """ 42 | Download a tfjs model with saved weights. 43 | 44 | :param model_cfg: The model configuration 45 | """ 46 | model_file_path = os.path.join(model_cfg['tfjs_dir'], model_cfg['filename']) 47 | if os.path.exists(model_file_path): 48 | print('Model file already exists: %s...' % model_file_path) 49 | return 50 | if not os.path.exists(model_cfg['tfjs_dir']): 51 | os.makedirs(model_cfg['tfjs_dir']) 52 | 53 | download_single_file(model_cfg['base_url'], model_cfg['filename'], model_cfg['tfjs_dir']) 54 | 55 | json_model_def = fix_model_file(model_cfg) 56 | 57 | shard_paths = json_model_def['weightsManifest'][0]['paths'] 58 | for shard in shard_paths: 59 | download_single_file(model_cfg['base_url'], shard, model_cfg['tfjs_dir']) 60 | -------------------------------------------------------------------------------- /posenet/decode.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from posenet.constants import * 4 | 5 | 6 | def traverse_to_targ_keypoint( 7 | edge_id, source_keypoint, target_keypoint_id, scores, offsets, output_stride, displacements 8 | ): 9 | height = scores.shape[0] 10 | width = scores.shape[1] 11 | 12 | source_keypoint_indices = np.clip( 13 | np.round(source_keypoint / output_stride), a_min=0, a_max=[height - 1, width - 1]).astype(np.int32) 14 | 15 | displaced_point = source_keypoint + displacements[ 16 | source_keypoint_indices[0], source_keypoint_indices[1], edge_id] 17 | 18 | displaced_point_indices = np.clip( 19 | np.round(displaced_point / output_stride), a_min=0, a_max=[height - 1, width - 1]).astype(np.int32) 20 | 21 | score = scores[displaced_point_indices[0], displaced_point_indices[1], target_keypoint_id] 22 | 23 | image_coord = displaced_point_indices * output_stride + offsets[ 24 | displaced_point_indices[0], displaced_point_indices[1], target_keypoint_id] 25 | 26 | return score, image_coord 27 | 28 | 29 | def decode_pose( 30 | root_score, root_id, root_image_coord, 31 | scores, 32 | offsets, 33 | output_stride, 34 | displacements_fwd, 35 | displacements_bwd 36 | ): 37 | num_parts = scores.shape[2] 38 | num_edges = len(PARENT_CHILD_TUPLES) 39 | 40 | instance_keypoint_scores = np.zeros(num_parts) 41 | instance_keypoint_coords = np.zeros((num_parts, 2)) 42 | instance_keypoint_scores[root_id] = root_score 43 | instance_keypoint_coords[root_id] = root_image_coord 44 | 45 | for edge in reversed(range(num_edges)): 46 | target_keypoint_id, source_keypoint_id = PARENT_CHILD_TUPLES[edge] 47 | if (instance_keypoint_scores[source_keypoint_id] > 0.0 and 48 | instance_keypoint_scores[target_keypoint_id] == 0.0): 49 | score, coords = traverse_to_targ_keypoint( 50 | edge, 51 | instance_keypoint_coords[source_keypoint_id], 52 | target_keypoint_id, 53 | scores, offsets, output_stride, displacements_bwd) 54 | instance_keypoint_scores[target_keypoint_id] = score 55 | instance_keypoint_coords[target_keypoint_id] = coords 56 | 57 | for edge in range(num_edges): 58 | source_keypoint_id, target_keypoint_id = PARENT_CHILD_TUPLES[edge] 59 | if (instance_keypoint_scores[source_keypoint_id] > 0.0 and 60 | instance_keypoint_scores[target_keypoint_id] == 0.0): 61 | score, coords = traverse_to_targ_keypoint( 62 | edge, 63 | instance_keypoint_coords[source_keypoint_id], 64 | target_keypoint_id, 65 | scores, offsets, output_stride, displacements_fwd) 66 | instance_keypoint_scores[target_keypoint_id] = score 67 | instance_keypoint_coords[target_keypoint_id] = coords 68 | 69 | return instance_keypoint_scores, instance_keypoint_coords 70 | -------------------------------------------------------------------------------- /posenet/decode_multi.py: -------------------------------------------------------------------------------- 1 | from posenet.decode import * 2 | from posenet.constants import * 3 | import time 4 | import scipy.ndimage as ndi 5 | 6 | 7 | def within_nms_radius(poses, squared_nms_radius, point, keypoint_id): 8 | for _, _, pose_coord in poses: 9 | if np.sum((pose_coord[keypoint_id] - point) ** 2) <= squared_nms_radius: 10 | return True 11 | return False 12 | 13 | 14 | def within_nms_radius_fast(pose_coords, squared_nms_radius, point): 15 | if not pose_coords.shape[0]: 16 | return False 17 | return np.any(np.sum((pose_coords - point) ** 2, axis=1) <= squared_nms_radius) 18 | 19 | 20 | def get_instance_score( 21 | existing_poses, squared_nms_radius, 22 | keypoint_scores, keypoint_coords): 23 | not_overlapped_scores = 0. 24 | for keypoint_id in range(len(keypoint_scores)): 25 | if not within_nms_radius( 26 | existing_poses, squared_nms_radius, 27 | keypoint_coords[keypoint_id], keypoint_id): 28 | not_overlapped_scores += keypoint_scores[keypoint_id] 29 | return not_overlapped_scores / len(keypoint_scores) 30 | 31 | 32 | def get_instance_score_fast( 33 | exist_pose_coords, 34 | squared_nms_radius, 35 | keypoint_scores, keypoint_coords): 36 | 37 | if exist_pose_coords.shape[0]: 38 | s = np.sum((exist_pose_coords - keypoint_coords) ** 2, axis=2) > squared_nms_radius 39 | not_overlapped_scores = np.sum(keypoint_scores[np.all(s, axis=0)]) 40 | else: 41 | not_overlapped_scores = np.sum(keypoint_scores) 42 | return not_overlapped_scores / len(keypoint_scores) 43 | 44 | 45 | def score_is_max_in_local_window(keypoint_id, score, hmy, hmx, local_max_radius, scores): 46 | height = scores.shape[0] 47 | width = scores.shape[1] 48 | 49 | y_start = max(hmy - local_max_radius, 0) 50 | y_end = min(hmy + local_max_radius + 1, height) 51 | x_start = max(hmx - local_max_radius, 0) 52 | x_end = min(hmx + local_max_radius + 1, width) 53 | 54 | for y in range(y_start, y_end): 55 | for x in range(x_start, x_end): 56 | if scores[y, x, keypoint_id] > score: 57 | return False 58 | return True 59 | 60 | 61 | def build_part_with_score(score_threshold, local_max_radius, scores): 62 | parts = [] 63 | height = scores.shape[0] 64 | width = scores.shape[1] 65 | num_keypoints = scores.shape[2] 66 | 67 | for hmy in range(height): 68 | for hmx in range(width): 69 | for keypoint_id in range(num_keypoints): 70 | score = scores[hmy, hmx, keypoint_id] 71 | if score < score_threshold: 72 | continue 73 | if score_is_max_in_local_window(keypoint_id, score, hmy, hmx, 74 | local_max_radius, scores): 75 | parts.append(( 76 | score, keypoint_id, np.array((hmy, hmx)) 77 | )) 78 | return parts 79 | 80 | 81 | def build_part_with_score_fast(score_threshold, local_max_radius, scores): 82 | parts = [] 83 | num_keypoints = scores.shape[2] 84 | lmd = 2 * local_max_radius + 1 85 | 86 | # NOTE it seems faster to iterate over the keypoints and perform maximum_filter 87 | # on each subarray vs doing the op on the full score array with size=(lmd, lmd, 1) 88 | for keypoint_id in range(num_keypoints): 89 | kp_scores = scores[:, :, keypoint_id].copy() 90 | kp_scores[kp_scores < score_threshold] = 0. 91 | max_vals = ndi.maximum_filter(kp_scores, size=lmd, mode='constant') 92 | max_loc = np.logical_and(kp_scores == max_vals, kp_scores > 0) 93 | max_loc_idx = max_loc.nonzero() 94 | for y, x in zip(*max_loc_idx): 95 | parts.append(( 96 | scores[y, x, keypoint_id], 97 | keypoint_id, 98 | np.array((y, x)) 99 | )) 100 | 101 | return parts 102 | 103 | 104 | def decode_multiple_poses( 105 | scores, offsets, displacements_fwd, displacements_bwd, output_stride, 106 | max_pose_detections=10, score_threshold=0.5, nms_radius=20, min_pose_score=0.5): 107 | 108 | pose_count = 0 109 | pose_scores = np.zeros(max_pose_detections) 110 | pose_keypoint_scores = np.zeros((max_pose_detections, NUM_KEYPOINTS)) 111 | pose_keypoint_coords = np.zeros((max_pose_detections, NUM_KEYPOINTS, 2)) 112 | 113 | squared_nms_radius = nms_radius ** 2 114 | 115 | scored_parts = build_part_with_score_fast(score_threshold, LOCAL_MAXIMUM_RADIUS, scores) 116 | scored_parts = sorted(scored_parts, key=lambda x: x[0], reverse=True) 117 | 118 | # change dimensions from (h, w, x) to (h, w, x//2, 2) to allow return of complete coord array 119 | height = scores.shape[0] 120 | width = scores.shape[1] 121 | offsets = offsets.reshape(height, width, 2, -1).swapaxes(2, 3) 122 | displacements_fwd = displacements_fwd.reshape(height, width, 2, -1).swapaxes(2, 3) 123 | displacements_bwd = displacements_bwd.reshape(height, width, 2, -1).swapaxes(2, 3) 124 | 125 | for root_score, root_id, root_coord in scored_parts: 126 | root_image_coords = root_coord * output_stride + offsets[ 127 | root_coord[0], root_coord[1], root_id] 128 | 129 | if within_nms_radius_fast( 130 | pose_keypoint_coords[:pose_count, root_id, :], squared_nms_radius, root_image_coords): 131 | continue 132 | 133 | keypoint_scores, keypoint_coords = decode_pose( 134 | root_score, root_id, root_image_coords, 135 | scores, offsets, output_stride, 136 | displacements_fwd, displacements_bwd) 137 | 138 | pose_score = get_instance_score_fast( 139 | pose_keypoint_coords[:pose_count, :, :], squared_nms_radius, keypoint_scores, keypoint_coords) 140 | 141 | # NOTE this isn't in the original implementation, but it appears that by initially ordering by 142 | # part scores, and having a max # of detections, we can end up populating the returned poses with 143 | # lower scored poses than if we discard 'bad' ones and continue (higher pose scores can still come later). 144 | # Set min_pose_score to 0. to revert to original behaviour 145 | if min_pose_score == 0. or pose_score >= min_pose_score: 146 | pose_scores[pose_count] = pose_score 147 | pose_keypoint_scores[pose_count, :] = keypoint_scores 148 | pose_keypoint_coords[pose_count, :, :] = keypoint_coords 149 | pose_count += 1 150 | 151 | if pose_count >= max_pose_detections: 152 | break 153 | 154 | return pose_scores, pose_keypoint_scores, pose_keypoint_coords 155 | -------------------------------------------------------------------------------- /posenet/mobilenet.py: -------------------------------------------------------------------------------- 1 | from posenet.base_model import BaseModel 2 | import numpy as np 3 | import cv2 4 | 5 | 6 | class MobileNet(BaseModel): 7 | 8 | def __init__(self, model_function, output_tensor_names, output_stride): 9 | super().__init__(model_function, output_tensor_names, output_stride) 10 | 11 | def preprocess_input(self, image): 12 | target_width, target_height = self.valid_resolution(image.shape[1], image.shape[0]) 13 | # the scale that can get us back to the original width and height: 14 | scale = np.array([image.shape[0] / target_height, image.shape[1] / target_width]) 15 | input_img = cv2.resize(image, (target_width, target_height), interpolation=cv2.INTER_LINEAR) 16 | input_img = cv2.cvtColor(input_img, cv2.COLOR_BGR2RGB).astype(np.float32) # to RGB colors 17 | 18 | input_img = input_img * (2.0 / 255.0) - 1.0 # normalize to [-1,1] 19 | input_img = input_img.reshape(1, target_height, target_width, 3) # NHWC 20 | return input_img, scale 21 | -------------------------------------------------------------------------------- /posenet/posenet.py: -------------------------------------------------------------------------------- 1 | from posenet.base_model import BaseModel 2 | import posenet 3 | 4 | 5 | class PoseNet: 6 | 7 | def __init__(self, model: BaseModel, min_score=0.25): 8 | self.model = model 9 | self.min_score = min_score 10 | 11 | def estimate_multiple_poses(self, image, max_pose_detections=10): 12 | heatmap_result, offsets_result, displacement_fwd_result, displacement_bwd_result, image_scale = \ 13 | self.model.predict(image) 14 | 15 | pose_scores, keypoint_scores, keypoint_coords = posenet.decode_multiple_poses( 16 | heatmap_result.numpy().squeeze(axis=0), 17 | offsets_result.numpy().squeeze(axis=0), 18 | displacement_fwd_result.numpy().squeeze(axis=0), 19 | displacement_bwd_result.numpy().squeeze(axis=0), 20 | output_stride=self.model.output_stride, 21 | max_pose_detections=max_pose_detections, 22 | min_pose_score=self.min_score) 23 | 24 | keypoint_coords *= image_scale 25 | 26 | return pose_scores, keypoint_scores, keypoint_coords 27 | 28 | def estimate_single_pose(self, image): 29 | return self.estimate_multiple_poses(image, max_pose_detections=1) 30 | 31 | def draw_poses(self, image, pose_scores, keypoint_scores, keypoint_coords): 32 | draw_image = posenet.draw_skel_and_kp( 33 | image, pose_scores, keypoint_scores, keypoint_coords, 34 | min_pose_score=self.min_score, min_part_score=self.min_score) 35 | 36 | return draw_image 37 | 38 | def print_scores(self, image_name, pose_scores, keypoint_scores, keypoint_coords): 39 | print() 40 | print("Results for image: %s" % image_name) 41 | for pi in range(len(pose_scores)): 42 | if pose_scores[pi] == 0.: 43 | break 44 | print('Pose #%d, score = %f' % (pi, pose_scores[pi])) 45 | for ki, (s, c) in enumerate(zip(keypoint_scores[pi, :], keypoint_coords[pi, :, :])): 46 | print('Keypoint %s, score = %f, coord = %s' % (posenet.PART_NAMES[ki], s, c)) 47 | -------------------------------------------------------------------------------- /posenet/posenet_factory.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import os 3 | import posenet.converter.config as config 4 | import posenet.converter.tfjs2tf as tfjs2tf 5 | from posenet.resnet import ResNet 6 | from posenet.mobilenet import MobileNet 7 | from posenet.posenet import PoseNet 8 | 9 | 10 | def load_model(model, stride, quant_bytes=4, multiplier=1.0): 11 | 12 | if model == config.RESNET50_MODEL: 13 | model_cfg = config.bodypix_resnet50_config(stride, quant_bytes) 14 | print('Loading ResNet50 model') 15 | else: 16 | model_cfg = config.bodypix_mobilenet_config(stride, quant_bytes, multiplier) 17 | print('Loading MobileNet model') 18 | 19 | model_path = model_cfg['tf_dir'] 20 | if not os.path.exists(model_path): 21 | print('Cannot find tf model path %s, converting from tfjs...' % model_path) 22 | tfjs2tf.convert(model_cfg) 23 | assert os.path.exists(model_path) 24 | 25 | loaded_model = tf.saved_model.load(model_path) 26 | 27 | signature_key = tf.compat.v1.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY 28 | print('We use the signature key %s It should be in the keys list:' % signature_key) 29 | for sig in loaded_model.signatures.keys(): 30 | print('signature key: %s' % sig) 31 | 32 | model_function = loaded_model.signatures[signature_key] 33 | print('model outputs: %s' % model_function.structured_outputs) 34 | 35 | output_tensor_names = model_cfg['output_tensors'] 36 | output_stride = model_cfg['output_stride'] 37 | 38 | if model == config.RESNET50_MODEL: 39 | net = ResNet(model_function, output_tensor_names, output_stride) 40 | else: 41 | net = MobileNet(model_function, output_tensor_names, output_stride) 42 | 43 | return PoseNet(net) 44 | -------------------------------------------------------------------------------- /posenet/resnet.py: -------------------------------------------------------------------------------- 1 | from posenet.base_model import BaseModel 2 | import numpy as np 3 | import cv2 4 | 5 | 6 | class ResNet(BaseModel): 7 | 8 | def __init__(self, model_function, output_tensor_names, output_stride): 9 | super().__init__(model_function, output_tensor_names, output_stride) 10 | self.image_net_mean = [-123.15, -115.90, -103.06] 11 | 12 | def preprocess_input(self, image): 13 | target_width, target_height = self.valid_resolution(image.shape[1], image.shape[0]) 14 | # the scale that can get us back to the original width and height: 15 | scale = np.array([image.shape[0] / target_height, image.shape[1] / target_width]) 16 | input_img = cv2.resize(image, (target_width, target_height), interpolation=cv2.INTER_LINEAR) 17 | input_img = cv2.cvtColor(input_img, cv2.COLOR_BGR2RGB).astype(np.float32) # to RGB colors 18 | # todo: test a variant that adds black bars to the image to match it to a valid resolution 19 | 20 | # See: https://github.com/tensorflow/tfjs-models/blob/master/body-pix/src/resnet.ts 21 | input_img = input_img + self.image_net_mean 22 | input_img = input_img.reshape(1, target_height, target_width, 3) # HWC to NHWC 23 | return input_img, scale 24 | -------------------------------------------------------------------------------- /posenet/utils.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | import posenet.constants 4 | 5 | 6 | def draw_keypoints( 7 | img, instance_scores, keypoint_scores, keypoint_coords, 8 | min_pose_confidence=0.5, min_part_confidence=0.5): 9 | cv_keypoints = [] 10 | for ii, score in enumerate(instance_scores): 11 | if score < min_pose_confidence: 12 | continue 13 | for ks, kc in zip(keypoint_scores[ii, :], keypoint_coords[ii, :, :]): 14 | if ks < min_part_confidence: 15 | continue 16 | cv_keypoints.append(cv2.KeyPoint(kc[1], kc[0], 10. * ks)) 17 | out_img = cv2.drawKeypoints(img, cv_keypoints, outImage=np.array([])) 18 | return out_img 19 | 20 | 21 | def get_adjacent_keypoints(keypoint_scores, keypoint_coords, min_confidence=0.1): 22 | results = [] 23 | for left, right in posenet.CONNECTED_PART_INDICES: 24 | if keypoint_scores[left] < min_confidence or keypoint_scores[right] < min_confidence: 25 | continue 26 | results.append( 27 | np.array([keypoint_coords[left][::-1], keypoint_coords[right][::-1]]).astype(np.int32), 28 | ) 29 | return results 30 | 31 | 32 | def draw_skeleton( 33 | img, instance_scores, keypoint_scores, keypoint_coords, 34 | min_pose_confidence=0.5, min_part_confidence=0.5): 35 | out_img = img 36 | adjacent_keypoints = [] 37 | for ii, score in enumerate(instance_scores): 38 | if score < min_pose_confidence: 39 | continue 40 | new_keypoints = get_adjacent_keypoints( 41 | keypoint_scores[ii, :], keypoint_coords[ii, :, :], min_part_confidence) 42 | adjacent_keypoints.extend(new_keypoints) 43 | out_img = cv2.polylines(out_img, adjacent_keypoints, isClosed=False, color=(255, 255, 0)) 44 | return out_img 45 | 46 | 47 | def draw_skel_and_kp( 48 | img, instance_scores, keypoint_scores, keypoint_coords, 49 | min_pose_score=0.5, min_part_score=0.5): 50 | out_img = img 51 | adjacent_keypoints = [] 52 | cv_keypoints = [] 53 | for ii, score in enumerate(instance_scores): 54 | if score < min_pose_score: 55 | continue 56 | 57 | new_keypoints = get_adjacent_keypoints( 58 | keypoint_scores[ii, :], keypoint_coords[ii, :, :], min_part_score) 59 | adjacent_keypoints.extend(new_keypoints) 60 | 61 | for ks, kc in zip(keypoint_scores[ii, :], keypoint_coords[ii, :, :]): 62 | if ks < min_part_score: 63 | continue 64 | cv_keypoints.append(cv2.KeyPoint(kc[1], kc[0], 10. * ks)) 65 | 66 | out_img = cv2.drawKeypoints( 67 | out_img, cv_keypoints, outImage=np.array([]), color=(255, 255, 255), 68 | flags=cv2.DRAW_MATCHES_FLAGS_DRAW_RICH_KEYPOINTS) 69 | out_img = cv2.polylines(out_img, adjacent_keypoints, isClosed=False, color=(255, 255, 0), thickness=2) 70 | return out_img 71 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | numpy==1.18.1 2 | opencv_python_headless==3.4.5.20 3 | scipy==1.4.1 4 | tensorflowjs==1.4.0 5 | # tensorflow==2.1.0 # uncomment when installing from scratch (not in docker) 6 | PyYAML==5.3 7 | tensorflow-hub==0.7.0 -------------------------------------------------------------------------------- /video_demo.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import cv2 3 | import time 4 | import argparse 5 | 6 | from posenet.posenet_factory import load_model 7 | from posenet.utils import draw_skel_and_kp 8 | 9 | parser = argparse.ArgumentParser() 10 | parser.add_argument('--model', type=str, default='resnet50') # mobilenet resnet50 11 | parser.add_argument('--stride', type=int, default=16) # 8, 16, 32 (max 16 for mobilenet) 12 | parser.add_argument('--quant_bytes', type=int, default=4) # 4 = float 13 | parser.add_argument('--multiplier', type=float, default=1.0) # only for mobilenet 14 | parser.add_argument('--scale_factor', type=float, default=0.7125) 15 | parser.add_argument('--input_file', type=str, help="Give the video file location") 16 | parser.add_argument('--output_file', type=str, help="Give the video file location") 17 | args = parser.parse_args() 18 | 19 | 20 | def main(): 21 | 22 | print('Tensorflow version: %s' % tf.__version__) 23 | assert tf.__version__.startswith('2.'), "Tensorflow version 2.x must be used!" 24 | 25 | model = args.model # mobilenet resnet50 26 | stride = args.stride # 8, 16, 32 (max 16 for mobilenet, min 16 for resnet50) 27 | quant_bytes = args.quant_bytes # float 28 | multiplier = args.multiplier # only for mobilenet 29 | 30 | posenet = load_model(model, stride, quant_bytes, multiplier) 31 | 32 | # for inspiration, see: https://www.programcreek.com/python/example/72134/cv2.VideoWriter 33 | if args.input_file is not None: 34 | cap = cv2.VideoCapture(args.input_file) 35 | else: 36 | raise IOError("video file not found") 37 | 38 | fps = cap.get(cv2.CAP_PROP_FPS) 39 | width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) 40 | height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) 41 | 42 | fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v') 43 | video_writer = cv2.VideoWriter(args.output_file, fourcc, fps, (width, height)) 44 | 45 | max_pose_detections = 20 46 | 47 | # Scaling the input image reduces the quality of the pose detections! 48 | # The speed gain is about the square of the scale factor. 49 | posenet_input_height = 540 # scale factor for the posenet input 50 | posenet_input_scale = 1.0 # posenet_input_height / height # 1.0 51 | posenet_input_width = int(width * posenet_input_scale) 52 | print("posenet_input_scale: %3.4f" % (posenet_input_scale)) 53 | 54 | 55 | start = time.time() 56 | frame_count = 0 57 | 58 | ret, frame = cap.read() 59 | 60 | while ret: 61 | if posenet_input_scale == 1.0: 62 | frame_rescaled = frame # no scaling 63 | else: 64 | frame_rescaled = \ 65 | cv2.resize(frame, (posenet_input_width, posenet_input_height), interpolation=cv2.INTER_LINEAR) 66 | 67 | pose_scores, keypoint_scores, keypoint_coords = posenet.estimate_multiple_poses(frame_rescaled, max_pose_detections) 68 | 69 | keypoint_coords_upscaled = keypoint_coords / posenet_input_scale 70 | overlay_frame = draw_skel_and_kp( 71 | frame, pose_scores, keypoint_scores, keypoint_coords_upscaled, 72 | min_pose_score=0.15, min_part_score=0.1) 73 | 74 | frame_count += 1 75 | # This is uncompressed video. cv2 has no way to write compressed videos, so we'll have to use ffmpeg to 76 | # compress it afterwards! See: 77 | # https://stackoverflow.com/questions/25998799/specify-compression-quality-in-python-for-opencv-video-object 78 | video_writer.write(overlay_frame) 79 | ret, frame = cap.read() 80 | 81 | print('Average FPS: ', frame_count / (time.time() - start)) 82 | 83 | video_writer.release() 84 | cap.release() 85 | 86 | if __name__ == "__main__": 87 | main() 88 | -------------------------------------------------------------------------------- /webcam_demo.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import cv2 3 | import time 4 | import argparse 5 | 6 | from posenet.posenet_factory import load_model 7 | from posenet.utils import draw_skel_and_kp 8 | 9 | parser = argparse.ArgumentParser() 10 | parser.add_argument('--model', type=str, default='resnet50') # mobilenet resnet50 11 | parser.add_argument('--stride', type=int, default=16) # 8, 16, 32 (max 16 for mobilenet) 12 | parser.add_argument('--quant_bytes', type=int, default=4) # 4 = float 13 | parser.add_argument('--multiplier', type=float, default=1.0) # only for mobilenet 14 | parser.add_argument('--cam_id', type=int, default=0) 15 | parser.add_argument('--cam_width', type=int, default=1280) 16 | parser.add_argument('--cam_height', type=int, default=720) 17 | parser.add_argument('--scale_factor', type=float, default=0.7125) 18 | parser.add_argument('--file', type=str, default=None, help="Optionally use a video file instead of a live camera") 19 | args = parser.parse_args() 20 | 21 | 22 | def main(): 23 | 24 | print('Tensorflow version: %s' % tf.__version__) 25 | assert tf.__version__.startswith('2.'), "Tensorflow version 2.x must be used!" 26 | 27 | model = args.model # mobilenet resnet50 28 | stride = args.stride # 8, 16, 32 (max 16 for mobilenet, min 16 for resnet50) 29 | quant_bytes = args.quant_bytes # float 30 | multiplier = args.multiplier # only for mobilenet 31 | 32 | posenet = load_model(model, stride, quant_bytes, multiplier) 33 | 34 | if args.file is not None: 35 | cap = cv2.VideoCapture(args.file) 36 | else: 37 | cap = cv2.VideoCapture(args.cam_id) 38 | cap.set(3, args.cam_width) 39 | cap.set(4, args.cam_height) 40 | 41 | start = time.time() 42 | frame_count = 0 43 | 44 | while True: 45 | res, img = cap.read() 46 | if not res: 47 | raise IOError("webcam failure") 48 | 49 | pose_scores, keypoint_scores, keypoint_coords = posenet.estimate_multiple_poses(img) 50 | 51 | overlay_image = draw_skel_and_kp( 52 | img, pose_scores, keypoint_scores, keypoint_coords, 53 | min_pose_score=0.15, min_part_score=0.1) 54 | 55 | cv2.imshow('posenet', overlay_image) 56 | frame_count += 1 57 | if cv2.waitKey(1) & 0xFF == ord('q'): 58 | break 59 | 60 | print('Average FPS: ', frame_count / (time.time() - start)) 61 | 62 | 63 | if __name__ == "__main__": 64 | main() --------------------------------------------------------------------------------