├── .github
└── workflows
│ └── python-publish.yml
├── .gitignore
├── LICENSE
├── README.md
├── demo.py
├── docs
├── Makefile
├── _static
│ └── styles
│ │ └── my_theme.css
├── _templates
│ └── autoapi
│ │ ├── index.rst
│ │ └── python
│ │ ├── attribute.rst
│ │ ├── class.rst
│ │ ├── data.rst
│ │ ├── exception.rst
│ │ ├── function.rst
│ │ ├── method.rst
│ │ ├── module.rst
│ │ ├── package.rst
│ │ └── property.rst
├── abbrev_long.bib
├── conf.py
├── conf_overrides.py
├── conf_spec.py
├── index.rst
├── make.bat
└── references.bib
├── docs_old
├── API.md
├── index.md
├── poseviz.md
└── stylesheets
│ └── extra.css
├── pyproject.toml
├── screenshot.jpg
├── screenshot2.jpg
├── screenshot_multicam.jpg
└── src
└── poseviz
├── __init__.py
├── colors.py
├── components
├── __init__.py
├── camera_viz.py
├── connected_points_viz.py
├── ground_plane_viz.py
├── main_viz.py
├── skeletons_viz.py
└── smpl_viz.py
├── draw2d.py
├── init.py
├── mayavi_util.py
├── messages.py
├── poseviz.py
├── video_writing.py
└── view_info.py
/.github/workflows/python-publish.yml:
--------------------------------------------------------------------------------
1 | name: Upload Python Package
2 |
3 | on:
4 | release:
5 | types: [published]
6 |
7 | permissions:
8 | contents: read
9 |
10 | jobs:
11 | pypi-publish:
12 | name: Upload release to PyPI
13 | runs-on: ubuntu-latest
14 | environment: pypi
15 | permissions:
16 | id-token: write
17 | steps:
18 | - name: Check out repository
19 | uses: actions/checkout@v4
20 | with:
21 | fetch-depth: 0
22 |
23 | - name: Set up Python
24 | uses: actions/setup-python@v5
25 | with:
26 | python-version: "3.x"
27 |
28 | - name: Install build dependencies
29 | run: python -m pip install --upgrade build
30 |
31 | - name: Build package distribution
32 | run: python -m build --sdist
33 |
34 | - name: Publish package distributions to PyPI
35 | uses: pypa/gh-action-pypi-publish@release/v1
36 |
--------------------------------------------------------------------------------
/.gitignore:
--------------------------------------------------------------------------------
1 | # Byte-compiled / optimized / DLL files
2 | __pycache__/
3 | *.py[cod]
4 | *$py.class
5 |
6 | # C extensions
7 | *.so
8 |
9 | # Distribution / packaging
10 | .Python
11 | build/
12 | develop-eggs/
13 | dist/
14 | downloads/
15 | eggs/
16 | .eggs/
17 | lib/
18 | lib64/
19 | parts/
20 | sdist/
21 | var/
22 | wheels/
23 | share/python-wheels/
24 | *.egg-info/
25 | .installed.cfg
26 | *.egg
27 | MANIFEST
28 |
29 | # PyInstaller
30 | # Usually these files are written by a python script from a template
31 | # before PyInstaller builds the exe, so as to inject date/other infos into it.
32 | *.manifest
33 | *.spec
34 |
35 | # Installer logs
36 | pip-log.txt
37 | pip-delete-this-directory.txt
38 |
39 | # Unit test / coverage reports
40 | htmlcov/
41 | .tox/
42 | .nox/
43 | .coverage
44 | .coverage.*
45 | .cache
46 | nosetests.xml
47 | coverage.xml
48 | *.cover
49 | *.py,cover
50 | .hypothesis/
51 | .pytest_cache/
52 | cover/
53 |
54 | # Translations
55 | *.mo
56 | *.pot
57 |
58 | # Django stuff:
59 | *.log
60 | local_settings.py
61 | db.sqlite3
62 | db.sqlite3-journal
63 |
64 | # Flask stuff:
65 | instance/
66 | .webassets-cache
67 |
68 | # Scrapy stuff:
69 | .scrapy
70 |
71 | # Sphinx documentation
72 | docs/_build/
73 |
74 | # PyBuilder
75 | .pybuilder/
76 | target/
77 |
78 | # Jupyter Notebook
79 | .ipynb_checkpoints
80 |
81 | # IPython
82 | profile_default/
83 | ipython_config.py
84 |
85 | # pyenv
86 | # For a library or package, you might want to ignore these files since the code is
87 | # intended to run in multiple environments; otherwise, check them in:
88 | # .python-version
89 |
90 | # pipenv
91 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
92 | # However, in case of collaboration, if having platform-specific dependencies or dependencies
93 | # having no cross-platform support, pipenv may install dependencies that don't work, or not
94 | # install all needed dependencies.
95 | #Pipfile.lock
96 |
97 | # poetry
98 | # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
99 | # This is especially recommended for binary packages to ensure reproducibility, and is more
100 | # commonly ignored for libraries.
101 | # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
102 | #poetry.lock
103 |
104 | # pdm
105 | # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
106 | #pdm.lock
107 | # pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
108 | # in version control.
109 | # https://pdm.fming.dev/#use-with-ide
110 | .pdm.toml
111 |
112 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
113 | __pypackages__/
114 |
115 | # Celery stuff
116 | celerybeat-schedule
117 | celerybeat.pid
118 |
119 | # SageMath parsed files
120 | *.sage.py
121 |
122 | # Environments
123 | .env
124 | .venv
125 | env/
126 | venv/
127 | ENV/
128 | env.bak/
129 | venv.bak/
130 |
131 | # Spyder project settings
132 | .spyderproject
133 | .spyproject
134 |
135 | # Rope project settings
136 | .ropeproject
137 |
138 | # mkdocs documentation
139 | /site
140 |
141 | # mypy
142 | .mypy_cache/
143 | .dmypy.json
144 | dmypy.json
145 |
146 | # Pyre type checker
147 | .pyre/
148 |
149 | # pytype static type analyzer
150 | .pytype/
151 |
152 | # Cython debug symbols
153 | cython_debug/
154 |
155 | *_cython.c
156 |
157 | # PyCharm
158 | # JetBrains specific template is maintained in a separate JetBrains.gitignore that can
159 | # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
160 | # and can be added to the global gitignore or merged into this file. For a more nuclear
161 | # option (not recommended) you can uncomment the following to ignore the entire idea folder.
162 | #.idea/
163 |
--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
1 | MIT License
2 |
3 | Copyright (c) 2023 István Sárándi
4 |
5 | Permission is hereby granted, free of charge, to any person obtaining a copy
6 | of this software and associated documentation files (the "Software"), to deal
7 | in the Software without restriction, including without limitation the rights
8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9 | copies of the Software, and to permit persons to whom the Software is
10 | furnished to do so, subject to the following conditions:
11 |
12 | The above copyright notice and this permission notice shall be included in all
13 | copies or substantial portions of the Software.
14 |
15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21 | SOFTWARE.
22 |
--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 | # PoseViz – 3D Human Pose Visualizer
2 |
3 |
4 | 
5 | 
6 | 
7 |
8 |
9 | Multi-person, multi-camera 3D human pose visualization tool built using
10 | [Mayavi](https://docs.enthought.com/mayavi/mayavi/). As used
11 | in [MeTRAbs](https://github.com/isarandi/metrabs) visualizations.
12 |
13 | **This repo does not contain pose estimation code, only the visualization part.**
14 |
15 | ## Gist of Usage
16 |
17 | ```python
18 | import poseviz
19 | import cameravision
20 |
21 | camera = cameravision.Camera(...)
22 |
23 | with poseviz.PoseViz(...) as viz:
24 | for frame in frames:
25 | bounding_boxes, poses3d = run_pose_estimation_model(frame)
26 | viz.update(frame=frame, boxes=bounding_boxes, poses=poses3d, camera=camera)
27 | ```
28 |
29 | See also [```demo.py```](demo.py).
30 |
31 | The main feature of this tool is that the graphical event loop is hidden from the library user. We
32 | want to write code in terms of the *prediction loop* of the human pose estimator, not from the point
33 | of view of the visualizer tool.
34 |
35 | Behind the scenes, this is achieved through forking a dedicated visualization process and passing
36 | new scene information via multiprocessing queues.
37 |
38 | ## Installation
39 |
40 | PoseViz is available on PyPI. However, in my experience it is better to install the Mayavi dependency through conda instead of pip.
41 |
42 | ```bash
43 | conda install mayavi -c conda-forge
44 | pip install poseviz
45 | ```
46 |
47 | Then run [demo.py](demo.py) to test if installation was successful.
48 |
49 |
--------------------------------------------------------------------------------
/demo.py:
--------------------------------------------------------------------------------
1 | import time
2 |
3 | import smpl.numpy
4 |
5 | import poseviz
6 |
7 | import numpy as np
8 | import cameravision
9 |
10 |
11 | def main():
12 | # Names of the body joints. Left joint names must start with 'l', right with 'r'.
13 | joint_names = ["l_wrist", "l_elbow"]
14 |
15 | # Joint index pairs specifying which ones should be connected with a line (i.e., the bones of
16 | # the body, e.g. wrist-elbow, elbow-shoulder)
17 | joint_edges = [[0, 1]]
18 | viz = poseviz.PoseViz(joint_names, joint_edges)
19 |
20 | # Iterate over the frames of e.g. a video
21 | for i in range(1):
22 | # Get the current frame
23 | frame = np.zeros([720, 1280, 3], np.uint8)
24 |
25 | # Make predictions here
26 | # ...
27 |
28 | # Update the visualization
29 | viz.update(
30 | frame=frame,
31 | boxes=np.array([[10, 20, 100, 100]], np.float32),
32 | poses=np.array([[[100, 100, 2000], [-100, 100, 2000]]], np.float32),
33 | camera=cameravision.Camera.from_fov(55, frame.shape[:2]),
34 | )
35 |
36 |
37 | def main_smpl():
38 | smpl_canonical = np.load(
39 | "/work/sarandi/projects/localizerfields/canonical_vertices_smpl.npy"
40 | ) * [1,-1,-1] * 1000 + [0, 0, 3000]
41 |
42 | # Names of the body joints. Left joint names must start with 'l', right with 'r'.
43 | joint_names = ["l_wrist", "l_elbow"]
44 |
45 | # Joint index pairs specifying which ones should be connected with a line (i.e., the bones of
46 | # the body, e.g. wrist-elbow, elbow-shoulder)
47 | joint_edges = [[0, 1]]
48 | faces = smpl.numpy.get_cached_body_model("smpl", "neutral").faces
49 | with poseviz.PoseViz(joint_names, joint_edges, body_model_faces=faces) as viz:
50 | # Iterate over the frames of e.g. a video
51 | for i in range(1):
52 | # Get the current frame
53 | frame = np.zeros([720, 1280, 3], np.uint8)
54 |
55 | # Make predictions here
56 | # ...
57 |
58 | # Update the visualization
59 | viz.update(
60 | frame=frame,
61 | boxes=np.array([[10, 20, 100, 100]], np.float32),
62 | vertices=np.array([smpl_canonical + [500, 0, 5000]]),
63 | camera=cameravision.Camera.from_fov(55, frame.shape[:2], world_up=(0, 1, 0)),
64 | )
65 |
66 | viz.update(
67 | frame=frame,
68 | boxes=np.array([[10, 20, 100, 100]], np.float32),
69 | vertices=np.array(
70 | [smpl_canonical + [-500, 0, 3000], smpl_canonical + [-500, 0, 5000]]
71 | ),
72 | camera=cameravision.Camera.from_fov(55, frame.shape[:2], world_up=(0, 1, 0)),
73 | )
74 |
75 | time.sleep(1000)
76 |
77 |
78 | if __name__ == "__main__":
79 | main_smpl()
80 |
--------------------------------------------------------------------------------
/docs/Makefile:
--------------------------------------------------------------------------------
1 | # Minimal makefile for Sphinx documentation
2 | #
3 |
4 | # You can set these variables from the command line, and also
5 | # from the environment for the first two.
6 | SPHINXOPTS ?=
7 | SPHINXBUILD ?= sphinx-build
8 | SOURCEDIR = .
9 | BUILDDIR = _build
10 |
11 | # Put it first so that "make" without argument is like "make help".
12 | help:
13 | @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
14 |
15 | .PHONY: help Makefile
16 |
17 | # Catch-all target: route all unknown targets to Sphinx using the new
18 | # "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS).
19 | %: Makefile
20 | @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
21 |
--------------------------------------------------------------------------------
/docs/_static/styles/my_theme.css:
--------------------------------------------------------------------------------
1 | @import url("theme.css");
2 | @import url("https://fonts.googleapis.com/css2?family=Mona+Sans:ital,wght@0,200..900;1,200..900&family=Geist:wght@100..900&&family=JetBrains+Mono:ital,wght@0,100..800;1,100..800&family=Outfit:wght@100..900&display=swap");
3 |
4 | @media (min-width: 960px) {
5 | .bd-page-width {
6 | max-width: 120rem;
7 | }
8 | }
9 |
10 | #rtd-footer-container {
11 | margin-top: 0 !important;
12 | }
13 |
14 | html[data-theme="light"] {
15 | --pst-color-table-row-hover-bg: #dfc6ff;
16 | --pst-color-link-hover: #845818;
17 | }
18 |
19 | html[data-theme="dark"] {
20 | --pst-color-table-row-hover-bg: #41296c;
21 | --pst-color-inline-code: #dd8cd4;
22 | }
23 |
24 |
25 | html[data-theme="dark"] dt:target {
26 | background-color: #4f4500;
27 | }
28 |
29 | html[data-theme="dark"] .linkcode-link {
30 | color: #9090ff;
31 | }
32 |
33 | html[data-theme="dark"] table.indextable tr.cap {
34 | background-color: #464646;
35 | }
36 |
37 | html[data-theme="dark"] a:visited {
38 | color: #9E67D0;
39 | }
40 |
41 | .navbar-brand .logo__title {
42 | font-family: "Mona Sans", sans-serif;
43 | font-size: 2.5rem;
44 | font-weight: 400;
45 | font-style: normal;
46 | }
47 |
48 | :root {
49 | --pst-font-family-monospace: "JetBrains Mono", monospace;
50 | --pst-font-family-heading: "Mona Sans", sans-serif;
51 | --pst-font-family-base: "Mona Sans", sans-serif;
52 | }
53 |
54 | body {
55 | font-weight: 450;
56 | }
57 |
58 | .bd-main .bd-content .bd-article-container {
59 | max-width: 100%; /* default is 60em */
60 | }
61 |
62 | .bd-sidebar-primary {
63 | width: 18%;
64 | }
--------------------------------------------------------------------------------
/docs/_templates/autoapi/index.rst:
--------------------------------------------------------------------------------
1 | API Reference
2 | =============
3 |
4 | .. toctree::
5 | :titlesonly:
6 |
7 | {% for page in pages|selectattr("is_top_level_object") %}
8 | {{ page.include_path }}
9 | {% endfor %}
10 |
--------------------------------------------------------------------------------
/docs/_templates/autoapi/python/attribute.rst:
--------------------------------------------------------------------------------
1 | {% extends "python/data.rst" %}
2 |
--------------------------------------------------------------------------------
/docs/_templates/autoapi/python/class.rst:
--------------------------------------------------------------------------------
1 | {% if obj.display %}
2 | {% if is_own_page %}
3 | {{ obj.name }}
4 | {{ "=" * obj.name | length }}
5 |
6 | {% endif %}
7 | {% set visible_children = obj.children|selectattr("display")|list %}
8 | {% set own_page_children = visible_children|selectattr("type", "in", own_page_types)|list %}
9 | {% if is_own_page and own_page_children %}
10 | .. toctree::
11 | :hidden:
12 |
13 | {% for child in own_page_children %}
14 | {{ child.include_path }}
15 | {% endfor %}
16 |
17 | {% endif %}
18 | .. py:{{ obj.type }}:: {% if is_own_page %}{{ obj.id }}{% else %}{{ obj.short_name }}{% endif %}{% if obj.args %}({{ obj.args }}){% endif %}
19 |
20 | {% for (args, return_annotation) in obj.overloads %}
21 | {{ " " * (obj.type | length) }} {{ obj.short_name }}{% if args %}({{ args }}){% endif %}
22 |
23 | {% endfor %}
24 | {% if obj.bases %}
25 | {% if "show-inheritance" in autoapi_options %}
26 |
27 | Bases: {% for base in obj.bases %}{{ base|link_objs }}{% if not loop.last %}, {% endif %}{% endfor %}
28 | {% endif %}
29 |
30 |
31 | {% if "show-inheritance-diagram" in autoapi_options and obj.bases != ["object"] %}
32 | .. autoapi-inheritance-diagram:: {{ obj.obj["full_name"] }}
33 | :parts: 1
34 | {% if "private-members" in autoapi_options %}
35 | :private-bases:
36 | {% endif %}
37 |
38 | {% endif %}
39 | {% endif %}
40 | {% if obj.docstring %}
41 |
42 | {{ obj.docstring|indent(3) }}
43 | {% endif %}
44 | {% for obj_item in visible_children %}
45 | {% if obj_item.type not in own_page_types %}
46 |
47 | {{ obj_item.render()|indent(3) }}
48 | {% endif %}
49 | {% endfor %}
50 | {% if is_own_page and own_page_children %}
51 | {% set visible_attributes = own_page_children|selectattr("type", "equalto", "attribute")|list %}
52 | {% if visible_attributes %}
53 | Attributes
54 | ----------
55 |
56 | .. autoapisummary::
57 |
58 | {% for attribute in visible_attributes %}
59 | {{ attribute.id }}
60 | {% endfor %}
61 |
62 |
63 | {% endif %}
64 | {% set visible_properties = own_page_children|selectattr("type", "equalto", "property")|list %}
65 | {% if visible_properties %}
66 | Properties
67 | ----------
68 |
69 | .. autoapisummary::
70 |
71 | {% for property in visible_properties %}
72 | {{ property.id }}
73 | {% endfor %}
74 |
75 |
76 | {% endif %}
77 | {% set visible_exceptions = own_page_children|selectattr("type", "equalto", "exception")|list %}
78 | {% if visible_exceptions %}
79 | Exceptions
80 | ----------
81 |
82 | .. autoapisummary::
83 |
84 | {% for exception in visible_exceptions %}
85 | {{ exception.id }}
86 | {% endfor %}
87 |
88 |
89 | {% endif %}
90 | {% set visible_classes = own_page_children|selectattr("type", "equalto", "class")|list %}
91 | {% if visible_classes %}
92 | Classes
93 | -------
94 |
95 | .. autoapisummary::
96 |
97 | {% for klass in visible_classes %}
98 | {{ klass.id }}
99 | {% endfor %}
100 |
101 |
102 | {% endif %}
103 | {% set visible_methods = own_page_children|selectattr("type", "equalto", "method")|list %}
104 | {% if visible_methods %}
105 | Methods
106 | -------
107 |
108 | .. autoapisummary::
109 |
110 | {% for method in visible_methods %}
111 | {{ method.id }}
112 | {% endfor %}
113 |
114 |
115 | {% endif %}
116 |
117 | {% endif %}
118 | {% endif %}
119 |
120 |
121 | .. footbibliography::
--------------------------------------------------------------------------------
/docs/_templates/autoapi/python/data.rst:
--------------------------------------------------------------------------------
1 | {% if obj.display %}
2 | {% if is_own_page %}
3 | {{ obj.name }}
4 | {{ "=" * obj.name | length }}
5 |
6 | {% endif %}
7 | .. py:{{ obj.type }}:: {% if is_own_page %}{{ obj.id }}{% else %}{{ obj.name }}{% endif %}
8 | {% if obj.annotation is not none %}
9 |
10 | :type: {% if obj.annotation %} {{ obj.annotation }}{% endif %}
11 | {% endif %}
12 | {% if obj.value is not none %}
13 |
14 | {% if obj.value.splitlines()|count > 1 %}
15 | :value: Multiline-String
16 |
17 | .. raw:: html
18 |
19 | Show Value
20 |
21 | .. code-block:: python
22 |
23 | {{ obj.value|indent(width=6,blank=true) }}
24 |
25 | .. raw:: html
26 |
27 |
4 |
5 |
6 |
7 |
8 |
9 | Multi-person, multi-camera 3D human pose visualization tool built using
10 | [Mayavi](https://docs.enthought.com/mayavi/mayavi/). As used
11 | in [MeTRAbs](https://github.com/isarandi/metrabs) visualizations.
12 |
13 | ## Gist of Usage
14 |
15 | ```python
16 | import poseviz
17 |
18 | viz = poseviz.PoseViz(...)
19 | camera = poseviz.Camera(...)
20 | for frame in frames:
21 | bounding_boxes, poses3d = run_pose_estimation_model(frame)
22 | viz.update(frame=frame, boxes=bounding_boxes, poses=poses3d, camera=camera)
23 | ```
24 |
25 | The main feature of this tool is that the graphical event loop is hidden from the library user. We
26 | want to write code in terms of the *prediction loop* of the human pose estimator, not from the point
27 | of view of the visualizer tool.
28 |
29 | Behind the scenes, this is achieved through forking a dedicated visualization process and passing
30 | new scene information via multiprocessing queues.
31 |
32 | ## Installation
33 |
34 | PoseViz is released as a conda package (experimental, tested only on Linux):
35 |
36 | ```bash
37 | conda install poseviz -c isarandi
38 | ```
39 |
40 | Alternatively, in case the above doesn't work, install Mayavi via conda (the Mayavi pip package has
41 | compilation problems), clone this repo and install PoseViz via pip:
42 |
43 | ```bash
44 | conda install mayavi -c conda-forge
45 | pip install git+https://github.com/isarandi/poseviz.git
46 | ```
--------------------------------------------------------------------------------
/docs_old/poseviz.md:
--------------------------------------------------------------------------------
1 | ::: poseviz.PoseViz
2 | rendering:
3 | show_root_heading: no
4 | members_order: source
5 |
--------------------------------------------------------------------------------
/docs_old/stylesheets/extra.css:
--------------------------------------------------------------------------------
1 | /* Indentation. */
2 | div.doc-contents:not(.first) {
3 | padding-left: 25px;
4 | border-left: 4px solid rgba(230, 230, 230);
5 | margin-bottom: 60px;
6 | }
--------------------------------------------------------------------------------
/pyproject.toml:
--------------------------------------------------------------------------------
1 | [build-system]
2 | requires = [
3 | "setuptools>=61",
4 | "wheel",
5 | "setuptools_scm[toml]>=8"
6 | ]
7 | build-backend = "setuptools.build_meta"
8 |
9 | [project]
10 | name = "poseviz"
11 | dynamic = ["version"]
12 | description = "3D human pose visualizer with multi-person, multi-view support, built on Mayavi."
13 | authors = [
14 | { name = "István Sárándi", email = "istvan.sarandi@uni-tuebingen.de" }
15 | ]
16 | readme = "README.md"
17 | requires-python = ">=3.8"
18 | license = { file = "LICENSE" }
19 |
20 | dependencies = [
21 | 'more-itertools',
22 | 'opencv-python',
23 | 'numpy',
24 | 'mayavi',
25 | 'imageio',
26 | 'cameravision',
27 | 'boxlib',
28 | 'framepump',
29 | ]
30 |
31 | [tool.setuptools]
32 | package-dir = { "" = "src" }
33 |
34 | [tool.setuptools.packages.find]
35 | where = ["src"]
36 |
37 | [tool.black]
38 | line-length = 99
39 | skip-string-normalization = true
40 |
--------------------------------------------------------------------------------
/screenshot.jpg:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/isarandi/poseviz/e642db70400e75e089607c319d70436919f5a190/screenshot.jpg
--------------------------------------------------------------------------------
/screenshot2.jpg:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/isarandi/poseviz/e642db70400e75e089607c319d70436919f5a190/screenshot2.jpg
--------------------------------------------------------------------------------
/screenshot_multicam.jpg:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/isarandi/poseviz/e642db70400e75e089607c319d70436919f5a190/screenshot_multicam.jpg
--------------------------------------------------------------------------------
/src/poseviz/__init__.py:
--------------------------------------------------------------------------------
1 | """PoseViz is a 3D visualization tool for human pose estimation."""
2 |
3 | from poseviz.poseviz import PoseViz
4 | from poseviz.view_info import ViewInfo
5 |
--------------------------------------------------------------------------------
/src/poseviz/colors.py:
--------------------------------------------------------------------------------
1 | import itertools
2 |
3 | green = (0.17254901960784313, 0.62745098039215685, 0.17254901960784313)
4 | red = (0.83921568627450982, 0.15294117647058825, 0.15686274509803921)
5 | gray = (0.49803921568627452, 0.49803921568627452, 0.49803921568627452)
6 | cyan = (0.090196078431372548, 0.74509803921568629, 0.81176470588235294)
7 | yellow = (0.73725490196078436, 0.74117647058823533, 0.13333333333333333)
8 | blue = (0.12156862745098039, 0.46666666666666667, 0.70588235294117652)
9 | orange = (1.0, 0.49803921568627452, 0.054901960784313725)
10 | black = (0.0, 0.0, 0.0)
11 |
12 |
13 | def cycle_over_colors(range_zero_one=False):
14 | """Returns a generator that cycles over a list of nice colors, indefinitely."""
15 | colors = (
16 | (0.12156862745098039, 0.46666666666666667, 0.70588235294117652),
17 | (1.0, 0.49803921568627452, 0.054901960784313725),
18 | (0.17254901960784313, 0.62745098039215685, 0.17254901960784313),
19 | (0.83921568627450982, 0.15294117647058825, 0.15686274509803921),
20 | (0.58039215686274515, 0.40392156862745099, 0.74117647058823533),
21 | (0.5490196078431373, 0.33725490196078434, 0.29411764705882354),
22 | (0.8901960784313725, 0.46666666666666667, 0.76078431372549016),
23 | (0.49803921568627452, 0.49803921568627452, 0.49803921568627452),
24 | (0.73725490196078436, 0.74117647058823533, 0.13333333333333333),
25 | (0.090196078431372548, 0.74509803921568629, 0.81176470588235294),
26 | )
27 |
28 | if not range_zero_one:
29 | colors = [[c * 255 for c in color] for color in colors]
30 |
31 | return itertools.cycle(colors)
32 |
--------------------------------------------------------------------------------
/src/poseviz/components/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/isarandi/poseviz/e642db70400e75e089607c319d70436919f5a190/src/poseviz/components/__init__.py
--------------------------------------------------------------------------------
/src/poseviz/components/camera_viz.py:
--------------------------------------------------------------------------------
1 | import numpy as np
2 | import tvtk.api
3 | from mayavi import mlab
4 | from scipy.spatial.transform import Rotation
5 | import poseviz.colors
6 | import poseviz.mayavi_util
7 |
8 |
9 | class CameraViz:
10 | def __init__(
11 | self, camera_type, show_image, show_field_of_view=True, show_camera_wireframe=True
12 | ):
13 | self.viz_im = None
14 | self.is_initialized = False
15 | self.prev_cam = None
16 | self.camera_type = camera_type
17 | self.show_image = show_image
18 | self.prev_imshape = None
19 | self.show_field_of_view = show_field_of_view
20 | self.show_camera_wireframe = show_camera_wireframe
21 | self.mesh = None
22 | self.mesh2 = None
23 | self.mesh3 = None
24 | self.prev_highlight = False
25 |
26 | def initial_update(self, camera, image, highlight=False):
27 | image_corners, far_corners = self.calculate_camera_vertices(camera, image.shape)
28 |
29 | if self.camera_type != "original":
30 | triangles = np.array([[0, 1, 2], [0, 2, 3], [0, 3, 4], [0, 4, 1]])
31 |
32 | if self.show_camera_wireframe:
33 | self.mesh = mlab.triangular_mesh(
34 | *image_corners[:5].T,
35 | triangles,
36 | color=poseviz.colors.cyan if highlight else poseviz.colors.black,
37 | tube_radius=0.04 if highlight else 0.02,
38 | line_width=10 if highlight else 1,
39 | tube_sides=3,
40 | representation="wireframe",
41 | reset_zoom=False,
42 | )
43 |
44 | if self.show_field_of_view:
45 | self.mesh2 = mlab.triangular_mesh(
46 | *far_corners.T,
47 | triangles,
48 | color=poseviz.colors.black,
49 | opacity=0.1,
50 | representation="surface",
51 | reset_zoom=False,
52 | )
53 | self.mesh3 = mlab.triangular_mesh(
54 | *far_corners.T,
55 | triangles,
56 | color=poseviz.colors.gray,
57 | opacity=0.1,
58 | tube_radius=0.01,
59 | tube_sides=3,
60 | representation="wireframe",
61 | reset_zoom=False,
62 | )
63 |
64 | if self.show_image:
65 | self.new_imshow(image.shape, image_corners)
66 | self.set_image_content(image)
67 | self.set_image_position(image_corners)
68 |
69 | self.prev_cam = camera
70 | self.prev_imshape = image.shape
71 | self.prev_highlight = highlight
72 | self.is_initialized = True
73 |
74 | def update(self, camera, image, highlight=False):
75 | if not self.is_initialized:
76 | return self.initial_update(camera, image, highlight)
77 |
78 | if (
79 | self.prev_cam.allclose(camera)
80 | and image.shape == self.prev_imshape
81 | and self.prev_highlight == highlight
82 | ):
83 | # Only the image content has changed
84 | if self.show_image:
85 | self.set_image_content(image)
86 | return
87 | elif self.prev_highlight == highlight:
88 | image_points, far_corners = self.calculate_camera_vertices(camera, image.shape)
89 | if self.camera_type != "original":
90 | if self.show_camera_wireframe:
91 | self.mesh.mlab_source.set(
92 | x=image_points[:5, 0], y=image_points[:5, 1], z=image_points[:5, 2]
93 | )
94 | if self.show_field_of_view:
95 | self.mesh2.mlab_source.set(
96 | x=far_corners[:, 0], y=far_corners[:, 1], z=far_corners[:, 2]
97 | )
98 | self.mesh3.mlab_source.set(
99 | x=far_corners[:, 0], y=far_corners[:, 1], z=far_corners[:, 2]
100 | )
101 |
102 | if self.show_image:
103 | if image.shape[:2] != self.prev_imshape[:2] or not np.allclose(
104 | camera.intrinsic_matrix, self.prev_cam.intrinsic_matrix
105 | ):
106 | self.viz_im.remove()
107 | self.new_imshow(image.shape, image_points)
108 | self.prev_imshape = image.shape
109 |
110 | self.set_image_content(image)
111 | self.set_image_position(image_points)
112 | self.prev_cam = camera
113 | else:
114 | # We change the highlight state by reinitializing it all
115 | for elem in [self.viz_im, self.mesh, self.mesh2, self.mesh3]:
116 | if elem is not None:
117 | elem.remove()
118 | self.initial_update(camera, image, highlight)
119 |
120 | def new_imshow(self, imshape, mayavi_image_points):
121 | mayavi_width = np.linalg.norm(mayavi_image_points[1] - mayavi_image_points[2])
122 | mayavi_height = np.linalg.norm(mayavi_image_points[2] - mayavi_image_points[3])
123 | extent = [0, mayavi_height, 0, mayavi_width, 0, 0]
124 | self.viz_im = mlab.imshow(
125 | np.ones(imshape[:2]), opacity=0.5, extent=extent, reset_zoom=False, interpolate=True
126 | )
127 |
128 | def calculate_camera_vertices(self, camera, imshape):
129 | h, w = imshape[:2]
130 | image_corners = np.array([[0, 0], [w - 1, 0], [w - 1, h - 1], [0, h - 1]], np.float32)
131 | image_center = np.array(imshape[:2][::-1]) / 2
132 | image_points = np.concatenate([image_corners, [image_center]], axis=0)
133 | image_points_world = camera.image_to_world(image_points, camera_depth=150)
134 | points = np.array([camera.t, *image_points_world])
135 | mayavi_image_points = poseviz.mayavi_util.world_to_mayavi(points)
136 | mayavi_far_corners = (
137 | mayavi_image_points[0] + (mayavi_image_points[:5] - mayavi_image_points[0]) * 20
138 | )
139 | return mayavi_image_points, mayavi_far_corners
140 |
141 | def set_image_content(self, image):
142 | reshaped = image.view().reshape([-1, 3], order="F")
143 | self.colors = tvtk.api.tvtk.UnsignedCharArray()
144 | self.colors.from_array(reshaped)
145 | self.viz_im.actor.input.point_data.scalars = self.colors
146 |
147 | def set_image_position(self, mayavi_image_points):
148 | down = unit_vector(mayavi_image_points[4] - mayavi_image_points[1])
149 | right = unit_vector(mayavi_image_points[2] - mayavi_image_points[1])
150 | R = np.column_stack([down, right, np.cross(down, right)])
151 | z, x, y = Rotation.from_matrix(R).as_euler("ZXY", degrees=True)
152 | self.viz_im.actor.orientation = [x, y, z]
153 | self.viz_im.actor.position = mayavi_image_points[5]
154 |
155 | def remove(self):
156 | if self.is_initialized:
157 | for elem in [self.viz_im, self.mesh, self.mesh2, self.mesh3]:
158 | if elem is not None:
159 | elem.remove()
160 | self.is_initialized = False
161 |
162 |
163 | def unit_vector(v):
164 | return v / np.linalg.norm(v)
165 |
--------------------------------------------------------------------------------
/src/poseviz/components/connected_points_viz.py:
--------------------------------------------------------------------------------
1 | import numpy as np
2 | from mayavi import mlab
3 |
4 |
5 | class ConnectedPoints:
6 | def __init__(
7 | self, color, line_color, mode, scale_factor, opacity_points, opacity_lines,
8 | point_names, high_quality=False):
9 | self.coords = []
10 | self.n_rendered_points = 0
11 | self.color = color
12 | self.mode = mode
13 | self.scale_factor = scale_factor
14 | self.line_color = line_color
15 | self.opacity_points = opacity_points
16 | self.opacity_lines = opacity_lines
17 | self.is_initialized = False
18 | self.points = None
19 | self.tube = None
20 | self.edges = []
21 | self.point_name_components = []
22 | self.point_names = point_names
23 | self.high_quality = high_quality
24 |
25 | def clear(self):
26 | self.coords.clear()
27 | self.edges.clear()
28 |
29 | def add_points(self, coords, edges, show_isolated_points=False):
30 | if not show_isolated_points:
31 | coords, edges = self.get_nonisolated(coords, edges)
32 | n_points = len(self.coords)
33 | self.coords.extend(coords)
34 | self.edges.extend([[n_points + i, n_points + j] for i, j in edges])
35 |
36 | def get_nonisolated(self, coords, edges):
37 | selected_coords = []
38 | new_point_ids = {}
39 | selected_edges = []
40 |
41 | for (i, j) in edges:
42 | if i not in new_point_ids:
43 | new_point_ids[i] = len(selected_coords)
44 | selected_coords.append(coords[i])
45 | if j not in new_point_ids:
46 | new_point_ids[j] = len(selected_coords)
47 | selected_coords.append(coords[j])
48 | selected_edges.append([new_point_ids[i], new_point_ids[j]])
49 | return selected_coords, selected_edges
50 |
51 | def update(self):
52 | if not self.is_initialized:
53 | return self.initial_update()
54 | c = np.asarray(self.coords)
55 |
56 | if not self.coords:
57 | c = c.reshape((0, 3))
58 |
59 | if len(self.coords) == self.n_rendered_points:
60 | self.points.mlab_source.set(x=c[:, 0], y=c[:, 1], z=c[:, 2])
61 | else:
62 | self.n_rendered_points = len(self.coords)
63 | self.points.mlab_source.reset(x=c[:, 0], y=c[:, 1], z=c[:, 2])
64 | self.points.mlab_source.dataset.lines = self.edges
65 |
66 | def initial_update(self):
67 | if not self.coords:
68 | return
69 | c = np.asarray(self.coords)
70 |
71 | self.points = mlab.points3d(
72 | *c.T, scale_factor=self.scale_factor, color=self.color,
73 | opacity=self.opacity_points,
74 | mode=self.mode, resolution=16 if self.high_quality else 3, scale_mode='vector',
75 | reset_zoom=False)
76 | self.points.mlab_source.dataset.lines = self.edges
77 | self.tube = mlab.pipeline.tube(
78 | self.points, tube_radius=self.scale_factor / 5,
79 | tube_sides=12 if self.high_quality else 3)
80 | mlab.pipeline.surface(
81 | self.tube, color=self.line_color, opacity=self.opacity_lines, reset_zoom=False)
82 |
83 | # for i, (point, name, color) in enumerate(
84 | # zip(c, self.point_names, colors.cycle_over_colors())):
85 | # self.point_name_components.append(mlab.text3d(
86 | # *point, f'{i}', scale=0.02, color=color))
87 |
88 | self.is_initialized = True
89 |
90 | def remove(self):
91 | if self.is_initialized:
92 | self.points.remove()
93 | self.tube.remove()
94 | for component in self.point_name_components:
95 | component.remove()
96 | self.is_initialized = False
97 | self.clear()
98 | self.point_name_components.clear()
99 | self.n_rendered_points = 0
--------------------------------------------------------------------------------
/src/poseviz/components/ground_plane_viz.py:
--------------------------------------------------------------------------------
1 | import numpy as np
2 | import poseviz.mayavi_util
3 | import tvtk.api
4 | from mayavi import mlab
5 |
6 |
7 | def draw_checkerboard(ground_plane_height):
8 | i, j = np.mgrid[:40, :40]
9 | image = ((i + j) % 2 == 0).astype(np.uint8)
10 | image[image == 0] = 0
11 | image[image == 1] = 255 - 96
12 |
13 | # image[image == 0] = 96
14 | # image[image == 1] = 256-96
15 | image = np.expand_dims(image, -1) * np.ones(3)
16 |
17 | size = 2 * 19
18 | extent = [0, size, 0, size, 0, 0]
19 | viz_im = mlab.imshow(
20 | np.ones(image.shape[:2]), opacity=0.3, extent=extent, reset_zoom=False, interpolate=False
21 | )
22 |
23 | reshaped = image.reshape([-1, 3], order="F")
24 | colors = tvtk.api.tvtk.UnsignedCharArray()
25 | colors.number_of_components = 3
26 | colors.from_array(reshaped)
27 | viz_im.actor.input.point_data.scalars = colors
28 | viz_im.actor.orientation = [0, 0, 90]
29 | viz_im.actor.position = np.array([-1, 1, ground_plane_height * poseviz.mayavi_util.MM_TO_UNIT])
30 |
--------------------------------------------------------------------------------
/src/poseviz/components/main_viz.py:
--------------------------------------------------------------------------------
1 | import logging
2 |
3 | import poseviz.colors as colors
4 | import poseviz.components.camera_viz
5 | import poseviz.components.skeletons_viz
6 | import poseviz.components.smpl_viz
7 |
8 |
9 | class MainViz:
10 | def __init__(
11 | self,
12 | joint_info_pred,
13 | joint_info_true,
14 | joint_info_alt,
15 | camera_type,
16 | show_image,
17 | high_quality,
18 | show_field_of_view=True,
19 | show_camera_wireframe=True,
20 | body_model_faces=None,
21 | ):
22 |
23 | if joint_info_pred[0] is not None:
24 | self.skeletons_pred = poseviz.components.skeletons_viz.SkeletonsViz(
25 | joint_info_pred,
26 | colors.blue,
27 | colors.cyan,
28 | colors.yellow,
29 | colors.green,
30 | 0.06,
31 | high_quality,
32 | opacity=0.95,
33 | )
34 | else:
35 | self.skeletons_pred = None
36 |
37 | if joint_info_true[0] is not None:
38 | self.skeletons_true = poseviz.components.skeletons_viz.SkeletonsViz(
39 | joint_info_true,
40 | colors.red,
41 | colors.red,
42 | colors.red,
43 | colors.red,
44 | 0.06,
45 | high_quality,
46 | opacity=0.95,
47 | )
48 | else:
49 | self.skeletons_true = None
50 |
51 | if joint_info_alt[0] is not None:
52 | self.skeletons_alt = poseviz.components.skeletons_viz.SkeletonsViz(
53 | joint_info_alt,
54 | colors.orange,
55 | colors.orange,
56 | colors.orange,
57 | colors.orange,
58 | 0.06,
59 | high_quality,
60 | opacity=0.95,
61 | )
62 | else:
63 | self.skeletons_alt = None
64 |
65 | if body_model_faces is not None:
66 | self.meshes_pred = poseviz.components.smpl_viz.SMPLViz(
67 | color=colors.blue, faces=body_model_faces, add_wireframe=True, colormap="Blues_r"
68 | )
69 | self.meshes_gt = poseviz.components.smpl_viz.SMPLViz(
70 | color=colors.red, faces=body_model_faces, add_wireframe=True
71 | )
72 | self.meshes_pred2 = poseviz.components.smpl_viz.SMPLViz(
73 | color=colors.orange,
74 | faces=body_model_faces,
75 | add_wireframe=True,
76 | colormap="Oranges_r",
77 | )
78 | else:
79 | self.meshes_pred = None
80 | self.meshes_gt = None
81 | self.meshes_pred2 = None
82 |
83 | self.camera_viz = poseviz.components.camera_viz.CameraViz(
84 | camera_type, show_image, show_field_of_view, show_camera_wireframe
85 | )
86 |
87 | def update(
88 | self,
89 | camera_display,
90 | image,
91 | poses_pred=None,
92 | poses_true=None,
93 | poses_alt=None,
94 | vertices_pred=None,
95 | vertices_true=None,
96 | vertices_alt=None,
97 | highlight=False,
98 | ):
99 |
100 | if poses_pred is not None and self.skeletons_pred is not None:
101 | self.skeletons_pred.update(poses_pred)
102 | if poses_true is not None and self.skeletons_true is not None:
103 | self.skeletons_true.update(poses_true)
104 | if poses_alt is not None and self.skeletons_alt is not None:
105 | self.skeletons_alt.update(poses_alt)
106 |
107 | if vertices_pred is not None and self.meshes_pred is not None:
108 | self.meshes_pred.update(vertices_pred)
109 | if vertices_true is not None and self.meshes_gt is not None:
110 | self.meshes_gt.update(vertices_true)
111 | if vertices_alt is not None and self.meshes_pred2 is not None:
112 | self.meshes_pred2.update(vertices_alt)
113 |
114 | self.camera_viz.update(camera_display, image, highlight=highlight)
115 |
116 | def remove(self):
117 | if self.is_initialized:
118 | for elem in [
119 | self.skeletons_pred,
120 | self.skeletons_true,
121 | self.skeletons_alt,
122 | self.meshes_pred,
123 | self.meshes_gt,
124 | self.meshes_pred2,
125 | ]:
126 | if elem is not None:
127 | elem.remove()
128 |
129 | self.camera_viz.remove()
130 |
131 | self.is_initialized = False
132 |
--------------------------------------------------------------------------------
/src/poseviz/components/skeletons_viz.py:
--------------------------------------------------------------------------------
1 | import poseviz.components.connected_points_viz
2 | import poseviz.mayavi_util
3 |
4 |
5 | class SkeletonsViz:
6 | def __init__(
7 | self,
8 | joint_info,
9 | left_color,
10 | mid_color,
11 | right_color,
12 | point_color,
13 | scale_factor,
14 | high_quality=False,
15 | opacity=0.7,
16 | ):
17 |
18 | joint_names, joint_edges = joint_info
19 | edge_colors = dict(left=left_color, mid=mid_color, right=right_color)
20 | sides = ("left", "mid", "right")
21 | self.pointsets = {
22 | side: poseviz.components.connected_points_viz.ConnectedPoints(
23 | point_color,
24 | edge_colors[side],
25 | "sphere",
26 | scale_factor,
27 | opacity,
28 | opacity,
29 | joint_names,
30 | high_quality,
31 | )
32 | for side in sides
33 | }
34 |
35 | def edge_side(edge):
36 | s1 = joint_side(edge[0])
37 | s2 = joint_side(edge[1])
38 | if s1 == "left" or s2 == "left":
39 | return "left"
40 | elif s1 == "right" or s2 == "right":
41 | return "right"
42 | else:
43 | return "mid"
44 |
45 | def joint_side(joint):
46 | name = joint_names[joint].lower()
47 | if name.startswith("l"):
48 | return "left"
49 | elif name.startswith("r"):
50 | return "right"
51 | else:
52 | return "mid"
53 |
54 | self.edges_per_side = {
55 | side: [e for e in joint_edges if edge_side(e) == side] for side in sides
56 | }
57 |
58 | self.indices_per_side = {
59 | side: sorted(
60 | set([index for edge in self.edges_per_side[side] for index in edge])
61 | | set(
62 | [index for index, name in enumerate(joint_names) if joint_side(index) == side]
63 | )
64 | )
65 | for side in sides
66 | }
67 |
68 | def index_within_side(global_index, side):
69 | return self.indices_per_side[side].index(global_index)
70 |
71 | self.edges_within_side = {
72 | side: [
73 | (index_within_side(i, side), index_within_side(j, side))
74 | for i, j in self.edges_per_side[side]
75 | ]
76 | for side in sides
77 | }
78 |
79 | def update(self, poses):
80 | mayavi_poses = [poseviz.mayavi_util.world_to_mayavi(pose) for pose in poses]
81 | for side, pointset in self.pointsets.items():
82 | pointset.clear()
83 | for coords in mayavi_poses:
84 | pointset.add_points(
85 | coords[self.indices_per_side[side]],
86 | self.edges_within_side[side],
87 | show_isolated_points=True,
88 | )
89 | pointset.update()
90 |
91 | def remove(self):
92 | for pointset in self.pointsets.values():
93 | pointset.remove()
94 |
--------------------------------------------------------------------------------
/src/poseviz/components/smpl_viz.py:
--------------------------------------------------------------------------------
1 | import matplotlib.pyplot as plt
2 | import numpy as np
3 | from mayavi import mlab
4 |
5 | import poseviz.mayavi_util
6 |
7 |
8 | class SMPLViz:
9 | def __init__(self, color, faces, add_wireframe=False, colormap="viridis"):
10 | self.opacity = 0.6
11 | self.faces = faces
12 | self.is_initialized = False
13 | self.n_rendered_verts = 0
14 | self.color = color
15 | self.add_wireframe = add_wireframe
16 | self.colormap = colormap
17 |
18 | def update(self, poses):
19 | if not self.is_initialized:
20 | return self.initial_update(poses)
21 |
22 | mayavi_poses = [poseviz.mayavi_util.world_to_mayavi(pose[:, :3]) for pose in poses]
23 | uncerts = [pose[:, 3] if pose.shape[1] == 4 else None for pose in poses]
24 |
25 | if len(mayavi_poses) == 0:
26 | c = np.array([[6, -3, 2]], np.float32)
27 | u = np.array([0], np.float32)
28 | else:
29 | c = np.concatenate(mayavi_poses, axis=0)
30 | if uncerts[0] is not None:
31 | u = np.concatenate(uncerts, axis=0)
32 | else:
33 | u = None
34 |
35 | if len(c) == self.n_rendered_verts:
36 | # Number of vertices is the same as prev frame, so we can just update the positions
37 | if u is not None:
38 | self.mesh.mlab_source.set(x=c[:, 0], y=c[:, 1], z=c[:, 2], scalars=u)
39 | else:
40 | self.mesh.mlab_source.set(x=c[:, 0], y=c[:, 1], z=c[:, 2])
41 | else:
42 | self.n_rendered_verts = len(c)
43 |
44 | if len(mayavi_poses) == 0:
45 | triangles = np.zeros([1, 3], np.int32)
46 | else:
47 | num_vertices = poses[0].shape[0]
48 | triangles = np.concatenate(
49 | [np.asarray(self.faces) + i * num_vertices for i in range(len(mayavi_poses))],
50 | axis=0,
51 | )
52 |
53 | if u is not None:
54 | self.mesh.mlab_source.reset(
55 | x=c[:, 0], y=c[:, 1], z=c[:, 2], triangles=triangles, scalars=u
56 | )
57 | else:
58 | self.mesh.mlab_source.reset(x=c[:, 0], y=c[:, 1], z=c[:, 2], triangles=triangles)
59 |
60 | def initial_update(self, poses):
61 | if len(poses) == 0:
62 | return
63 |
64 | mayavi_poses = [poseviz.mayavi_util.world_to_mayavi(pose[:, :3]) for pose in poses]
65 | uncerts = [pose[:, 3] if pose.shape[1] == 4 else None for pose in poses]
66 |
67 | num_vertices = poses[0].shape[0]
68 | triangles = np.concatenate(
69 | [np.asarray(self.faces) + i * num_vertices for i in range(len(mayavi_poses))], axis=0
70 | )
71 | c = np.concatenate(mayavi_poses, axis=0)
72 | if uncerts[0] is not None:
73 | u = np.concatenate(uncerts, axis=0)
74 | else:
75 | u = None
76 |
77 | if u is not None:
78 | self.mesh = mlab.triangular_mesh(
79 | *c.T,
80 | triangles,
81 | scalars=u,
82 | colormap="cool",
83 | vmin=0,
84 | vmax=0.1,
85 | opacity=self.opacity,
86 | representation="surface",
87 | reset_zoom=False,
88 | )
89 |
90 | cmap = plt.get_cmap(self.colormap)
91 | cmaplist = np.array([cmap(i) for i in range(cmap.N)]) * 255
92 | self.mesh.module_manager.scalar_lut_manager.lut.table = cmaplist
93 | else:
94 | self.mesh = mlab.triangular_mesh(
95 | *c.T,
96 | triangles,
97 | color=self.color,
98 | opacity=self.opacity,
99 | representation="surface",
100 | reset_zoom=False,
101 | )
102 |
103 | attrs = dict(backface_culling=True)
104 |
105 | if self.add_wireframe:
106 | attrs.update(edge_visibility=True, edge_color=(0.0, 0.0, 0.0), line_width=0.1)
107 |
108 | for key, value in attrs.items():
109 | setattr(self.mesh.actor.property, key, value)
110 |
111 | self.is_initialized = True
112 | self.n_rendered_meshes = len(mayavi_poses)
113 |
114 | def remove(self):
115 | if self.is_initialized:
116 | self.mesh.remove()
117 | self.is_initialized = False
118 |
--------------------------------------------------------------------------------
/src/poseviz/draw2d.py:
--------------------------------------------------------------------------------
1 | import cv2
2 | import numpy as np
3 |
4 |
5 | def rounded_int_tuple(p):
6 | return tuple(np.round(p).astype(int))
7 |
8 |
9 | def line(im, p1, p2, *args, **kwargs):
10 | cv2.line(im, rounded_int_tuple(p1), rounded_int_tuple(p2), *args, **kwargs)
11 |
12 |
13 | def draw_stick_figure_2d_inplace(im, coords, joint_edges, thickness=3, color=None):
14 | for i_joint1, i_joint2 in joint_edges:
15 | relevant_coords = coords[[i_joint1, i_joint2]]
16 | if not np.isnan(relevant_coords).any() and not np.isclose(0, relevant_coords).any():
17 | line(
18 | im,
19 | coords[i_joint1],
20 | coords[i_joint2],
21 | color=color,
22 | thickness=thickness,
23 | lineType=cv2.LINE_AA,
24 | )
25 |
26 |
27 | def rectangle(im, pt1, pt2, color, thickness):
28 | cv2.rectangle(im, rounded_int_tuple(pt1), rounded_int_tuple(pt2), color, thickness)
29 |
30 |
31 | def draw_box(im, box, color=(255, 0, 0), thickness=5):
32 | box = np.array(box)
33 | rectangle(im, box[:2], box[:2] + box[2:4], color, thickness)
34 |
35 |
36 | def resize_by_factor(im, factor, interp=None, dst=None):
37 | """Returns a copy of `im` resized by `factor`, using bilinear interp for up and area interp
38 | for downscaling.
39 | """
40 | new_size = rounded_int_tuple([im.shape[1] * factor, im.shape[0] * factor])
41 | if interp is None:
42 | interp = cv2.INTER_LINEAR if factor > 1.0 else cv2.INTER_AREA
43 | return cv2.resize(im, new_size, fx=factor, fy=factor, interpolation=interp, dst=dst)
44 |
45 |
46 |
47 | def resize(im, dst, interp=None):
48 | """Returns a copy of `im` resized by `factor`, using bilinear interp for up and area interp
49 | for downscaling.
50 | """
51 | if interp is None:
52 | interp = cv2.INTER_LINEAR if dst.shape[0] > im.shape[0] else cv2.INTER_AREA
53 | return cv2.resize(im, (dst.shape[1], dst.shape[0]), interpolation=interp, dst=dst)
54 |
--------------------------------------------------------------------------------
/src/poseviz/init.py:
--------------------------------------------------------------------------------
1 | import vtk
2 | from mayavi import mlab
3 | from tvtk.api import tvtk
4 | import ctypes
5 | import signal
6 |
7 |
8 | def initialize(size=(1280, 720)):
9 | fig = mlab.figure('PoseViz', bgcolor=(1, 1, 1), fgcolor=(0, 0, 0), size=size)
10 | fig.scene.renderer.render_window.set(alpha_bit_planes=1, multi_samples=0, full_screen=False)
11 | fig.scene.renderer.set(use_depth_peeling=True, maximum_number_of_peels=5, occlusion_ratio=0.01)
12 |
13 | fig.scene._tool_bar.setVisible(False)
14 | fig.scene.interactor.interactor_style = tvtk.InteractorStyleTerrain()
15 | # Suppress warning outputs that have no effect on the visualization.
16 | output = vtk.vtkFileOutputWindow()
17 | output.SetFileName('/dev/null')
18 | vtk.vtkOutputWindow().SetInstance(output)
19 | vtk.vtkObject.GlobalWarningDisplayOff()
20 | fig.scene.camera.clipping_range = [0.1, 1000]
21 | fig.scene.anti_aliasing_frames = 0
22 | return fig
23 |
24 |
--------------------------------------------------------------------------------
/src/poseviz/mayavi_util.py:
--------------------------------------------------------------------------------
1 | import numpy as np
2 | import transforms3d
3 | from mayavi import mlab
4 |
5 | import cameravision
6 |
7 | MM_TO_UNIT = 1 / 1000
8 | UNIT_TO_MM = 1000
9 | WORLD_UP = np.array([0, 0, 1])
10 | WORLD_TO_MAYAVI_ROTATION_MAT = np.array([
11 | [1, 0, 0],
12 | [0, 0, 1],
13 | [0, -1, 0]], dtype=np.float32)
14 |
15 | CAM_TO_MAYCAM_ROTATION_MAT = np.array([
16 | [0, 1, 0],
17 | [1, 0, 0],
18 | [0, 0, -1]], dtype=np.float32)
19 |
20 |
21 | def rotation_mat(up):
22 | up = unit_vector(up)
23 | rightlike = np.array([1, 0, 0])
24 | if np.allclose(up, rightlike):
25 | rightlike = np.array([0, 1, 0])
26 |
27 | forward = unit_vector(np.cross(up, rightlike))
28 | right = np.cross(forward, up)
29 | return np.row_stack([right, forward, up])
30 |
31 |
32 | def unit_vector(vectors, axis=-1):
33 | norm = np.linalg.norm(vectors, axis=axis, keepdims=True)
34 | return vectors / norm
35 |
36 |
37 | def world_to_mayavi(points):
38 | points = np.asarray(points)
39 | if points.ndim == 1:
40 | return WORLD_TO_MAYAVI_ROTATION_MAT @ points * MM_TO_UNIT
41 | else:
42 | return points @ WORLD_TO_MAYAVI_ROTATION_MAT.T * MM_TO_UNIT
43 |
44 |
45 | def mayavi_to_world(points):
46 | points = np.asarray(points)
47 | if points.ndim == 1:
48 | return WORLD_TO_MAYAVI_ROTATION_MAT.T @ points * UNIT_TO_MM
49 | else:
50 | return points @ WORLD_TO_MAYAVI_ROTATION_MAT * UNIT_TO_MM
51 |
52 |
53 | def set_world_up(world_up):
54 | global WORLD_UP
55 | WORLD_UP = np.asarray(world_up)
56 | global WORLD_TO_MAYAVI_ROTATION_MAT
57 | WORLD_TO_MAYAVI_ROTATION_MAT = rotation_mat(WORLD_UP)
58 |
59 |
60 | def set_view_to_camera(
61 | camera, pivot=None, image_size=None, view_angle=None, allow_roll=True,
62 | camera_view_padding=0.2):
63 | if pivot is None:
64 | pivot = camera.t + camera.R[2] * 2500
65 |
66 | fig = mlab.gcf()
67 | mayavi_cam = fig.scene.camera
68 | if image_size is not None:
69 | half_height = 0.5 * image_size[1]
70 | offset = np.abs(camera.intrinsic_matrix[1, 2] - half_height)
71 | tan = (1 + camera_view_padding) * (half_height + offset) / camera.intrinsic_matrix[1, 1]
72 | mayavi_cam.view_angle = np.rad2deg(2 * np.arctan(tan))
73 | elif view_angle is not None:
74 | mayavi_cam.view_angle = view_angle
75 |
76 | mayavi_cam.focal_point = world_to_mayavi(pivot)
77 | mayavi_cam.position = world_to_mayavi(camera.t)
78 | if allow_roll:
79 | mayavi_cam.view_up = -WORLD_TO_MAYAVI_ROTATION_MAT @ camera.R[1]
80 | mayavi_cam.compute_view_plane_normal()
81 | fig.scene.renderer.reset_camera_clipping_range()
82 |
83 |
84 | def get_current_view_as_camera():
85 | azimuth, elevation, distance, focalpoint = mlab.view()
86 | azi = -azimuth - 90
87 | elev = -elevation + 90
88 | total_rotation_mat = transforms3d.euler.euler2mat(np.deg2rad(azi), np.deg2rad(elev), 0, 'szxy')
89 | R = WORLD_TO_MAYAVI_ROTATION_MAT.T @ total_rotation_mat @ WORLD_TO_MAYAVI_ROTATION_MAT
90 |
91 | fig = mlab.gcf()
92 | t = mayavi_to_world(fig.scene.camera.position)
93 |
94 | width, height = fig.scene.get_size()
95 | f = height / (np.tan(np.deg2rad(fig.scene.camera.view_angle) / 2) * 2)
96 | intrinsics = np.array(
97 | [[f, 0, (width - 1) / 2],
98 | [0, f, (height - 1) / 2],
99 | [0, 0, 1]], np.float32)
100 | return cameravision.Camera(
101 | intrinsic_matrix=intrinsics, rot_world_to_cam=R, optical_center=t, world_up=WORLD_UP)
102 |
--------------------------------------------------------------------------------
/src/poseviz/messages.py:
--------------------------------------------------------------------------------
1 | from dataclasses import dataclass
2 |
3 |
4 | @dataclass
5 | class UpdateScene:
6 | view_infos: any
7 | viz_camera: any
8 | viz_imshape: any
9 |
10 |
11 | @dataclass
12 | class AppendFrame:
13 | frame: any
14 |
15 |
16 | @dataclass
17 | class StartSequence:
18 | video_path: any =None
19 | fps: any = 30
20 | resolution: any = None
21 | camera_trajectory_path: any = None
22 | audio_source_path: any = None
23 |
24 |
25 | class EndSequence:
26 | pass
27 |
28 |
29 | class Pause:
30 | pass
31 |
32 |
33 | class Resume:
34 | pass
35 |
36 |
37 | class ReinitCameraView:
38 | pass
39 |
40 |
41 | class Quit:
42 | pass
43 |
44 |
45 | class Nothing:
46 | pass
47 |
48 |
49 | @dataclass
50 | class NewRingBuffers:
51 | ringbuffers: any
52 |
--------------------------------------------------------------------------------
/src/poseviz/poseviz.py:
--------------------------------------------------------------------------------
1 | import contextlib
2 | import ctypes
3 | import itertools
4 | import multiprocessing as mp
5 | import multiprocessing.shared_memory
6 | import multiprocessing.sharedctypes
7 | import os
8 | import queue
9 | import signal
10 | import threading
11 | from collections.abc import Sequence
12 | from contextlib import AbstractContextManager
13 | from typing import Optional, TYPE_CHECKING
14 |
15 | import cameravision
16 | import numpy as np
17 | import poseviz.colors
18 | import poseviz.draw2d
19 | import poseviz.view_info
20 | import simplepyutils as spu
21 | import framepump
22 | from poseviz import messages
23 |
24 | if TYPE_CHECKING:
25 | from poseviz import ViewInfo
26 | else:
27 | from poseviz.view_info import ViewInfo
28 |
29 | # We must not globally import mlab, because then it will be imported in the main process
30 | # and the new process will not be able to open a mayavi window.
31 | # This means, we also cannot import various modules like poseviz.components.main_viz
32 | # We import these in the methods of PoseVizMayaviSide.
33 |
34 |
35 | class PoseViz(AbstractContextManager):
36 | def __init__(
37 | self,
38 | joint_names=None,
39 | joint_edges=None,
40 | camera_type="free",
41 | n_views=1,
42 | world_up=(0, -1, 0),
43 | ground_plane_height=-1000,
44 | downscale=None,
45 | downscale_main=None,
46 | viz_fps=100,
47 | queue_size=64,
48 | draw_detections=True,
49 | multicolor_detections=False,
50 | snap_to_cam_on_scene_change=True,
51 | high_quality=True,
52 | draw_2d_pose=False,
53 | show_camera_wireframe=True,
54 | show_field_of_view=True,
55 | resolution=(1280, 720),
56 | use_virtual_display=False,
57 | show_virtual_display=True,
58 | show_ground_plane=True,
59 | paused=False,
60 | camera_view_padding=0.2,
61 | body_model_faces=None,
62 | show_image=True,
63 | out_video_path=None,
64 | out_fps=None,
65 | audio_path=None,
66 | max_pixels_per_frame=1920 * 1080,
67 | ):
68 | """The main class that creates a visualizer process, and provides an interface to update
69 | the visualization state.
70 | This class supports the Python **context manager** protocol to close the visualization at
71 | the end.
72 |
73 | Args:
74 | joint_names (iterable of strings): Names of the human body joints that will be
75 | visualized. Left joints must start with 'l', right joints with 'r', mid-body joints
76 | with something else. Currently only the first character is inspected in each joint
77 | name, to color the left, right and middle part of the stick figure differently.
78 | joint_edges (iterable of int pairs): joint index pairs, describing the
79 | bone-connectivity of the stick figure.
80 | camera_type (string): One of 'original', 'free' or 'bird' (experimental). 'original'
81 | forces the view-camera (i.e. from which we see the visualization) to stay or move
82 | where the displayed camera goes, typically so that we can see the scene through the
83 | recording camera from which the video was recorded. 'free' means that the
84 | view-camera is free to be moved around by the user. 'bird' ( experimental) tries to
85 | automatically move to a nice location to be able to see both the person and the
86 | recording camera from a third-person perspective.
87 | n_views (int): The number of cameras that will be displayed.
88 | world_up (3-vector): A 3-vector, the up vector in the world coordinate system in
89 | which the poses will be specified.
90 | ground_plane_height (float): The vertical position of the ground plane in the world
91 | coordinate system, along the up-vector.
92 | downscale (int): Image downscaling factor for display, to speed up the
93 | visualization.
94 | viz_fps (int): Target frames-per-second of the visualization. If the updates come
95 | faster than this, the visualizer will block and wait, to ensure that
96 | visualization does not
97 | happen faster than this FPS. Of course, if the speed of updates do not deliver this
98 | fps, the visualization will also be slower.
99 | queue_size (int): Size of the internal queue used to communicate with the
100 | visualizer process.
101 | draw_detections (bool): Whether to draw detection boxes on the images.
102 | multicolor_detections (bool): Whether to color each detection box with a different
103 | color.
104 | This is useful when tracking people with consistent person IDs, and has no effect
105 | if `draw_detections` is False.
106 | snap_to_cam_on_scene_change (bool): Whether to reinitialize the view camera to the
107 | original camera on each change of sequence (through a call to
108 | ```viz.reinit_camera_view()```).
109 | high_quality (bool): Whether to use high-resolution spheres and tubes for the
110 | skeletons (set to False for better speed).
111 | draw_2d_pose (bool): Whether to draw the 2D skeleton on the displayed camera image.
112 | show_camera_wireframe (bool): Whether to visualize each camera as a pyramid-like
113 | wireframe object.
114 | show_field_of_view (bool): Whether to visualize an extended pyramid shape indicating the
115 | field of view of the cameras. Recommended to turn off in multi-camera
116 | setups, as otherwise the visualization can get crowded.
117 | resolution ((int, int)): The resolution of the visualization window (width,
118 | height) pair.
119 | use_virtual_display (bool): Whether to use a virtual display for visualization.
120 | There may be two reasons to do this. First, to allow higher-resolution visualization
121 | than the screen resolution. Windows that are larger than the display screen can be
122 | difficult to create under certain GUI managers like GNOME. Second, this can be a
123 | way to do off-screen rendering.
124 | show_virtual_display (bool): Whether to show the virtual display or to hide it (
125 | off-screen rendering). This has no effect if ```use_virtual_display``` is False.
126 | show_ground_plane (bool): Whether to visualize a checkerboard ground plane.
127 | paused (bool): Whether to start the visualization in paused state.
128 | camera_view_padding (float): When viewing the scence from a visualized camera position,
129 | it is often useful to also see beyond the edges of the video frame. The
130 | ```cameera_view_padding```
131 | value adjusts what fraction of the frame size is applied as padding around it.
132 | Example with [```camera_view_padding=0```](/poseviz/images/padding_0.jpg) and [
133 | ```camera_view_padding=0.2```](/poseviz/images/padding_0.2.jpg)
134 | out_video_path: Path where the output video should be generated. (It can also be started
135 | later with ```new_sequence_output```). If None, no video will be generated for now.
136 | out_fps: 'frames per second' setting of the output video to be generated.
137 | audio_path: Path to the audio source file, which may also be a video file whose audio
138 | will be copied over to the output video.
139 | """
140 |
141 | n_threads_undist = 12
142 | queue_size_undist = min(max(16, n_views), 2 * n_views)
143 | queue_size_post = 2
144 | queue_size_waiter = queue_size
145 | self.q_messages_pre = queue.Queue(queue_size_waiter)
146 | if paused:
147 | self.pause()
148 | self.q_messages_post = mp.JoinableQueue(queue_size_post)
149 |
150 | self.undistort_pool = spu.ThrottledPool(
151 | n_threads_undist, use_threads=True, task_buffer_size=queue_size_undist
152 | )
153 |
154 | self.ringbuffer_size = queue_size_undist + queue_size_waiter + queue_size_post + 1
155 | self.raw_arrays = [
156 | multiprocessing.sharedctypes.RawArray(
157 | ctypes.c_uint8, self.ringbuffer_size * max_pixels_per_frame * 3
158 | )
159 | for _ in range(n_views)
160 | ]
161 | self.ring_index = 0
162 |
163 | self.posedata_waiter_thread = threading.Thread(
164 | target=main_posedata_waiter,
165 | args=(self.q_messages_pre, self.q_messages_post),
166 | daemon=True,
167 | )
168 | self.posedata_waiter_thread.start()
169 |
170 | self.downscale_main = downscale_main or downscale or (1 if high_quality else 2)
171 | self.downscale = downscale or 4
172 | self.snap_to_cam_on_scene_change = snap_to_cam_on_scene_change
173 | self.main_cam_value = mp.Value("i", 0)
174 |
175 | self.visualizer_process = mp.Process(
176 | target=_main_visualize,
177 | args=(
178 | self.q_messages_post,
179 | joint_names,
180 | joint_edges,
181 | camera_type,
182 | n_views,
183 | world_up,
184 | ground_plane_height,
185 | viz_fps,
186 | draw_detections,
187 | multicolor_detections,
188 | snap_to_cam_on_scene_change,
189 | high_quality,
190 | draw_2d_pose,
191 | show_camera_wireframe,
192 | show_field_of_view,
193 | resolution,
194 | use_virtual_display,
195 | show_virtual_display,
196 | show_ground_plane,
197 | self.main_cam_value,
198 | camera_view_padding,
199 | body_model_faces,
200 | show_image,
201 | self.raw_arrays,
202 | ),
203 | daemon=True,
204 | )
205 | self.visualizer_process.start()
206 |
207 | if resolution[1] > 1080:
208 | input("Move the window to be partially outside the screen, then press Enter...")
209 | # Spent a lot of time trying to fix this, but it's very tricky and frustrating!
210 | # We have to manually drag the window partially off-screen, else it's impossible to
211 | # set the size
212 | # larger than the display resolution! If you try, it will just snap to the screen
213 | # borders (get maximized).
214 | # An alternative would be like https://unix.stackexchange.com/a/680848/291533, ie
215 | # we can take the window away from the window manager and then freely set its size
216 | # and position.
217 | # The downside is that we are no longer able to move the window with the mouse and it
218 | # is stuck in the foreground.
219 | # So here we rather go the manual route.
220 | # Further info:
221 | # https://www.reddit.com/r/kde/comments/mzza4d
222 | # /programmatically_moving_windows_off_the_screen/
223 | # https://unix.stackexchange.com/questions/517396/is-it-possible-to-move-window-out
224 | # -of-the-screen-border#comment1396167_517396
225 | # https://github.com/jordansissel/xdotool/issues/186#issuecomment-470032261
226 | # Another option is some kind of virtual screen, off-screen rendering etc.
227 | # But then there are issues with hardware acceleration... It's quite complex.
228 | os.system(
229 | "/usr/bin/xdotool search --name ^PoseViz$"
230 | f" windowsize {resolution[0]} {resolution[1]}"
231 | )
232 |
233 | if out_video_path is not None:
234 | self.new_sequence_output(out_video_path, out_fps, audio_path)
235 |
236 | def update(
237 | self,
238 | frame: np.ndarray,
239 | boxes=(),
240 | poses=(),
241 | camera: Optional[cameravision.Camera] = None,
242 | poses_true=(),
243 | poses_alt=(),
244 | vertices=(),
245 | vertices_true=(),
246 | vertices_alt=(),
247 | viz_camera=None,
248 | viz_imshape=None,
249 | block=True,
250 | uncerts=None,
251 | uncerts_alt=None,
252 | ):
253 | """Update the visualization for a new timestep, assuming a single-camera setup.
254 |
255 | Args:
256 | frame (uint8): RGB image frame [H, W, 3], the image to be displayed on the camera.
257 | boxes (seq of np.ndarray): bounding boxes of the persons
258 | poses (seq of np.ndarray): the world
259 | camera (poseviz.Camera):
260 | poses_true (seq of np.ndarray):
261 | poses_alt (seq of np.ndarray):
262 | block (bool): decides what to do if the buffer queue is full because the visualizer
263 | is slower than the update calls. If true, the thread will block (wait). If false,
264 | the current update call is ignored (frame dropping).
265 | """
266 | if len(boxes) == 0:
267 | boxes = np.zeros((0, 4), np.float32)
268 |
269 | if uncerts is not None:
270 | vertices = [
271 | np.concatenate([v, uncert[..., np.newaxis]], axis=-1)
272 | for v, uncert in zip(vertices, uncerts)
273 | ]
274 | if uncerts_alt is not None:
275 | vertices_alt = [
276 | np.concatenate([v, uncert[..., np.newaxis]], axis=-1)
277 | for v, uncert in zip(vertices_alt, uncerts_alt)
278 | ]
279 |
280 | viewinfo = ViewInfo(
281 | frame,
282 | boxes,
283 | poses,
284 | camera,
285 | poses_true,
286 | poses_alt,
287 | vertices,
288 | vertices_true,
289 | vertices_alt,
290 | )
291 | self.update_multiview([viewinfo], viz_camera, viz_imshape, block=block)
292 |
293 | def update_multiview(
294 | self,
295 | view_infos: Sequence[ViewInfo],
296 | viz_camera: Optional[cameravision.Camera] = None,
297 | viz_imshape=None,
298 | block: bool = True,
299 | ):
300 | """Update the visualization for a new timestep, with multi-view data.
301 |
302 | Args:
303 | view_infos: The view information for each view.
304 | viz_camera: The camera to be used for the visualization.
305 | viz_imshape: The shape of the image of the visualization camera.
306 | block: decides what to do if the buffer queue is full because the visualizer
307 | is slower than the update calls. If true, the thread will block (wait). If false,
308 | the current update call is ignored (frame dropping).
309 | """
310 |
311 | for i in range(len(view_infos)):
312 | if len(view_infos[i].boxes) == 0:
313 | view_infos[i].boxes = np.zeros((0, 4), np.float32)
314 |
315 | d = self.downscale_main if i == self.main_cam_value.value else self.downscale
316 | if view_infos[i].frame is not None:
317 | new_imshape = poseviz.draw2d.rounded_int_tuple(
318 | np.array(view_infos[i].frame.shape[:2], np.float32) / d
319 | )
320 | dst = np_from_raw_array(
321 | self.raw_arrays[i],
322 | (new_imshape[0], new_imshape[1], 3),
323 | self.ring_index,
324 | np.uint8,
325 | )
326 | view_infos[i] = self.undistort_pool.apply_async(
327 | poseviz.view_info.downscale_and_undistort_view_info,
328 | (view_infos[i], dst, self.ring_index),
329 | )
330 | else:
331 | view_infos[i] = self.undistort_pool.apply_async(lambda x: x, (view_infos[i],))
332 | self.ring_index = (self.ring_index + 1) % self.ringbuffer_size
333 |
334 | try:
335 | self.q_messages_pre.put(
336 | messages.UpdateScene(
337 | view_infos=view_infos, viz_camera=viz_camera, viz_imshape=viz_imshape
338 | ),
339 | block=block,
340 | )
341 | except queue.Full:
342 | pass
343 |
344 | def reinit_camera_view(self):
345 | """Waits until the current sequence has finished visualizing, then notifies PoseViz that
346 | a new sequence is now being visualized, so the view camera may need to be
347 | reinitialized.
348 | """
349 | self.q_messages_pre.put(messages.ReinitCameraView())
350 |
351 | def pause(self):
352 | self.q_messages_pre.put(messages.Pause())
353 |
354 | def resume(self):
355 | self.q_messages_pre.put(messages.Resume())
356 |
357 | def new_sequence_output(
358 | self,
359 | new_video_path: Optional[str] = None,
360 | fps: Optional[int] = None,
361 | new_camera_trajectory_path: Optional[str] = None,
362 | audio_source_path: Optional[str] = None,
363 | ):
364 | """Waits until the buffered information is visualized and then starts a new output video
365 | for the upcoming timesteps.
366 |
367 | Args:
368 | new_video_path: Path where the new video should be generated.
369 | fps: 'frames per second' setting of the new video to be generated
370 | new_camera_trajectory_path: Path where the camera trajectory should be saved.
371 | """
372 | if new_video_path is not None or new_camera_trajectory_path is not None:
373 | self.q_messages_pre.put(
374 | messages.StartSequence(
375 | video_path=new_video_path,
376 | fps=fps,
377 | camera_trajectory_path=new_camera_trajectory_path,
378 | audio_source_path=audio_source_path,
379 | )
380 | )
381 |
382 | def finalize_sequence_output(self):
383 | """Notifies PoseViz that the current sequence is finished and the output video should be
384 | finalized."""
385 | self.q_messages_pre.put(messages.EndSequence())
386 |
387 | def close(self):
388 | """Closes the visualization process and window."""
389 | self.undistort_pool.finish()
390 | self.q_messages_pre.put(messages.Quit())
391 | self._join()
392 | self.posedata_waiter_thread.join()
393 | self.visualizer_process.join()
394 |
395 | def _join(self):
396 | self.undistort_pool.join()
397 | self.q_messages_pre.join()
398 | self.q_messages_post.join()
399 |
400 | def __exit__(self, exc_type, exc_val, exc_tb):
401 | """Context manager protocol."""
402 | self.close()
403 |
404 |
405 | def np_from_raw_array(raw_array, elem_shape, index, dtype):
406 | return np.frombuffer(
407 | buffer=raw_array,
408 | dtype=dtype,
409 | count=np.prod(elem_shape),
410 | offset=index * np.prod(elem_shape) * np.dtype(dtype).itemsize,
411 | ).reshape(elem_shape)
412 |
413 |
414 | class PoseVizMayaviSide:
415 | def __init__(
416 | self,
417 | q_messages,
418 | joint_names,
419 | joint_edges,
420 | camera_type,
421 | n_views,
422 | world_up,
423 | ground_plane_height,
424 | fps,
425 | draw_detections,
426 | multicolor_detections,
427 | snap_to_cam_on_scene_change,
428 | high_quality,
429 | draw_2d_pose,
430 | show_camera_wireframe,
431 | show_field_of_view,
432 | resolution,
433 | use_virtual_display,
434 | show_virtual_display,
435 | show_ground_plane,
436 | main_cam_value,
437 | camera_view_padding,
438 | body_model_faces,
439 | show_image,
440 | raw_arrays,
441 | ):
442 | self.q_messages = q_messages
443 | self.video_writer = framepump.VideoWriter()
444 |
445 | self.camera_type = camera_type
446 | self.joint_info = (joint_names, joint_edges)
447 | self.initialized_camera = False
448 | self.n_views = n_views
449 | self.world_up = world_up
450 | self.draw_detections = draw_detections
451 | self.multicolor_detections = multicolor_detections
452 | self.main_cam = 0
453 | self.pose_displayed_cam_id = None
454 | self.fps = fps
455 | self.snap_to_cam_on_scene_change = snap_to_cam_on_scene_change
456 | self.ground_plane_height = ground_plane_height
457 | self.show_ground_plane = show_ground_plane
458 | self.step_one_by_one = False
459 | self.paused = False
460 | self.current_viewinfos = None
461 | self.high_quality = high_quality
462 | self.draw_2d_pose = draw_2d_pose
463 | self.show_camera_wireframe = show_camera_wireframe
464 | self.show_field_of_view = show_field_of_view
465 | self.resolution = resolution
466 | self.use_virtual_display = use_virtual_display
467 | self.show_virtual_display = show_virtual_display
468 | self.main_cam_value = main_cam_value
469 | self.camera_view_padding = camera_view_padding
470 | self.body_model_faces = body_model_faces
471 | self.i_pred_frame = 0
472 | self.camera_trajectory = []
473 | self.camera_trajectory_path = None
474 | self.show_image = show_image
475 | self.raw_arrays = raw_arrays
476 |
477 | def run_loop(self):
478 | if self.use_virtual_display:
479 | import pyvirtualdisplay
480 |
481 | display = pyvirtualdisplay.Display(
482 | visible=self.show_virtual_display, size=self.resolution
483 | )
484 | else:
485 | display = contextlib.nullcontext()
486 |
487 | with display:
488 | # Imports are here so Mayavi is loaded in the visualizer process,
489 | # and potentially under the virtual display
490 | from mayavi import mlab
491 | import poseviz.init
492 | import poseviz.components.main_viz
493 | import poseviz.components.ground_plane_viz
494 | import poseviz.mayavi_util
495 |
496 | poseviz.mayavi_util.set_world_up(self.world_up)
497 | fig = poseviz.init.initialize(self.resolution)
498 | fig.scene.interactor.add_observer("KeyPressEvent", self._on_keypress)
499 | if self.show_ground_plane:
500 | poseviz.components.ground_plane_viz.draw_checkerboard(
501 | ground_plane_height=self.ground_plane_height
502 | )
503 | self.view_visualizers = [
504 | poseviz.components.main_viz.MainViz(
505 | self.joint_info,
506 | self.joint_info,
507 | self.joint_info,
508 | self.camera_type,
509 | show_image=self.show_image,
510 | high_quality=self.high_quality,
511 | show_field_of_view=self.show_field_of_view,
512 | show_camera_wireframe=self.show_camera_wireframe,
513 | body_model_faces=self.body_model_faces,
514 | )
515 | for _ in range(self.n_views)
516 | ]
517 | delay = max(10, int(round(1000 / self.fps)))
518 |
519 | # Need to store it in a variable to make sure it's not
520 | # destructed by the garbage collector!
521 | _ = mlab.animate(delay=delay, ui=False)(self.animate)(fig)
522 | mlab.show()
523 |
524 | def animate(self, fig):
525 | # The main animation loop of the visualizer
526 | # The infinite loop pops an incoming command or frame information from the queue
527 | # and processes it. If there's no incoming information, we just render the scene and
528 | # continue to iterate.
529 | from mayavi import mlab
530 |
531 | while True:
532 | # Get a new message (command) from the multiprocessing queue
533 | msg = self.receive_message()
534 | if isinstance(msg, messages.Nothing):
535 | if self.current_viewinfos is not None:
536 | self.update_view_camera()
537 | fig.scene.render()
538 | if self.paused:
539 | self.capture_frame()
540 | yield
541 | elif isinstance(msg, messages.UpdateScene):
542 | self.update_visu(msg.view_infos)
543 | self.update_view_camera(msg.viz_camera, msg.viz_imshape)
544 | self.capture_frame()
545 |
546 | if self.step_one_by_one:
547 | self.step_one_by_one = False
548 | self.paused = True
549 |
550 | self.i_pred_frame += 1
551 | self.q_messages.task_done()
552 | yield
553 | elif isinstance(msg, messages.ReinitCameraView):
554 | if self.snap_to_cam_on_scene_change:
555 | self.initialized_camera = False
556 | self.current_viewinfos = None
557 | self.q_messages.task_done()
558 | elif isinstance(msg, messages.StartSequence):
559 | if msg.video_path is not None:
560 | self.video_writer.start_sequence(
561 | video_path=msg.video_path,
562 | fps=msg.fps,
563 | audio_source_path=msg.audio_source_path,
564 | )
565 | self.camera_trajectory_path = msg.camera_trajectory_path
566 | self.i_pred_frame = 0
567 |
568 | # os.system('/usr/bin/xdotool search --name ^PoseViz$'
569 | # f' windowsize --sync {msg.resolution[0]} {msg.resolution[1]}')
570 | self.q_messages.task_done()
571 | # yield
572 | elif isinstance(msg, messages.Pause):
573 | self.paused = True
574 | self.q_messages.task_done()
575 | elif isinstance(msg, messages.Resume):
576 | self.paused = False
577 | self.q_messages.task_done()
578 | elif isinstance(msg, messages.EndSequence):
579 | if self.camera_trajectory_path is not None:
580 | spu.dump_pickle(self.camera_trajectory, self.camera_trajectory_path)
581 | self.video_writer.end_sequence()
582 | self.q_messages.task_done()
583 | elif isinstance(msg, messages.Quit):
584 | if self.camera_trajectory_path is not None:
585 | spu.dump_pickle(self.camera_trajectory, self.camera_trajectory_path)
586 | mlab.close(all=True)
587 | self.video_writer.close()
588 | self.q_messages.task_done()
589 | return
590 | else:
591 | raise ValueError("Unknown message:", msg)
592 |
593 | def receive_message(self):
594 | """Gets a new command or frame visualization data from the multiprocessing queue.
595 | Returns messages.Nothing if the queue is empty or the visualization is paused."""
596 | if self.paused:
597 | return messages.Nothing()
598 |
599 | try:
600 | return self.q_messages.get_nowait()
601 | except queue.Empty:
602 | return messages.Nothing()
603 |
604 | def update_visu(self, view_infos):
605 | self.update_num_views(len(view_infos)) # Update the number of views if necessary
606 |
607 | for i, v in enumerate(view_infos):
608 | if v.frame is not None:
609 | shape, dtype, index = v.frame
610 | v.frame = np_from_raw_array(self.raw_arrays[i], shape, index, dtype)
611 |
612 | pose_displayed_view_infos = (
613 | view_infos
614 | if self.pose_displayed_cam_id is None
615 | else [view_infos[self.pose_displayed_cam_id]]
616 | )
617 |
618 | all_poses = [p for v in pose_displayed_view_infos for p in v.poses]
619 | all_poses_true = [p for v in pose_displayed_view_infos for p in v.poses_true]
620 | all_poses_alt = [p for v in pose_displayed_view_infos for p in v.poses_alt]
621 | all_vertices = [p for v in pose_displayed_view_infos for p in v.vertices]
622 | all_vertices_true = [p for v in pose_displayed_view_infos for p in v.vertices_true]
623 | all_vertices_alt = [p for v in pose_displayed_view_infos for p in v.vertices_alt]
624 |
625 | self.current_viewinfos = view_infos
626 |
627 | if self.multicolor_detections:
628 | box_colors = poseviz.colors.cycle_over_colors(False)
629 | else:
630 | box_colors = itertools.repeat((31, 119, 180))
631 | joint_names, joint_edges = self.joint_info
632 | for i_viz, (view_info, viz) in enumerate(zip(view_infos, self.view_visualizers)):
633 | poses = all_poses if viz is self.view_visualizers[0] else None
634 | poses_true = all_poses_true if viz is self.view_visualizers[0] else None
635 | poses_alt = all_poses_alt if viz is self.view_visualizers[0] else None
636 | vertices = all_vertices if viz is self.view_visualizers[0] else None
637 | vertices_true = all_vertices_true if viz is self.view_visualizers[0] else None
638 | vertices_alt = all_vertices_alt if viz is self.view_visualizers[0] else None
639 |
640 | max_size = np.max(view_info.frame.shape[:2])
641 | if max_size < 512:
642 | thickness = 1
643 | elif max_size < 1024:
644 | thickness = 2
645 | else:
646 | thickness = 3
647 |
648 | if self.draw_detections:
649 | for color, box in zip(box_colors, view_info.boxes):
650 | poseviz.draw2d.draw_box(view_info.frame, box, color, thickness=thickness)
651 |
652 | if self.draw_2d_pose:
653 | pose_groups = [view_info.poses, view_info.poses_true, view_info.poses_alt]
654 | colors = [poseviz.colors.green, poseviz.colors.red, poseviz.colors.orange]
655 | for poses, color in zip(pose_groups, colors):
656 | for pose in poses:
657 | pose2d = view_info.camera.world_to_image(pose)
658 | poseviz.draw2d.draw_stick_figure_2d_inplace(
659 | view_info.frame, pose2d, joint_edges, thickness, color=color
660 | )
661 |
662 | viz.update(
663 | view_info.camera,
664 | view_info.frame,
665 | poses,
666 | poses_true,
667 | poses_alt,
668 | vertices,
669 | vertices_true,
670 | vertices_alt,
671 | highlight=self.pose_displayed_cam_id == i_viz,
672 | )
673 |
674 | def update_num_views(self, new_n_views):
675 | if new_n_views > self.n_views:
676 | self.view_visualizers += [
677 | poseviz.components.main_viz.MainViz(
678 | self.joint_info,
679 | self.joint_info,
680 | self.joint_info,
681 | self.camera_type,
682 | show_image=True,
683 | high_quality=self.high_quality,
684 | show_field_of_view=self.show_field_of_view,
685 | show_camera_wireframe=self.show_camera_wireframe,
686 | body_model_faces=self.body_model_faces,
687 | )
688 | for _ in range(new_n_views - self.n_views)
689 | ]
690 | self.n_views = new_n_views
691 | elif new_n_views < self.n_views:
692 | for viz in self.view_visualizers[new_n_views:]:
693 | viz.remove()
694 | del self.view_visualizers[new_n_views:]
695 | self.n_views = new_n_views
696 |
697 | def update_view_camera(self, camera=None, imshape=None):
698 | import poseviz.mayavi_util
699 |
700 | main_view_info = self.current_viewinfos[self.main_cam]
701 |
702 | if camera is not None:
703 | if imshape is None:
704 | imshape = main_view_info.frame.shape
705 | poseviz.mayavi_util.set_view_to_camera(
706 | camera, image_size=(imshape[1], imshape[0]), allow_roll=True, camera_view_padding=0
707 | )
708 | elif self.camera_type == "original" or (
709 | self.camera_type == "free" and not self.initialized_camera
710 | ):
711 | poseviz.mayavi_util.set_view_to_camera(
712 | main_view_info.camera,
713 | image_size=(main_view_info.frame.shape[1], main_view_info.frame.shape[0]),
714 | allow_roll=False,
715 | camera_view_padding=self.camera_view_padding,
716 | )
717 | elif self.camera_type == "bird" and not self.initialized_camera:
718 | pivot = main_view_info.camera.t + main_view_info.camera.R[2] * 2000
719 | camera_view = main_view_info.camera.copy()
720 | camera_view.t = (camera_view.t - pivot) * 1.35 + pivot
721 | camera_view.orbit_around(pivot, np.deg2rad(-25), "vertical")
722 | camera_view.orbit_around(pivot, np.deg2rad(-15), "horizontal")
723 | poseviz.mayavi_util.set_view_to_camera(
724 | camera_view,
725 | pivot=pivot,
726 | view_angle=55,
727 | allow_roll=False,
728 | camera_view_padding=self.camera_view_padding,
729 | )
730 |
731 | self.initialized_camera = True
732 |
733 | def capture_frame(self):
734 | from mayavi import mlab
735 |
736 | if self.camera_trajectory_path is not None:
737 | from poseviz.mayavi_util import get_current_view_as_camera
738 |
739 | viz_cam = get_current_view_as_camera()
740 | self.camera_trajectory.append((self.i_pred_frame, viz_cam))
741 |
742 | if self.video_writer.is_active():
743 | out_frame = mlab.screenshot(antialiased=False)
744 | # fig = mlab.gcf()
745 | # fig.scene.disable_render = True
746 | # fig.scene.off_screen_rendering = True
747 | # mlab.savefig(self.temp_image_path, size=self.resolution)
748 | # fig.scene.off_screen_rendering = False
749 | # fig.scene.disable_render = False
750 | # out_frame = imageio.imread(self.temp_image_path)
751 | # factor = 1 / 3
752 | # if self.camera_type != 'original' and False:
753 | # image = main_view_info.frame
754 | # image = np.clip(image.astype(np.float32) + 50, 0, 255).astype(np.uint8)
755 | # for pred in main_view_info.poses:
756 | # image = poseviz.draw2d.draw_stick_figure_2d(
757 | # image, main_view_info.camera.world_to_image(pred), self.joint_info, 3)
758 | # out_frame[:illust.shape[0], :illust.shape[1]] = illust
759 | self.video_writer.append_data(out_frame)
760 |
761 | def _on_keypress(self, obj, ev):
762 | key = obj.GetKeySym()
763 | if key == "x":
764 | # Pause
765 | self.paused = not self.paused
766 | elif key == "t":
767 | # Toggle view poses predicted from all views or just one
768 | if self.pose_displayed_cam_id is None:
769 | self.pose_displayed_cam_id = self.main_cam
770 | else:
771 | self.pose_displayed_cam_id = None
772 | elif key == "c":
773 | # Frame by frame stepping on key press
774 | self.step_one_by_one = True
775 | self.paused = False
776 | elif key == "n":
777 | # Cycle the view camera
778 | self.main_cam = (self.main_cam + 1) % self.n_views
779 | self.main_cam_value.value = self.main_cam
780 | self.initialized_camera = False
781 | elif key == "m":
782 | # Cycle which camera's predicted pose is displayed
783 | if self.pose_displayed_cam_id is None:
784 | self.pose_displayed_cam_id = self.main_cam
785 | else:
786 | self.pose_displayed_cam_id = (self.pose_displayed_cam_id + 1) % self.n_views
787 | elif key == "o":
788 | # Cycle both the view camera and which camera's predicted pose is displayed
789 | self.main_cam = (self.main_cam + 1) % self.n_views
790 | self.main_cam_value.value = self.main_cam
791 | self.pose_displayed_cam_id = self.main_cam
792 | self.initialized_camera = False
793 | elif key in ("d", "g"):
794 | # Snap to nearest camera
795 | import poseviz.mayavi_util
796 |
797 | display_cameras = [vi.camera for vi in self.current_viewinfos]
798 | viewing_camera = poseviz.mayavi_util.get_current_view_as_camera()
799 | self.main_cam = np.argmin(
800 | [np.linalg.norm(viewing_camera.t - c.t) for c in display_cameras]
801 | )
802 | self.main_cam_value.value = self.main_cam
803 | self.initialized_camera = False
804 | if key == "g":
805 | # Also display the predicted pose from this nearest camera
806 | self.pose_displayed_cam_id = self.main_cam
807 | elif key == "u":
808 | # Show just the main cam pred
809 | self.pose_displayed_cam_id = self.main_cam
810 | elif key == "z":
811 | # Show just the main cam pred
812 | self.camera_type = "original" if self.camera_type != "original" else "free"
813 | self.initialized_camera = False
814 | else:
815 | try:
816 | self.main_cam = max(0, min(int(key) - 1, self.n_views - 1))
817 | self.main_cam_value.value = self.main_cam
818 | self.initialized_camera = False
819 | except ValueError:
820 | pass
821 |
822 |
823 | def _main_visualize(*args, **kwargs):
824 | _terminate_on_parent_death()
825 | remove_qt_paths_added_by_opencv()
826 | monkey_patched_numpy_types()
827 | viz = PoseVizMayaviSide(*args, **kwargs)
828 | viz.run_loop()
829 |
830 |
831 | def _terminate_on_parent_death():
832 | # Make sure the visualizer process is terminated when the parent process dies, for whatever reason
833 | prctl = ctypes.CDLL("libc.so.6").prctl
834 | PR_SET_PDEATHSIG = 1
835 | prctl(PR_SET_PDEATHSIG, signal.SIGTERM)
836 |
837 | # ignore sigterm and sigint
838 | # without this, a single CTRL+C leaves the process hanging
839 | #signal.signal(signal.SIGTERM, signal.SIG_IGN)
840 | signal.signal(signal.SIGINT, signal.SIG_IGN)
841 |
842 |
843 | def monkey_patched_numpy_types():
844 | """VTK tries to import np.bool etc which are not available anymore."""
845 | for name in ["bool", "int", "float", "complex", "object", "str"]:
846 | if name not in dir(np):
847 | setattr(np, name, getattr(np, name + "_"))
848 |
849 |
850 | def remove_qt_paths_added_by_opencv():
851 | """Remove Qt paths added by OpenCV, which may cause conflicts with Qt as used by Mayavi."""
852 | # See also https://forum.qt.io/post/654289
853 | import sys
854 |
855 | # noinspection PyUnresolvedReferences
856 | import cv2
857 |
858 | try:
859 | from cv2.version import ci_build, headless
860 |
861 | ci_and_not_headless = ci_build and not headless
862 | except:
863 | ci_and_not_headless = False
864 |
865 | if sys.platform.startswith("linux") and ci_and_not_headless:
866 | os.environ.pop("QT_QPA_PLATFORM_PLUGIN_PATH")
867 | os.environ.pop("QT_QPA_FONTDIR")
868 |
869 |
870 | def main_posedata_waiter(q_posedata_pre, q_posedata_post):
871 | while True:
872 | msg = q_posedata_pre.get()
873 | if isinstance(msg, messages.UpdateScene):
874 | msg.view_infos = [v.get() for v in msg.view_infos]
875 |
876 | q_posedata_post.put(msg)
877 | q_posedata_pre.task_done()
878 | if isinstance(msg, messages.Quit):
879 | return
880 |
--------------------------------------------------------------------------------
/src/poseviz/video_writing.py:
--------------------------------------------------------------------------------
1 | # import os
2 | # import os.path as osp
3 | # import queue
4 | # import threading
5 | #
6 | # import imageio.v2 as imageio
7 | # from poseviz import messages
8 | #
9 | #
10 | # def main_video_writer(q):
11 | # writer = None
12 | # while True:
13 | # msg = q.get()
14 | #
15 | # if isinstance(msg, messages.AppendFrame) and writer is not None:
16 | # writer.append_data(msg.frame)
17 | # elif isinstance(msg, messages.StartSequence):
18 | # if writer is not None:
19 | # writer.close()
20 | #
21 | # os.makedirs(osp.dirname(msg.video_path), exist_ok=True)
22 | # writer = imageio.get_writer(
23 | # msg.video_path,
24 | # codec="h264",
25 | # input_params=["-r", str(msg.fps), '-thread_queue_size', '4096'],
26 | # output_params=["-crf", "15"],
27 | # audio_path=msg.audio_source_path,
28 | # audio_codec='copy' if msg.audio_source_path is not None else None,
29 | # macro_block_size=None,
30 | # )
31 | # elif isinstance(msg, (messages.EndSequence, messages.Quit)):
32 | # if writer is not None:
33 | # writer.close()
34 | # writer = None
35 | #
36 | # q.task_done()
37 | #
38 | # if isinstance(msg, messages.Quit):
39 | # return
40 | #
41 | #
42 | # class VideoWriter:
43 | # def __init__(self, queue_size=32):
44 | # self.q = queue.Queue(queue_size)
45 | # self.thread = threading.Thread(target=main_video_writer, args=(self.q,), daemon=True)
46 | # self.thread.start()
47 | # self.active = False
48 | #
49 | # def start_sequence(self, msg: messages.StartSequence):
50 | # self.q.put(msg)
51 | # self.active = True
52 | #
53 | # def is_active(self):
54 | # return self.active
55 | #
56 | # def append_data(self, frame):
57 | # self.q.put(messages.AppendFrame(frame))
58 | #
59 | # def end_sequence(self):
60 | # self.q.put(messages.EndSequence())
61 | # self.active = False
62 | #
63 | # def finish(self):
64 | # self.q.put(messages.Quit())
65 | # self.q.join()
66 | # self.thread.join()
67 | # self.active = False
68 |
--------------------------------------------------------------------------------
/src/poseviz/view_info.py:
--------------------------------------------------------------------------------
1 | import dataclasses
2 |
3 | import cameravision
4 | import numpy as np
5 | import poseviz.draw2d
6 |
7 |
8 | @dataclasses.dataclass
9 | class ViewInfo:
10 | """A container for storing information about a single view of a scene"""
11 |
12 | frame: any = None
13 | boxes: any = ()
14 | poses: any = ()
15 | camera: cameravision.Camera = None
16 | poses_true: any = ()
17 | poses_alt: any = ()
18 | vertices: any = ()
19 | vertices_true: any = ()
20 | vertices_alt: any = ()
21 |
22 |
23 | def downscale_and_undistort_view_info(view_info, dst, index):
24 | scale_factor = np.array(dst.shape[:2], np.float32) / np.array(
25 | view_info.frame.shape[:2], np.float32
26 | )
27 |
28 | if view_info.frame.dtype == np.uint16:
29 | view_info.frame = np.ascontiguousarray(view_info.frame.view(np.uint8)[..., ::2])
30 |
31 | if np.any(scale_factor != 1):
32 | resize_dst = np.empty_like(dst) if view_info.camera.has_distortion() else dst
33 |
34 | view_info.frame = poseviz.draw2d.resize(view_info.frame, dst=resize_dst)
35 | if view_info.boxes is not None:
36 | view_info.boxes = np.asarray(view_info.boxes)
37 | view_info.boxes[:, :2] *= scale_factor
38 | view_info.boxes[:, 2:4] *= scale_factor
39 | view_info.camera = view_info.camera.scale_output(scale_factor, inplace=False)
40 |
41 | if view_info.camera.has_distortion():
42 | old_camera = view_info.camera
43 | view_info.camera = old_camera.undistort(
44 | alpha_balance=0.8, imshape=dst.shape, inplace=False
45 | )
46 |
47 | cameravision.reproject_image(
48 | view_info.frame,
49 | old_camera,
50 | view_info.camera,
51 | dst.shape,
52 | precomp_undist_maps=False,
53 | cache_maps=True,
54 | dst=dst,
55 | )
56 |
57 | if view_info.boxes is not None:
58 | for i in range(len(view_info.boxes)):
59 | view_info.boxes[i, :4] = cameravision.reproject_box(
60 | view_info.boxes[i, :4], old_camera, view_info.camera
61 | )
62 | elif view_info.frame is not dst:
63 | np.copyto(dst, view_info.frame)
64 |
65 | view_info.frame = (dst.shape, dst.dtype, index)
66 | return view_info
67 |
--------------------------------------------------------------------------------