├── .gitattributes ├── .gitignore ├── LICENSE ├── README.md ├── gen_unified_interface.py ├── pyproject.toml ├── test ├── io_ │ └── write_ply.py ├── numpy_ │ ├── mesh │ │ ├── compute_face_angle.py │ │ ├── compute_face_normal.py │ │ ├── compute_vertex_normal.py │ │ ├── compute_vertex_normal_weighted.py │ │ ├── merge_duplicate_vertices.py │ │ ├── remove_corrupted_faces.py │ │ └── triangulate.py │ ├── rasterization │ │ └── warp_image_by_depth.py │ ├── transforms │ │ ├── crop_intrinsic.py │ │ ├── extrinsic_look_at.py │ │ ├── extrinsic_to_view.py │ │ ├── intrinsic.py │ │ ├── intrinsic_from_fov.py │ │ ├── intrinsic_from_fov_xy.py │ │ ├── intrinsic_to_perspective.py │ │ ├── linearize_depth.py │ │ ├── normalize_intrinsic.py │ │ ├── perspective.py │ │ ├── perspective_from_fov.py │ │ ├── perspective_from_fov_xy.py │ │ ├── perspective_to_intrinsic.py │ │ ├── pixel_to_ndc.py │ │ ├── pixel_to_uv.py │ │ ├── project_cv.py │ │ ├── project_depth.py │ │ ├── project_gl.py │ │ ├── project_gl_cv.py │ │ ├── unproject_cv.py │ │ ├── unproject_gl.py │ │ ├── view_look_at.py │ │ └── view_to_extrinsic.py │ └── utils │ │ └── image_mesh.py ├── rasterization_ │ └── gl │ │ ├── basic.py │ │ └── rasterize_uv.py ├── test.py └── torch_ │ ├── mesh │ ├── compute_face_angle.py │ ├── compute_face_normal.py │ ├── compute_vertex_normal.py │ ├── compute_vertex_normal_weighted.py │ ├── merge_duplicate_vertices.py │ ├── remove_corrupted_faces.py │ └── triangulate.py │ ├── rasterization │ └── warp_image_by_depth.py │ ├── transforms │ ├── crop_intrinsic.py │ ├── extrinsic_look_at.py │ ├── extrinsic_to_view.py │ ├── intrinsic.py │ ├── intrinsic_from_fov.py │ ├── intrinsic_from_fov_xy.py │ ├── intrinsic_to_perspective.py │ ├── linearize_depth.py │ ├── matrix_to_quaternion.py │ ├── normalize_intrinsic.py │ ├── perspective.py │ ├── perspective_from_fov.py │ ├── perspective_from_fov_xy.py │ ├── perspective_to_intrinsic.py │ ├── pixel_to_ndc.py │ ├── pixel_to_uv.py │ ├── project_cv.py │ ├── project_depth.py │ ├── project_gl.py │ ├── project_gl_cv.py │ ├── quaternion_to_matrix.py │ ├── slerp.py │ ├── unproject_cv.py │ ├── unproject_gl.py │ ├── view_look_at.py │ └── view_to_extrinsic.py │ └── utils │ └── image_mesh.py └── utils3d ├── __init__.py ├── _helpers.py ├── _unified ├── __init__.py └── __init__.pyi ├── io ├── __init__.py ├── colmap.py ├── obj.py └── ply.py ├── numpy ├── __init__.py ├── _helpers.py ├── mesh.py ├── quadmesh.py ├── rasterization.py ├── shaders │ ├── texture.fsh │ ├── texture.vsh │ ├── vertex_attribute.fsh │ └── vertex_attribute.vsh ├── spline.py ├── transforms.py └── utils.py └── torch ├── __init__.py ├── _helpers.py ├── mesh.py ├── nerf.py ├── rasterization.py ├── transforms.py └── utils.py /.gitattributes: -------------------------------------------------------------------------------- 1 | # Auto detect text files and perform LF normalization 2 | * text=auto 3 | -------------------------------------------------------------------------------- /.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 | *.egg-info/ 24 | .installed.cfg 25 | *.egg 26 | MANIFEST 27 | 28 | # PyInstaller 29 | # Usually these files are written by a python script from a template 30 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 31 | *.manifest 32 | *.spec 33 | 34 | # Installer logs 35 | pip-log.txt 36 | pip-delete-this-directory.txt 37 | 38 | # Unit test / coverage reports 39 | htmlcov/ 40 | .tox/ 41 | .nox/ 42 | .coverage 43 | .coverage.* 44 | .cache 45 | nosetests.xml 46 | coverage.xml 47 | *.cover 48 | .hypothesis/ 49 | .pytest_cache/ 50 | 51 | # Translations 52 | *.mo 53 | *.pot 54 | 55 | # Django stuff: 56 | *.log 57 | local_settings.py 58 | db.sqlite3 59 | 60 | # Flask stuff: 61 | instance/ 62 | .webassets-cache 63 | 64 | # Scrapy stuff: 65 | .scrapy 66 | 67 | # Sphinx documentation 68 | docs/_build/ 69 | 70 | # PyBuilder 71 | target/ 72 | 73 | # Jupyter Notebook 74 | .ipynb_checkpoints 75 | 76 | # IPython 77 | profile_default/ 78 | ipython_config.py 79 | 80 | # pyenv 81 | .python-version 82 | 83 | # celery beat schedule file 84 | celerybeat-schedule 85 | 86 | # SageMath parsed files 87 | *.sage.py 88 | 89 | # Environments 90 | .env 91 | .venv 92 | env/ 93 | venv/ 94 | ENV/ 95 | env.bak/ 96 | venv.bak/ 97 | 98 | # Spyder project settings 99 | .spyderproject 100 | .spyproject 101 | 102 | # Rope project settings 103 | .ropeproject 104 | 105 | # mkdocs documentation 106 | /site 107 | 108 | # mypy 109 | .mypy_cache/ 110 | .dmypy.json 111 | dmypy.json 112 | 113 | # Pyre type checker 114 | .pyre/ 115 | 116 | .vscode 117 | test/results_to_check 118 | timetest.py 119 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2022 EasternJournalist 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 | # utils3d 2 | Easy 3D python utilities for computer vision and graphics researchers. 3 | 4 | Supports: 5 | * Transformation between OpenCV and OpenGL coordinate systems, **no more confusion** 6 | * Easy rasterization, **no worries about OpenGL objects and buffers** 7 | * Some mesh processing utilities, **all vectorized for effciency; some differentiable** 8 | * Projection, unprojection, depth-based image warping, flow-based image warping... 9 | * Easy Reading and writing .obj, .ply files 10 | * Reading and writing Colmap format camera parameters 11 | * NeRF/MipNeRF utilities 12 | 13 | For most functions, there are both numpy (indifferentiable) and pytorch implementations (differentiable). 14 | 15 | Pytorch is not required for using this package, but if you want to use the differentiable functions, you will need to install pytorch (and nvdiffrast if you want to use the pytorch rasterization functions). 16 | 17 | ## Install 18 | 19 | Install by git 20 | 21 | ```bash 22 | pip install git+https://github.com/EasternJournalist/utils3d.git#egg=utils3d 23 | ``` 24 | 25 | or clone the repo and install with `-e` option for convenient updating and modifying. 26 | 27 | ```bash 28 | git clone https://github.com/EasternJournalist/utils3d.git 29 | pip install -e ./utils3d 30 | ``` 31 | 32 | ## Topics (TODO) 33 | 34 | ### Camera 35 | 36 | ### Rotations 37 | 38 | ### Mesh 39 | 40 | ### Rendering 41 | 42 | ### Projection 43 | 44 | ### Image warping 45 | 46 | ### NeRF -------------------------------------------------------------------------------- /gen_unified_interface.py: -------------------------------------------------------------------------------- 1 | import inspect 2 | import textwrap 3 | import re 4 | import itertools 5 | import numbers 6 | import importlib 7 | import sys 8 | import functools 9 | from pathlib import Path 10 | 11 | from utils3d._helpers import suppress_traceback 12 | 13 | 14 | def _contains_tensor(obj): 15 | if isinstance(obj, (list, tuple)): 16 | return any(_contains_tensor(item) for item in obj) 17 | elif isinstance(obj, dict): 18 | return any(_contains_tensor(value) for value in obj.values()) 19 | else: 20 | import torch 21 | return isinstance(obj, torch.Tensor) 22 | 23 | @suppress_traceback 24 | def _call_based_on_args(fname, args, kwargs): 25 | if 'torch' in sys.modules: 26 | if any(_contains_tensor(arg) for arg in args) or any(_contains_tensor(v) for v in kwargs.values()): 27 | fn = getattr(utils3d.torch, fname, None) 28 | if fn is None: 29 | raise NotImplementedError(f"Function {fname} has no torch implementation.") 30 | return fn(*args, **kwargs) 31 | fn = getattr(utils3d.numpy, fname, None) 32 | if fn is None: 33 | raise NotImplementedError(f"Function {fname} has no numpy implementation.") 34 | return fn(*args, **kwargs) 35 | 36 | 37 | def extract_signature(fn): 38 | signature = inspect.signature(fn) 39 | 40 | signature_str = str(signature) 41 | 42 | signature_str = re.sub(r"", lambda m: m.group(0).split('\'')[1], signature_str) 43 | signature_str = re.sub(r"(?=61.0", "wheel"] 3 | build-backend = "setuptools.build_meta" 4 | 5 | [project] 6 | name = "utils3d" 7 | version = "0.0.2" 8 | description = "A small package for 3D graphics" 9 | readme = "README.md" 10 | authors = [ 11 | {name = "EasternJournalist", email = "wangrc2081cs@mail.ustc.edu.cn"} 12 | ] 13 | license = {text = "MIT"} 14 | classifiers = [ 15 | "Programming Language :: Python :: 3", 16 | "License :: OSI Approved :: MIT License", 17 | "Operating System :: OS Independent" 18 | ] 19 | dependencies = [ 20 | "moderngl", 21 | "numpy", 22 | "plyfile", 23 | "scipy" 24 | ] 25 | requires-python = ">=3.7" 26 | 27 | [project.urls] 28 | Homepage = "https://github.com/EasternJournalist/utils3d" 29 | 30 | [tool.setuptools.packages.find] 31 | where = ["."] 32 | include = ["utils3d*"] 33 | 34 | [tool.setuptools.package-data] 35 | "utils3d.numpy.shaders" = ["*"] -------------------------------------------------------------------------------- /test/io_/write_ply.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | 7 | def run(): 8 | image_uv, image_mesh = utils3d.numpy.utils.image_mesh(128, 128) 9 | image_mesh = image_mesh.reshape(-1, 4) 10 | depth = np.ones((128, 128), dtype=float) * 2 11 | depth[32:96, 32:96] = 1 12 | depth = depth.reshape(-1) 13 | intrinsics = utils3d.numpy.transforms.intrinsics_from_fov(1.0, 128, 128) 14 | intrinsics = utils3d.numpy.transforms.normalize_intrinsics(intrinsics, 128, 128) 15 | extrinsics = utils3d.numpy.transforms.extrinsics_look_at([0, 0, 1], [0, 0, 0], [0, 1, 0]) 16 | pts = utils3d.numpy.transforms.unproject_cv(image_uv, depth, extrinsics, intrinsics) 17 | pts = pts.reshape(-1, 3) 18 | image_mesh = utils3d.numpy.mesh.triangulate(image_mesh, vertices=pts) 19 | utils3d.io.write_ply(os.path.join(os.path.dirname(__file__), '..', 'results_to_check', 'write_ply.ply'), pts, image_mesh) 20 | 21 | if __name__ == '__main__': 22 | run() 23 | -------------------------------------------------------------------------------- /test/numpy_/mesh/compute_face_angle.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | 7 | def run(): 8 | for i in range(100): 9 | if i == 0: 10 | spatial = [] 11 | vertices = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0]], dtype=float) 12 | faces = np.array([[0, 1, 2]]) 13 | expected = np.array([[np.pi/2, np.pi/4, np.pi/4]]) 14 | else: 15 | dim = np.random.randint(4) 16 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 17 | N = np.random.randint(100, 1000) 18 | vertices = np.random.rand(*spatial, N, 3) 19 | L = np.random.randint(1, 1000) 20 | faces = np.random.randint(0, N, size=(*spatial, L, 3)) 21 | faces[..., 1] = (faces[..., 0] + 1) % N 22 | faces[..., 2] = (faces[..., 0] + 2) % N 23 | 24 | faces_ = faces.reshape(-1, L, 3) 25 | vertices_ = vertices.reshape(-1, N, 3) 26 | N = vertices_.shape[0] 27 | expected = np.zeros((N, L, 3), dtype=float) 28 | for i in range(3): 29 | edge0 = vertices_[np.arange(N)[:, None], faces_[..., (i+1)%3]] - vertices_[np.arange(N)[:, None], faces_[..., i]] 30 | edge1 = vertices_[np.arange(N)[:, None], faces_[..., (i+2)%3]] - vertices_[np.arange(N)[:, None], faces_[..., i]] 31 | expected[..., i] = np.arccos(np.sum( 32 | edge0 / np.linalg.norm(edge0, axis=-1, keepdims=True) * \ 33 | edge1 / np.linalg.norm(edge1, axis=-1, keepdims=True), 34 | axis=-1 35 | )) 36 | expected = expected.reshape(*spatial, L, 3) 37 | 38 | actual = utils3d.numpy.compute_face_angle(vertices, faces) 39 | 40 | assert np.allclose(expected, actual), '\n' + \ 41 | 'Input:\n' + \ 42 | f'{faces}\n' + \ 43 | 'Actual:\n' + \ 44 | f'{actual}\n' + \ 45 | 'Expected:\n' + \ 46 | f'{expected}' 47 | -------------------------------------------------------------------------------- /test/numpy_/mesh/compute_face_normal.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | 7 | def run(): 8 | for i in range(100): 9 | if i == 0: 10 | spatial = [] 11 | vertices = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0]], dtype=float) 12 | faces = np.array([[0, 1, 2]]) 13 | expected = np.array([[0, 0, 1]]) 14 | else: 15 | dim = np.random.randint(4) 16 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 17 | N = np.random.randint(100, 1000) 18 | vertices = np.random.rand(*spatial, N, 3) 19 | L = np.random.randint(1, 1000) 20 | faces = np.random.randint(0, N, size=(*spatial, L, 3)) 21 | faces[..., 1] = (faces[..., 0] + 1) % N 22 | faces[..., 2] = (faces[..., 0] + 2) % N 23 | 24 | faces_ = faces.reshape(-1, L, 3) 25 | vertices_ = vertices.reshape(-1, N, 3) 26 | N = vertices_.shape[0] 27 | expected = np.cross( 28 | vertices_[np.arange(N)[:, None], faces_[..., 1]] - vertices_[np.arange(N)[:, None], faces_[..., 0]], 29 | vertices_[np.arange(N)[:, None], faces_[..., 2]] - vertices_[np.arange(N)[:, None], faces_[..., 0]] 30 | ).reshape(*spatial, L, 3) 31 | expected = np.nan_to_num(expected / np.linalg.norm(expected, axis=-1, keepdims=True)) 32 | 33 | actual = utils3d.numpy.compute_face_normal(vertices, faces) 34 | 35 | assert np.allclose(expected, actual), '\n' + \ 36 | 'Input:\n' + \ 37 | f'{faces}\n' + \ 38 | 'Actual:\n' + \ 39 | f'{actual}\n' + \ 40 | 'Expected:\n' + \ 41 | f'{expected}' 42 | -------------------------------------------------------------------------------- /test/numpy_/mesh/compute_vertex_normal.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | from trimesh import geometry 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | vertices = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0]], dtype=float) 13 | faces = np.array([[0, 1, 2]]) 14 | expected = np.array([[0, 0, 1], [0, 0, 1], [0, 0, 1]]) 15 | else: 16 | dim = np.random.randint(4) 17 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 18 | N = np.random.randint(100, 1000) 19 | vertices = np.random.rand(*spatial, N, 3) 20 | L = np.random.randint(1, 1000) 21 | faces = np.random.randint(0, N, size=(*spatial, L, 3)) 22 | faces[..., 1] = (faces[..., 0] + 1) % N 23 | faces[..., 2] = (faces[..., 0] + 2) % N 24 | face_normals = utils3d.numpy.compute_face_normal(vertices, faces) 25 | 26 | faces_ = faces.reshape(-1, L, 3) 27 | face_normals = face_normals.reshape(-1, L, 3) 28 | vertices_normals = [] 29 | for face, face_normal in zip(faces_, face_normals): 30 | vertices_normals.append( 31 | geometry.mean_vertex_normals(N, face, face_normal) 32 | ) 33 | expected = np.array(vertices_normals).reshape(*spatial, N, 3) 34 | 35 | actual = utils3d.numpy.compute_vertex_normal(vertices, faces) 36 | 37 | assert np.allclose(expected, actual), '\n' + \ 38 | 'Input:\n' + \ 39 | f'{faces}\n' + \ 40 | 'Actual:\n' + \ 41 | f'{actual}\n' + \ 42 | 'Expected:\n' + \ 43 | f'{expected}' 44 | -------------------------------------------------------------------------------- /test/numpy_/mesh/compute_vertex_normal_weighted.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | from trimesh import geometry 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | vertices = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0]], dtype=float) 13 | faces = np.array([[0, 1, 2]]) 14 | expected = np.array([[0, 0, 1], [0, 0, 1], [0, 0, 1]]) 15 | else: 16 | dim = np.random.randint(4) 17 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 18 | N = np.random.randint(100, 1000) 19 | vertices = np.random.rand(*spatial, N, 3) 20 | L = np.random.randint(1, 1000) 21 | faces = np.random.randint(0, N, size=(*spatial, L, 3)) 22 | faces[..., 1] = (faces[..., 0] + 1) % N 23 | faces[..., 2] = (faces[..., 0] + 2) % N 24 | face_normals = utils3d.numpy.compute_face_normal(vertices, faces) 25 | face_angles = utils3d.numpy.compute_face_angle(vertices, faces) 26 | 27 | faces_ = faces.reshape(-1, L, 3) 28 | face_normals = face_normals.reshape(-1, L, 3) 29 | face_angles = face_angles.reshape(-1, L, 3) 30 | vertices_normals = [] 31 | for face, face_normal, face_angle in zip(faces_, face_normals, face_angles): 32 | vertices_normals.append( 33 | geometry.weighted_vertex_normals(N, face, face_normal, face_angle) 34 | ) 35 | expected = np.array(vertices_normals).reshape(*spatial, N, 3) 36 | 37 | actual = utils3d.numpy.compute_vertex_normal_weighted(vertices, faces) 38 | 39 | assert np.allclose(expected, actual), '\n' + \ 40 | 'Input:\n' + \ 41 | f'{faces}\n' + \ 42 | 'Actual:\n' + \ 43 | f'{actual}\n' + \ 44 | 'Expected:\n' + \ 45 | f'{expected}' 46 | -------------------------------------------------------------------------------- /test/numpy_/mesh/merge_duplicate_vertices.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | 7 | def run(): 8 | for i in range(100): 9 | if i == 0: 10 | spatial = [] 11 | vertices = np.array([[0, 0, 0], [1, 0, 0], [1, 0, 0]], dtype=float) 12 | faces = np.array([[0, 1, 2]]) 13 | expected_vertices = np.array([[0, 0, 0], [1, 0, 0]]) 14 | expected_faces = np.array([[0, 1, 1]]) 15 | expected = expected_vertices[expected_faces] 16 | else: 17 | N = np.random.randint(100, 1000) 18 | vertices = np.random.rand(N, 3) 19 | L = np.random.randint(1, 1000) 20 | faces = np.random.randint(0, N, size=(L, 3)) 21 | faces[..., 1] = (faces[..., 0] + 1) % N 22 | faces[..., 2] = (faces[..., 0] + 2) % N 23 | vertices[-(N//2):] = vertices[:N//2] 24 | 25 | expected_vertices = vertices[:-(N//2)].copy() 26 | expected_faces = faces.copy() 27 | expected_faces[expected_faces >= N - N//2] -= N - N//2 28 | expected = expected_vertices[expected_faces] 29 | 30 | actual_vertices, actual_faces = utils3d.numpy.merge_duplicate_vertices(vertices, faces) 31 | actual = actual_vertices[actual_faces] 32 | 33 | assert expected_vertices.shape == actual_vertices.shape and np.allclose(expected, actual), '\n' + \ 34 | 'Input:\n' + \ 35 | f'{faces}\n' + \ 36 | 'Actual:\n' + \ 37 | f'{actual}\n' + \ 38 | 'Expected:\n' + \ 39 | f'{expected}' 40 | -------------------------------------------------------------------------------- /test/numpy_/mesh/remove_corrupted_faces.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | 7 | def run(): 8 | for i in range(100): 9 | if i == 0: 10 | faces = np.array([[0, 1, 2], [0, 2, 2], [0, 2, 3]]) 11 | expected = np.array([[0, 1, 2], [0, 2, 3]]) 12 | else: 13 | L = np.random.randint(1, 1000) 14 | N = np.random.randint(100, 1000) 15 | faces = np.random.randint(0, N, size=(L, 3)) 16 | faces[..., 1] = (faces[..., 0] + 1) % N 17 | faces[..., 2] = (faces[..., 0] + 2) % N 18 | corrupted = np.random.randint(0, 2, size=L).astype(bool) 19 | faces[corrupted, 1] = faces[corrupted, 0] 20 | expected = faces[~corrupted] 21 | 22 | actual = utils3d.numpy.remove_corrupted_faces(faces) 23 | 24 | assert np.allclose(expected, actual), '\n' + \ 25 | 'Input:\n' + \ 26 | f'{faces}\n' + \ 27 | 'Actual:\n' + \ 28 | f'{actual}\n' + \ 29 | 'Expected:\n' + \ 30 | f'{expected}' 31 | -------------------------------------------------------------------------------- /test/numpy_/mesh/triangulate.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | 7 | def run(): 8 | for i in range(100): 9 | if i == 0: 10 | spatial = [] 11 | L = 1 12 | N = 5 13 | faces = np.array([[0, 1, 2, 3, 4]]) 14 | expected = np.array([[0, 1, 2], [0, 2, 3], [0, 3, 4]]) 15 | else: 16 | dim = np.random.randint(4) 17 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 18 | L = np.random.randint(1, 1000) 19 | N = np.random.randint(3, 10) 20 | faces = np.random.randint(0, 10000, size=(*spatial, L, N)) 21 | 22 | loop_indices = [[0, i, i + 1] for i in range(1, N - 1)] 23 | expected = faces[..., loop_indices].reshape((*spatial, L * (N - 2), 3)) 24 | 25 | actual = utils3d.numpy.triangulate(faces) 26 | 27 | assert np.allclose(expected, actual), '\n' + \ 28 | 'Input:\n' + \ 29 | f'{faces}\n' + \ 30 | 'Actual:\n' + \ 31 | f'{actual}\n' + \ 32 | 'Expected:\n' + \ 33 | f'{expected}' 34 | -------------------------------------------------------------------------------- /test/numpy_/rasterization/warp_image_by_depth.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import imageio 7 | 8 | def run(): 9 | depth = np.ones((128, 128), dtype=np.float32) * 2 10 | depth[32:48, 32:48] = 1 11 | intrinsics = utils3d.numpy.transforms.intrinsics(1.0, 1.0, 0.5, 0.5).astype(np.float32) 12 | extrinsics_src = utils3d.numpy.transforms.extrinsics_look_at([0, 0, 1], [0, 0, 0], [0, 1, 0]).astype(np.float32) 13 | extrinsics_tgt = utils3d.numpy.transforms.extrinsics_look_at([1, 0, 1], [0, 0, 0], [0, 1, 0]).astype(np.float32) 14 | ctx = utils3d.numpy.rasterization.RastContext( 15 | standalone=True, 16 | backend='egl', 17 | device_index=0, 18 | ) 19 | uv, _ = utils3d.numpy.rasterization.warp_image_by_depth( 20 | ctx, 21 | depth, 22 | extrinsics_src=extrinsics_src, 23 | extrinsics_tgt=extrinsics_tgt, 24 | intrinsics_src=intrinsics 25 | ) 26 | uv = (np.concatenate([uv, np.zeros((128, 128, 1), dtype=np.float32)], axis=-1) * 255).astype(np.uint8) 27 | imageio.imwrite(os.path.join(os.path.dirname(__file__), '..', '..', 'results_to_check', 'warp_image_uv.png'), uv) 28 | 29 | if __name__ == '__main__': 30 | run() 31 | -------------------------------------------------------------------------------- /test/numpy_/transforms/crop_intrinsic.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | 7 | def run(): 8 | for i in range(100): 9 | if i == 0: 10 | spatial = [] 11 | else: 12 | dim = np.random.randint(4) 13 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 14 | fov = np.random.uniform(5 / 180 * np.pi, 175 / 180 * np.pi, spatial) 15 | width = np.random.uniform(1, 10000, spatial) 16 | height = np.random.uniform(1, 10000, spatial) 17 | left = np.random.uniform(0, width, spatial) 18 | top = np.random.uniform(0, height, spatial) 19 | crop_width = np.random.uniform(0, width - left, spatial) 20 | crop_height = np.random.uniform(0, height - top, spatial) 21 | 22 | focal = np.maximum(width, height) / (2 * np.tan(fov / 2)) 23 | cx = width / 2 - left 24 | cy = height / 2 - top 25 | expected = np.zeros((*spatial, 3, 3)) 26 | expected[..., 0, 0] = focal 27 | expected[..., 1, 1] = focal 28 | expected[..., 0, 2] = cx 29 | expected[..., 1, 2] = cy 30 | expected[..., 2, 2] = 1 31 | expected = utils3d.numpy.normalize_intrinsics(expected, crop_width, crop_height) 32 | 33 | actual = utils3d.numpy.crop_intrinsics( 34 | utils3d.numpy.normalize_intrinsics( 35 | utils3d.numpy.intrinsics_from_fov(fov, width, height), 36 | width, height 37 | ), 38 | width, height, left, top, crop_width, crop_height 39 | ) 40 | 41 | assert np.allclose(expected, actual), '\n' + \ 42 | 'Input:\n' + \ 43 | f'\tfov: {fov}\n' + \ 44 | f'\twidth: {width}\n' + \ 45 | f'\theight: {height}\n' + \ 46 | f'\tleft: {left}\n' + \ 47 | f'\ttop: {top}\n' + \ 48 | f'\tcrop_width: {crop_width}\n' + \ 49 | f'\tcrop_height: {crop_height}\n' + \ 50 | 'Actual:\n' + \ 51 | f'{actual}\n' + \ 52 | 'Expected:\n' + \ 53 | f'{expected}' 54 | -------------------------------------------------------------------------------- /test/numpy_/transforms/extrinsic_look_at.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import glm 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | else: 13 | dim = np.random.randint(4) 14 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 15 | eye = np.random.uniform(-10, 10, [*spatial, 3]).astype(np.float32) 16 | lookat = np.random.uniform(-10, 10, [*spatial, 3]).astype(np.float32) 17 | up = np.random.uniform(-10, 10, [*spatial, 3]).astype(np.float32) 18 | 19 | expected = [] 20 | for i in range(np.prod(spatial) if len(spatial) > 0 else 1): 21 | expected.append(utils3d.numpy.view_to_extrinsics(np.array(glm.lookAt( 22 | glm.vec3(eye.reshape([-1, 3])[i]), 23 | glm.vec3(lookat.reshape([-1, 3])[i]), 24 | glm.vec3(up.reshape([-1, 3])[i]) 25 | )))) 26 | expected = np.concatenate(expected, axis=0).reshape([*spatial, 4, 4]) 27 | 28 | actual = utils3d.numpy.extrinsics_look_at(eye, lookat, up) 29 | 30 | assert np.allclose(expected, actual, 1e-5, 1e-5), '\n' + \ 31 | 'Input:\n' + \ 32 | f'eye: {eye}\n' + \ 33 | f'lookat: {lookat}\n' + \ 34 | f'up: {up}\n' + \ 35 | 'Actual:\n' + \ 36 | f'{actual}\n' + \ 37 | 'Expected:\n' + \ 38 | f'{expected}' 39 | -------------------------------------------------------------------------------- /test/numpy_/transforms/extrinsic_to_view.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import glm 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | else: 13 | dim = np.random.randint(4) 14 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 15 | eye = np.random.uniform(-10, 10, [*spatial, 3]).astype(np.float32) 16 | lookat = np.random.uniform(-10, 10, [*spatial, 3]).astype(np.float32) 17 | up = np.random.uniform(-10, 10, [*spatial, 3]).astype(np.float32) 18 | 19 | expected = utils3d.numpy.view_look_at(eye, lookat, up) 20 | 21 | actual = utils3d.numpy.view_to_extrinsics(utils3d.numpy.extrinsics_look_at(eye, lookat, up)) 22 | 23 | assert np.allclose(expected, actual, 1e-5, 1e-5), '\n' + \ 24 | 'Input:\n' + \ 25 | f'eye: {eye}\n' + \ 26 | f'lookat: {lookat}\n' + \ 27 | f'up: {up}\n' + \ 28 | 'Actual:\n' + \ 29 | f'{actual}\n' + \ 30 | 'Expected:\n' + \ 31 | f'{expected}' 32 | -------------------------------------------------------------------------------- /test/numpy_/transforms/intrinsic.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | 7 | def run(): 8 | for i in range(100): 9 | if i == 0: 10 | spatial = [] 11 | else: 12 | dim = np.random.randint(4) 13 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 14 | focal_x = np.random.uniform(1, 10000, spatial) 15 | focal_y = np.random.uniform(1, 10000, spatial) 16 | center_x = np.random.uniform(1, 10000, spatial) 17 | center_y = np.random.uniform(1, 10000, spatial) 18 | 19 | expected = np.zeros((*spatial, 3, 3)) 20 | expected[..., 0, 0] = focal_x 21 | expected[..., 1, 1] = focal_y 22 | expected[..., 0, 2] = center_x 23 | expected[..., 1, 2] = center_y 24 | expected[..., 2, 2] = 1 25 | 26 | actual = utils3d.numpy.intrinsics(focal_x, focal_y, center_x, center_y) 27 | 28 | assert np.allclose(expected, actual), '\n' + \ 29 | 'Input:\n' + \ 30 | f'\tfocal_x: {focal_x}\n' + \ 31 | f'\tfocal_y: {focal_y}\n' + \ 32 | f'\tcenter_x: {center_x}\n' + \ 33 | f'\tcenter_y: {center_y}\n' + \ 34 | 'Actual:\n' + \ 35 | f'{actual}\n' + \ 36 | 'Expected:\n' + \ 37 | f'{expected}' 38 | -------------------------------------------------------------------------------- /test/numpy_/transforms/intrinsic_from_fov.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | 7 | def run(): 8 | for i in range(100): 9 | if i == 0: 10 | spatial = [] 11 | else: 12 | dim = np.random.randint(4) 13 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 14 | fov = np.random.uniform(5 / 180 * np.pi, 175 / 180 * np.pi, spatial) 15 | width = np.random.uniform(1, 10000, spatial) 16 | height = np.random.uniform(1, 10000, spatial) 17 | 18 | focal = np.maximum(width, height) / (2 * np.tan(fov / 2)) 19 | cx = width / 2 20 | cy = height / 2 21 | expected = np.zeros((*spatial, 3, 3)) 22 | expected[..., 0, 0] = focal 23 | expected[..., 1, 1] = focal 24 | expected[..., 0, 2] = cx 25 | expected[..., 1, 2] = cy 26 | expected[..., 2, 2] = 1 27 | 28 | actual = utils3d.numpy.intrinsics_from_fov(fov, width, height) 29 | 30 | assert np.allclose(expected, actual), '\n' + \ 31 | 'Input:\n' + \ 32 | f'\tfov: {fov}\n' + \ 33 | f'\twidth: {width}\n' + \ 34 | f'\theight: {height}\n' + \ 35 | 'Actual:\n' + \ 36 | f'{actual}\n' + \ 37 | 'Expected:\n' + \ 38 | f'{expected}' 39 | -------------------------------------------------------------------------------- /test/numpy_/transforms/intrinsic_from_fov_xy.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | 7 | def run(): 8 | for i in range(100): 9 | if i == 0: 10 | spatial = [] 11 | else: 12 | dim = np.random.randint(4) 13 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 14 | fov_x = np.random.uniform(5 / 180 * np.pi, 175 / 180 * np.pi, spatial) 15 | fov_y = np.random.uniform(5 / 180 * np.pi, 175 / 180 * np.pi, spatial) 16 | 17 | focal_x = 0.5 / np.tan(fov_x / 2) 18 | focal_y = 0.5 / np.tan(fov_y / 2) 19 | cx = cy = 0.5 20 | expected = np.zeros((*spatial, 3, 3)) 21 | expected[..., 0, 0] = focal_x 22 | expected[..., 1, 1] = focal_y 23 | expected[..., 0, 2] = cx 24 | expected[..., 1, 2] = cy 25 | expected[..., 2, 2] = 1 26 | 27 | actual = utils3d.numpy.intrinsics_from_fov_xy(fov_x, fov_y) 28 | 29 | assert np.allclose(expected, actual), '\n' + \ 30 | 'Input:\n' + \ 31 | f'\tfov_x: {fov_x}\n' + \ 32 | f'\tfov_y: {fov_y}\n' + \ 33 | 'Actual:\n' + \ 34 | f'{actual}\n' + \ 35 | 'Expected:\n' + \ 36 | f'{expected}' 37 | -------------------------------------------------------------------------------- /test/numpy_/transforms/intrinsic_to_perspective.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import glm 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | else: 13 | dim = np.random.randint(4) 14 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 15 | fov_x = np.random.uniform(5 / 180 * np.pi, 175 / 180 * np.pi, spatial) 16 | fov_y = np.random.uniform(5 / 180 * np.pi, 175 / 180 * np.pi, spatial) 17 | near = np.random.uniform(0.1, 100, spatial) 18 | far = np.random.uniform(near*2, 1000, spatial) 19 | 20 | expected = utils3d.numpy.perspective_from_fov_xy(fov_x, fov_y, near, far) 21 | 22 | actual = utils3d.numpy.intrinsics_to_perspective( 23 | utils3d.numpy.intrinsics_from_fov_xy(fov_x, fov_y), 24 | near, 25 | far 26 | ) 27 | 28 | assert np.allclose(expected, actual), '\n' + \ 29 | 'Input:\n' + \ 30 | f'\tfov_x: {fov_x}\n' + \ 31 | f'\tfov_y: {fov_y}\n' + \ 32 | f'\tnear: {near}\n' + \ 33 | f'\tfar: {far}\n' + \ 34 | 'Actual:\n' + \ 35 | f'{actual}\n' + \ 36 | 'Expected:\n' + \ 37 | f'{expected}' 38 | -------------------------------------------------------------------------------- /test/numpy_/transforms/linearize_depth.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import glm 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | else: 13 | dim = np.random.randint(4) 14 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 15 | near = np.random.uniform(0.1, 100, spatial) 16 | far = np.random.uniform(near*2, 1000, spatial) 17 | depth = np.random.uniform(near, far, spatial) 18 | 19 | expected = depth 20 | 21 | actual = utils3d.numpy.depth_buffer_to_linear( 22 | utils3d.numpy.project_depth(depth, near, far), 23 | near, far 24 | ) 25 | 26 | assert np.allclose(expected, actual), '\n' + \ 27 | 'Input:\n' + \ 28 | f'\tdepth: {depth}\n' + \ 29 | f'\tnear: {near}\n' + \ 30 | f'\tfar: {far}\n' + \ 31 | 'Actual:\n' + \ 32 | f'{actual}\n' + \ 33 | 'Expected:\n' + \ 34 | f'{expected}' 35 | -------------------------------------------------------------------------------- /test/numpy_/transforms/normalize_intrinsic.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | 7 | def run(): 8 | for i in range(100): 9 | if i == 0: 10 | spatial = [] 11 | else: 12 | dim = np.random.randint(4) 13 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 14 | fov = np.random.uniform(5 / 180 * np.pi, 175 / 180 * np.pi, spatial) 15 | width = np.random.uniform(1, 10000, spatial) 16 | height = np.random.uniform(1, 10000, spatial) 17 | fov_x = np.where(width >= height, fov, 2 * np.arctan(np.tan(fov / 2) * width / height)) 18 | fov_y = np.where(width >= height, 2 * np.arctan(np.tan(fov / 2) * height / width), fov) 19 | 20 | expected = utils3d.numpy.intrinsics_from_fov_xy(fov_x, fov_y) 21 | 22 | actual = utils3d.numpy.normalize_intrinsics(utils3d.numpy.intrinsics_from_fov(fov, width, height), width, height) 23 | 24 | assert np.allclose(expected, actual), '\n' + \ 25 | 'Input:\n' + \ 26 | f'\tfov: {fov}\n' + \ 27 | f'\twidth: {width}\n' + \ 28 | f'\theight: {height}\n' + \ 29 | 'Actual:\n' + \ 30 | f'{actual}\n' + \ 31 | 'Expected:\n' + \ 32 | f'{expected}' 33 | -------------------------------------------------------------------------------- /test/numpy_/transforms/perspective.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import glm 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | else: 13 | dim = np.random.randint(4) 14 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 15 | fovy = np.random.uniform(5 / 180 * np.pi, 175 / 180 * np.pi, spatial) 16 | aspect = np.random.uniform(0.01, 100, spatial) 17 | near = np.random.uniform(0.1, 100, spatial) 18 | far = np.random.uniform(near*2, 1000, spatial) 19 | 20 | expected = [] 21 | for i in range(np.prod(spatial) if len(spatial) > 0 else 1): 22 | expected.append(glm.perspective(fovy.flat[i], aspect.flat[i], near.flat[i], far.flat[i])) 23 | expected = np.concatenate(expected, axis=0).reshape(*spatial, 4, 4) 24 | 25 | actual = utils3d.numpy.perspective(fovy, aspect, near, far) 26 | 27 | assert np.allclose(expected, actual), '\n' + \ 28 | 'Input:\n' + \ 29 | f'\tfovy: {fovy}\n' + \ 30 | f'\taspect: {aspect}\n' + \ 31 | f'\tnear: {near}\n' + \ 32 | f'\tfar: {far}\n' + \ 33 | 'Actual:\n' + \ 34 | f'{actual}\n' + \ 35 | 'Expected:\n' + \ 36 | f'{expected}' 37 | -------------------------------------------------------------------------------- /test/numpy_/transforms/perspective_from_fov.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import glm 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | else: 13 | dim = np.random.randint(4) 14 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 15 | fov = np.random.uniform(5 / 180 * np.pi, 175 / 180 * np.pi, spatial) 16 | width = np.random.uniform(1, 10000, spatial) 17 | height = np.random.uniform(1, 10000, spatial) 18 | near = np.random.uniform(0.1, 100, spatial) 19 | far = np.random.uniform(near*2, 1000, spatial) 20 | 21 | fov_y = 2 * np.arctan(np.tan(fov / 2) * height / np.maximum(width, height)) 22 | expected = [] 23 | for i in range(np.prod(spatial) if len(spatial) > 0 else 1): 24 | expected.append(glm.perspective(fov_y.flat[i], width.flat[i] / height.flat[i], near.flat[i], far.flat[i])) 25 | expected = np.concatenate(expected, axis=0).reshape(*spatial, 4, 4) 26 | 27 | actual = utils3d.numpy.perspective_from_fov(fov, width, height, near, far) 28 | 29 | assert np.allclose(expected, actual), '\n' + \ 30 | 'Input:\n' + \ 31 | f'\tfov: {fov}\n' + \ 32 | f'\twidth: {width}\n' + \ 33 | f'\theight: {height}\n' + \ 34 | f'\tnear: {near}\n' + \ 35 | f'\tfar: {far}\n' + \ 36 | 'Actual:\n' + \ 37 | f'{actual}\n' + \ 38 | 'Expected:\n' + \ 39 | f'{expected}' 40 | -------------------------------------------------------------------------------- /test/numpy_/transforms/perspective_from_fov_xy.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import glm 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | else: 13 | dim = np.random.randint(4) 14 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 15 | fov_x = np.random.uniform(5 / 180 * np.pi, 175 / 180 * np.pi, spatial) 16 | fov_y = np.random.uniform(5 / 180 * np.pi, 175 / 180 * np.pi, spatial) 17 | near = np.random.uniform(0.1, 100, spatial) 18 | far = np.random.uniform(near*2, 1000, spatial) 19 | 20 | aspect = np.tan(fov_x / 2) / np.tan(fov_y / 2) 21 | expected = [] 22 | for i in range(np.prod(spatial) if len(spatial) > 0 else 1): 23 | expected.append(glm.perspective(fov_y.flat[i], aspect.flat[i], near.flat[i], far.flat[i])) 24 | expected = np.concatenate(expected, axis=0).reshape(*spatial, 4, 4) 25 | 26 | actual = utils3d.numpy.perspective_from_fov_xy(fov_x, fov_y, near, far) 27 | 28 | assert np.allclose(expected, actual), '\n' + \ 29 | 'Input:\n' + \ 30 | f'\tfov_x: {fov_x}\n' + \ 31 | f'\tfov_y: {fov_y}\n' + \ 32 | f'\tnear: {near}\n' + \ 33 | f'\tfar: {far}\n' + \ 34 | 'Actual:\n' + \ 35 | f'{actual}\n' + \ 36 | 'Expected:\n' + \ 37 | f'{expected}' 38 | -------------------------------------------------------------------------------- /test/numpy_/transforms/perspective_to_intrinsic.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import glm 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | else: 13 | dim = np.random.randint(4) 14 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 15 | fov_x = np.random.uniform(5 / 180 * np.pi, 175 / 180 * np.pi, spatial) 16 | fov_y = np.random.uniform(5 / 180 * np.pi, 175 / 180 * np.pi, spatial) 17 | near = np.random.uniform(0.1, 100) 18 | far = np.random.uniform(near*2, 1000) 19 | 20 | expected = utils3d.numpy.intrinsics_from_fov_xy(fov_x, fov_y) 21 | 22 | actual = utils3d.numpy.perspective_to_intrinsics( 23 | utils3d.numpy.perspective_from_fov_xy(fov_x, fov_y, near, far) 24 | ) 25 | 26 | assert np.allclose(expected, actual), '\n' + \ 27 | 'Input:\n' + \ 28 | f'\tfov_x: {fov_x}\n' + \ 29 | f'\tfov_y: {fov_y}\n' + \ 30 | f'\tnear: {near}\n' + \ 31 | f'\tfar: {far}\n' + \ 32 | 'Actual:\n' + \ 33 | f'{actual}\n' + \ 34 | 'Expected:\n' + \ 35 | f'{expected}' 36 | -------------------------------------------------------------------------------- /test/numpy_/transforms/pixel_to_ndc.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | 7 | def run(): 8 | for i in range(100): 9 | H = np.random.randint(1, 1000) 10 | W = np.random.randint(1, 1000) 11 | x, y = np.meshgrid(np.arange(W), np.arange(H), indexing='xy') 12 | pixel = np.stack([x, y], axis=-1) 13 | 14 | expected = np.stack( 15 | np.meshgrid( 16 | np.linspace(-1 + 1 / W, 1 - 1 / W, W), 17 | np.linspace(1 - 1 / H, -1 + 1 / H, H), 18 | indexing='xy' 19 | ), 20 | axis=-1 21 | ) 22 | 23 | actual = utils3d.numpy.pixel_to_ndc(pixel, W, H) 24 | 25 | assert np.allclose(expected, actual), '\n' + \ 26 | 'Input:\n' + \ 27 | f'\tH: {H}\n' + \ 28 | f'\tW: {W}\n' + \ 29 | 'Actual:\n' + \ 30 | f'{actual}\n' + \ 31 | 'Expected:\n' + \ 32 | f'{expected}' 33 | -------------------------------------------------------------------------------- /test/numpy_/transforms/pixel_to_uv.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | 7 | def run(): 8 | for i in range(100): 9 | H = np.random.randint(1, 1000) 10 | W = np.random.randint(1, 1000) 11 | x, y = np.meshgrid(np.arange(W), np.arange(H), indexing='xy') 12 | pixel = np.stack([x, y], axis=-1) 13 | 14 | expected = np.stack( 15 | np.meshgrid( 16 | np.linspace(0.5 / W, 1 - 0.5 / W, W), 17 | np.linspace(0.5 / H, 1 - 0.5 / H, H), 18 | indexing='xy' 19 | ), 20 | axis=-1 21 | ) 22 | 23 | actual = utils3d.numpy.pixel_to_uv(pixel, W, H) 24 | 25 | assert np.allclose(expected, actual), '\n' + \ 26 | 'Input:\n' + \ 27 | f'\tH: {H}\n' + \ 28 | f'\tW: {W}\n' + \ 29 | 'Actual:\n' + \ 30 | f'{actual}\n' + \ 31 | 'Expected:\n' + \ 32 | f'{expected}' 33 | -------------------------------------------------------------------------------- /test/numpy_/transforms/project_cv.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import glm 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | N = 1 13 | else: 14 | dim = np.random.randint(4) 15 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 16 | N = np.random.randint(1, 10) 17 | focal_x = np.random.uniform(0, 10, spatial) 18 | focal_y = np.random.uniform(0, 10, spatial) 19 | center_x = np.random.uniform(0, 1, spatial) 20 | center_y = np.random.uniform(0, 1, spatial) 21 | eye = np.random.uniform(-10, 10, [*spatial, 3]) 22 | lookat = np.random.uniform(-10, 10, [*spatial, 3]) 23 | up = np.random.uniform(-10, 10, [*spatial, 3]) 24 | points = np.random.uniform(-10, 10, [*spatial, N, 3]) 25 | 26 | pts = points - eye[..., None, :] 27 | z_axis = lookat - eye 28 | x_axis = np.cross(-up, z_axis) 29 | y_axis = np.cross(z_axis, x_axis) 30 | x_axis = x_axis / np.linalg.norm(x_axis, axis=-1, keepdims=True) 31 | y_axis = y_axis / np.linalg.norm(y_axis, axis=-1, keepdims=True) 32 | z_axis = z_axis / np.linalg.norm(z_axis, axis=-1, keepdims=True) 33 | z = (pts * z_axis[..., None, :]).sum(axis=-1) 34 | x = (pts * x_axis[..., None, :]).sum(axis=-1) 35 | y = (pts * y_axis[..., None, :]).sum(axis=-1) 36 | x = (x / z * focal_x[..., None] + center_x[..., None]) 37 | y = (y / z * focal_y[..., None] + center_y[..., None]) 38 | expected = np.stack([x, y], axis=-1) 39 | 40 | actual, _ = utils3d.numpy.transforms.project_cv(points, 41 | utils3d.numpy.extrinsics_look_at(eye, lookat, up), 42 | utils3d.numpy.intrinsics(focal_x, focal_y, center_x, center_y)) 43 | 44 | assert np.allclose(expected, actual), '\n' + \ 45 | 'Input:\n' + \ 46 | f'\tfocal_x: {focal_x}\n' + \ 47 | f'\tfocal_y: {focal_y}\n' + \ 48 | f'\tcenter_x: {center_x}\n' + \ 49 | f'\tcenter_y: {center_y}\n' + \ 50 | f'\teye: {eye}\n' + \ 51 | f'\tlookat: {lookat}\n' + \ 52 | f'\tup: {up}\n' + \ 53 | f'\tpoints: {points}\n' + \ 54 | 'Actual:\n' + \ 55 | f'{actual}\n' + \ 56 | 'Expected:\n' + \ 57 | f'{expected}' 58 | -------------------------------------------------------------------------------- /test/numpy_/transforms/project_depth.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import glm 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | else: 13 | dim = np.random.randint(4) 14 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 15 | near = np.random.uniform(0.1, 100, spatial) 16 | far = np.random.uniform(near*2, 1000, spatial) 17 | depth = np.random.uniform(near, far, spatial) 18 | 19 | proj = utils3d.numpy.perspective(1.0, 1.0, near, far)[..., 2, 2:4] 20 | expected = ((proj[..., 0] * -depth + proj[..., 1]) / depth) * 0.5 + 0.5 21 | 22 | actual = utils3d.numpy.project_depth(depth, near, far) 23 | 24 | assert np.allclose(expected, actual), '\n' + \ 25 | 'Input:\n' + \ 26 | f'\tdepth: {depth}\n' + \ 27 | f'\tnear: {near}\n' + \ 28 | f'\tfar: {far}\n' + \ 29 | 'Actual:\n' + \ 30 | f'{actual}\n' + \ 31 | 'Expected:\n' + \ 32 | f'{expected}' 33 | -------------------------------------------------------------------------------- /test/numpy_/transforms/project_gl.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import glm 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | N = 1 13 | else: 14 | dim = np.random.randint(4) 15 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 16 | N = np.random.randint(1, 10) 17 | fovy = np.random.uniform(5 / 180 * np.pi, 175 / 180 * np.pi, spatial) 18 | aspect = np.random.uniform(0.01, 100, spatial) 19 | near = np.random.uniform(0.1, 100, spatial) 20 | far = np.random.uniform(near*2, 1000, spatial) 21 | eye = np.random.uniform(-10, 10, [*spatial, 3]) 22 | lookat = np.random.uniform(-10, 10, [*spatial, 3]) 23 | up = np.random.uniform(-10, 10, [*spatial, 3]) 24 | points = np.random.uniform(-10, 10, [*spatial, N, 3]) 25 | 26 | pts = points - eye[..., None, :] 27 | z_axis = (eye - lookat) 28 | x_axis = np.cross(up, z_axis) 29 | y_axis = np.cross(z_axis, x_axis) 30 | x_axis = x_axis / np.linalg.norm(x_axis, axis=-1, keepdims=True) 31 | y_axis = y_axis / np.linalg.norm(y_axis, axis=-1, keepdims=True) 32 | z_axis = z_axis / np.linalg.norm(z_axis, axis=-1, keepdims=True) 33 | z = (pts * z_axis[..., None, :]).sum(axis=-1) 34 | x = (pts * x_axis[..., None, :]).sum(axis=-1) 35 | y = (pts * y_axis[..., None, :]).sum(axis=-1) 36 | x = (x / -z / np.tan(fovy[..., None] / 2) / aspect[..., None]) * 0.5 + 0.5 37 | y = (y / -z / np.tan(fovy[..., None] / 2)) * 0.5 + 0.5 38 | z = utils3d.numpy.project_depth(-z, near[..., None], far[..., None]) 39 | expected = np.stack([x, y, z], axis=-1) 40 | 41 | actual, _ = utils3d.numpy.transforms.project_gl(points, None, 42 | utils3d.numpy.view_look_at(eye, lookat, up), 43 | utils3d.numpy.perspective(fovy, aspect, near, far)) 44 | 45 | assert np.allclose(expected, actual), '\n' + \ 46 | 'Input:\n' + \ 47 | f'\tfovy: {fovy}\n' + \ 48 | f'\taspect: {aspect}\n' + \ 49 | f'\tnear: {near}\n' + \ 50 | f'\tfar: {far}\n' + \ 51 | f'\teye: {eye}\n' + \ 52 | f'\tlookat: {lookat}\n' + \ 53 | f'\tup: {up}\n' + \ 54 | f'\tpoints: {points}\n' + \ 55 | 'Actual:\n' + \ 56 | f'{actual}\n' + \ 57 | 'Expected:\n' + \ 58 | f'{expected}' 59 | -------------------------------------------------------------------------------- /test/numpy_/transforms/project_gl_cv.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import glm 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | N = 1 13 | else: 14 | dim = np.random.randint(4) 15 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 16 | N = np.random.randint(1, 10) 17 | fovy = np.random.uniform(5 / 180 * np.pi, 175 / 180 * np.pi, spatial) 18 | aspect = np.random.uniform(0.01, 100, spatial) 19 | focal_x = 0.5 / (np.tan(fovy / 2) * aspect) 20 | focal_y = 0.5 / np.tan(fovy / 2) 21 | near = np.random.uniform(0.1, 100, spatial) 22 | far = np.random.uniform(near*2, 1000, spatial) 23 | eye = np.random.uniform(-10, 10, [*spatial, 3]) 24 | lookat = np.random.uniform(-10, 10, [*spatial, 3]) 25 | up = np.random.uniform(-10, 10, [*spatial, 3]) 26 | points = np.random.uniform(-10, 10, [*spatial, N, 3]) 27 | 28 | gl = utils3d.numpy.transforms.project_gl(points, None, 29 | utils3d.numpy.view_look_at(eye, lookat, up), 30 | utils3d.numpy.perspective(fovy, aspect, near, far)) 31 | gl_uv = gl[0][..., :2] 32 | gl_uv[..., 1] = 1 - gl_uv[..., 1] 33 | gl_depth = gl[1] 34 | 35 | cv = utils3d.numpy.transforms.project_cv(points, 36 | utils3d.numpy.extrinsics_look_at(eye, lookat, up), 37 | utils3d.numpy.intrinsics(focal_x, focal_y, 0.5, 0.5)) 38 | cv_uv = cv[0][..., :2] 39 | cv_depth = cv[1] 40 | 41 | assert np.allclose(gl_uv, cv_uv) and np.allclose(gl_depth, cv_depth), '\n' + \ 42 | 'Input:\n' + \ 43 | f'\tfovy: {fovy}\n' + \ 44 | f'\taspect: {aspect}\n' + \ 45 | f'\teye: {eye}\n' + \ 46 | f'\tlookat: {lookat}\n' + \ 47 | f'\tup: {up}\n' + \ 48 | f'\tpoints: {points}\n' + \ 49 | 'GL:\n' + \ 50 | f'{gl}\n' + \ 51 | 'CV:\n' + \ 52 | f'{cv}' 53 | -------------------------------------------------------------------------------- /test/numpy_/transforms/unproject_cv.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import glm 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | N = 1 13 | else: 14 | dim = np.random.randint(4) 15 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 16 | N = np.random.randint(1, 10) 17 | focal_x = np.random.uniform(0, 10, spatial) 18 | focal_y = np.random.uniform(0, 10, spatial) 19 | center_x = np.random.uniform(0, 1, spatial) 20 | center_y = np.random.uniform(0, 1, spatial) 21 | eye = np.random.uniform(-10, 10, [*spatial, 3]) 22 | lookat = np.random.uniform(-10, 10, [*spatial, 3]) 23 | up = np.random.uniform(-10, 10, [*spatial, 3]) 24 | points = np.random.uniform(-10, 10, [*spatial, N, 3]) 25 | 26 | expected = points 27 | 28 | actual = utils3d.numpy.transforms.unproject_cv( 29 | *utils3d.numpy.transforms.project_cv(points, 30 | utils3d.numpy.transforms.extrinsics_look_at(eye, lookat, up), 31 | utils3d.numpy.transforms.intrinsics(focal_x, focal_y, center_x, center_y)), 32 | utils3d.numpy.transforms.extrinsics_look_at(eye, lookat, up), 33 | utils3d.numpy.transforms.intrinsics(focal_x, focal_y, center_x, center_y) 34 | ) 35 | 36 | assert np.allclose(expected, actual), '\n' + \ 37 | 'Input:\n' + \ 38 | f'\tfocal_x: {focal_x}\n' + \ 39 | f'\tfocal_y: {focal_y}\n' + \ 40 | f'\tcenter_x: {center_x}\n' + \ 41 | f'\tcenter_y: {center_y}\n' + \ 42 | f'\teye: {eye}\n' + \ 43 | f'\tlookat: {lookat}\n' + \ 44 | f'\tup: {up}\n' + \ 45 | f'\tpoints: {points}\n' + \ 46 | 'Actual:\n' + \ 47 | f'{actual}\n' + \ 48 | 'Expected:\n' + \ 49 | f'{expected}' 50 | -------------------------------------------------------------------------------- /test/numpy_/transforms/unproject_gl.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import glm 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | N = 1 13 | else: 14 | dim = np.random.randint(4) 15 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 16 | N = np.random.randint(1, 10) 17 | fovy = np.random.uniform(5 / 180 * np.pi, 175 / 180 * np.pi, spatial) 18 | aspect = np.random.uniform(0.01, 100, spatial) 19 | near = np.random.uniform(0.1, 100, spatial) 20 | far = np.random.uniform(near*2, 1000, spatial) 21 | eye = np.random.uniform(-10, 10, [*spatial, 3]) 22 | lookat = np.random.uniform(-10, 10, [*spatial, 3]) 23 | up = np.random.uniform(-10, 10, [*spatial, 3]) 24 | points = np.random.uniform(-10, 10, [*spatial, N, 3]) 25 | 26 | expected = points 27 | 28 | actual = utils3d.numpy.transforms.unproject_gl( 29 | utils3d.numpy.transforms.project_gl(points, None, 30 | utils3d.numpy.view_look_at(eye, lookat, up), 31 | utils3d.numpy.perspective(fovy, aspect, near, far))[0], 32 | None, 33 | utils3d.numpy.view_look_at(eye, lookat, up), 34 | utils3d.numpy.perspective(fovy, aspect, near, far) 35 | ) 36 | 37 | assert np.allclose(expected, actual), '\n' + \ 38 | 'Input:\n' + \ 39 | f'\tfovy: {fovy}\n' + \ 40 | f'\taspect: {aspect}\n' + \ 41 | f'\tnear: {near}\n' + \ 42 | f'\tfar: {far}\n' + \ 43 | f'\teye: {eye}\n' + \ 44 | f'\tlookat: {lookat}\n' + \ 45 | f'\tup: {up}\n' + \ 46 | f'\tpoints: {points}\n' + \ 47 | 'Actual:\n' + \ 48 | f'{actual}\n' + \ 49 | 'Expected:\n' + \ 50 | f'{expected}' 51 | -------------------------------------------------------------------------------- /test/numpy_/transforms/view_look_at.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import glm 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | else: 13 | dim = np.random.randint(4) 14 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 15 | eye = np.random.uniform(-10, 10, [*spatial, 3]).astype(np.float32) 16 | lookat = np.random.uniform(-10, 10, [*spatial, 3]).astype(np.float32) 17 | up = np.random.uniform(-10, 10, [*spatial, 3]).astype(np.float32) 18 | 19 | expected = [] 20 | for i in range(np.prod(spatial) if len(spatial) > 0 else 1): 21 | expected.append(np.array(glm.lookAt( 22 | glm.vec3(eye.reshape([-1, 3])[i]), 23 | glm.vec3(lookat.reshape([-1, 3])[i]), 24 | glm.vec3(up.reshape([-1, 3])[i]) 25 | ))) 26 | expected = np.concatenate(expected, axis=0).reshape([*spatial, 4, 4]) 27 | 28 | actual = utils3d.numpy.view_look_at(eye, lookat, up) 29 | 30 | assert np.allclose(expected, actual, 1e-5, 1e-5), '\n' + \ 31 | 'Input:\n' + \ 32 | f'eye: {eye}\n' + \ 33 | f'lookat: {lookat}\n' + \ 34 | f'up: {up}\n' + \ 35 | 'Actual:\n' + \ 36 | f'{actual}\n' + \ 37 | 'Expected:\n' + \ 38 | f'{expected}' 39 | -------------------------------------------------------------------------------- /test/numpy_/transforms/view_to_extrinsic.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import glm 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | else: 13 | dim = np.random.randint(4) 14 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 15 | eye = np.random.uniform(-10, 10, [*spatial, 3]).astype(np.float32) 16 | lookat = np.random.uniform(-10, 10, [*spatial, 3]).astype(np.float32) 17 | up = np.random.uniform(-10, 10, [*spatial, 3]).astype(np.float32) 18 | 19 | expected = utils3d.numpy.extrinsics_look_at(eye, lookat, up) 20 | 21 | actual = utils3d.numpy.view_to_extrinsics(utils3d.numpy.view_look_at(eye, lookat, up)) 22 | 23 | assert np.allclose(expected, actual, 1e-5, 1e-5), '\n' + \ 24 | 'Input:\n' + \ 25 | f'eye: {eye}\n' + \ 26 | f'lookat: {lookat}\n' + \ 27 | f'up: {up}\n' + \ 28 | 'Actual:\n' + \ 29 | f'{actual}\n' + \ 30 | 'Expected:\n' + \ 31 | f'{expected}' 32 | -------------------------------------------------------------------------------- /test/numpy_/utils/image_mesh.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | 7 | def run(): 8 | args = [ 9 | {'W':2, 'H':2, 'backslash': np.array([False])}, 10 | {'W':2, 'H':2, 'backslash': np.array([True])}, 11 | {'H':2, 'W':3, 'backslash': np.array([True, False])}, 12 | ] 13 | 14 | expected = [ 15 | np.array([[0, 2, 1], [1, 2, 3]]), 16 | np.array([[0, 2, 3], [0, 3, 1]]), 17 | np.array([[0, 3, 4], [0, 4, 1], [1, 4, 2], [2, 4, 5]]), 18 | ] 19 | 20 | for args, expected in zip(args, expected): 21 | actual = utils3d.numpy.triangulate( 22 | utils3d.numpy.image_mesh(args['H'], args['W'])[1], 23 | backslash=args.get('backslash', None), 24 | ) 25 | 26 | assert np.allclose(expected, actual), '\n' + \ 27 | 'Input:\n' + \ 28 | f'{args}\n' + \ 29 | 'Actual:\n' + \ 30 | f'{actual}\n' + \ 31 | 'Expected:\n' + \ 32 | f'{expected}' 33 | -------------------------------------------------------------------------------- /test/rasterization_/gl/basic.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import moderngl 6 | import numpy as np 7 | from PIL import Image 8 | from pyrr import Matrix44 9 | 10 | # ------------------- 11 | # CREATE CONTEXT HERE 12 | # ------------------- 13 | 14 | import moderngl 15 | 16 | def run(): 17 | ctx = moderngl.create_context( 18 | standalone=True, 19 | backend='egl', 20 | # These are OPTIONAL if you want to load a specific version 21 | libgl='libGL.so.1', 22 | libegl='libEGL.so.1', 23 | ) 24 | 25 | prog = ctx.program(vertex_shader=""" 26 | #version 330 27 | uniform mat4 model; 28 | in vec2 in_vert; 29 | in vec3 in_color; 30 | out vec3 color; 31 | void main() { 32 | gl_Position = model * vec4(in_vert, 0.0, 1.0); 33 | color = in_color; 34 | } 35 | """, 36 | fragment_shader=""" 37 | #version 330 38 | in vec3 color; 39 | out vec4 fragColor; 40 | void main() { 41 | fragColor = vec4(color, 1.0); 42 | } 43 | """) 44 | 45 | vertices = np.array([ 46 | -0.6, -0.6, 47 | 1.0, 0.0, 0.0, 48 | 0.6, -0.6, 49 | 0.0, 1.0, 0.0, 50 | 0.0, 0.6, 51 | 0.0, 0.0, 1.0, 52 | ], dtype='f4') 53 | 54 | vbo = ctx.buffer(vertices) 55 | vao = ctx.simple_vertex_array(prog, vbo, 'in_vert', 'in_color') 56 | fbo = ctx.framebuffer(color_attachments=[ctx.texture((512, 512), 4)]) 57 | 58 | fbo.use() 59 | ctx.clear() 60 | prog['model'].write(Matrix44.from_eulers((0.0, 0.1, 0.0), dtype='f4')) 61 | vao.render(moderngl.TRIANGLES) 62 | 63 | data = fbo.read(components=3) 64 | image = Image.frombytes('RGB', fbo.size, data) 65 | image = image.transpose(Image.FLIP_TOP_BOTTOM) 66 | image.save(os.path.join(os.path.dirname(__file__), '..', '..', 'results_to_check', 'output.png')) 67 | 68 | 69 | if __name__ == '__main__': 70 | run() 71 | -------------------------------------------------------------------------------- /test/rasterization_/gl/rasterize_uv.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import imageio 7 | 8 | def run(): 9 | image_uv, image_mesh = utils3d.numpy.utils.image_mesh(128, 128) 10 | image_mesh = image_mesh.reshape(-1, 4) 11 | depth = np.ones((128, 128), dtype=np.float32) * 2 12 | depth[32:96, 32:96] = 1 13 | depth = depth.reshape(-1) 14 | intrinsics = utils3d.numpy.transforms.intrinsics_from_fov(1.0, 128, 128).astype(np.float32) 15 | intrinsics = utils3d.numpy.transforms.normalize_intrinsics(intrinsics, 128, 128) 16 | extrinsics = utils3d.numpy.transforms.extrinsics_look_at([0, 0, 1], [0, 0, 0], [0, 1, 0]).astype(np.float32) 17 | pts = utils3d.numpy.transforms.unproject_cv(image_uv, depth, extrinsics, intrinsics) 18 | pts = pts.reshape(-1, 3) 19 | image_mesh = utils3d.numpy.mesh.triangulate(image_mesh, vertices=pts) 20 | 21 | perspective = utils3d.numpy.transforms.perspective(1.0, 1.0, 0.1, 10) 22 | view = utils3d.numpy.transforms.view_look_at([1, 0, 1], [0, 0, 0], [0, 1, 0]) 23 | mvp = np.matmul(perspective, view) 24 | ctx = utils3d.numpy.rasterization.RastContext( 25 | standalone=True, 26 | backend='egl', 27 | device_index=0, 28 | ) 29 | uv = utils3d.numpy.rasterization.rasterize_triangle_faces( 30 | ctx, 31 | pts, 32 | image_mesh, 33 | image_uv, 34 | width=128, 35 | height=128, 36 | mvp=mvp, 37 | )[0] 38 | uv = (np.concatenate([uv, np.zeros((128, 128, 1), dtype=np.float32)], axis=-1) * 255).astype(np.uint8) 39 | imageio.imwrite(os.path.join(os.path.dirname(__file__), '..', '..', 'results_to_check', 'rasterize_uv.png'), uv) 40 | 41 | if __name__ == '__main__': 42 | run() 43 | -------------------------------------------------------------------------------- /test/test.py: -------------------------------------------------------------------------------- 1 | import importlib 2 | import os 3 | import torch 4 | import traceback 5 | 6 | CRED = '\033[91m' 7 | CGREEN = '\033[92m' 8 | CEND = '\033[0m' 9 | 10 | if __name__ == '__main__': 11 | # list all tests 12 | tests = [] 13 | for root, dirs, files in os.walk('test'): 14 | if root == 'test': 15 | continue 16 | for file in files: 17 | if file.endswith('.py'): 18 | root = root.replace('test/', '').replace('test\\', '') 19 | test = os.path.join(root, file) 20 | test = test.replace('/', '.').replace('\\', '.').replace('.py', '') 21 | tests.append(test) 22 | tests.sort() 23 | print(f'Found {len(tests)} tests:') 24 | for test in tests: 25 | print(f' {test}') 26 | print() 27 | 28 | # disable torch optimizations 29 | torch.backends.cudnn.enabled = False 30 | torch.backends.cuda.matmul.allow_tf32 = False 31 | 32 | # import and run 33 | passed = 0 34 | for test in tests: 35 | print(f'Running test: {test}... ', end='') 36 | test = importlib.import_module(test, '.'.join(test.split('.')[:-1])) 37 | try: 38 | test.run() 39 | except Exception as e: 40 | print(CRED, end='') 41 | print('Failed') 42 | traceback.print_exc() 43 | else: 44 | print(CGREEN, end='') 45 | print('Passed') 46 | passed += 1 47 | print(CEND, end='') 48 | 49 | print(f'Passed {passed}/{len(tests)} tests') 50 | -------------------------------------------------------------------------------- /test/torch_/mesh/compute_face_angle.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import torch 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | vertices = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0]], dtype=float) 13 | faces = np.array([[0, 1, 2]]) 14 | expected = np.array([[np.pi/2, np.pi/4, np.pi/4]]) 15 | else: 16 | dim = np.random.randint(4) 17 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 18 | N = np.random.randint(100, 1000) 19 | vertices = np.random.rand(*spatial, N, 3) 20 | L = np.random.randint(1, 1000) 21 | faces = np.random.randint(0, N, size=(*spatial, L, 3)) 22 | faces[..., 1] = (faces[..., 0] + 1) % N 23 | faces[..., 2] = (faces[..., 0] + 2) % N 24 | 25 | expected = utils3d.numpy.compute_face_angle(vertices, faces) 26 | 27 | device = [torch.device('cpu'), torch.device('cuda')][np.random.randint(2)] 28 | vertices = torch.tensor(vertices, device=device) 29 | faces = torch.tensor(faces, device=device) 30 | 31 | actual = utils3d.torch.compute_face_angle(vertices, faces).cpu().numpy() 32 | 33 | assert np.allclose(expected, actual), '\n' + \ 34 | 'Input:\n' + \ 35 | f'{faces}\n' + \ 36 | 'Actual:\n' + \ 37 | f'{actual}\n' + \ 38 | 'Expected:\n' + \ 39 | f'{expected}' 40 | -------------------------------------------------------------------------------- /test/torch_/mesh/compute_face_normal.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import torch 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | vertices = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0]], dtype=float) 13 | faces = np.array([[0, 1, 2]]) 14 | expected = np.array([[0, 0, 1]]) 15 | else: 16 | dim = np.random.randint(4) 17 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 18 | N = np.random.randint(100, 1000) 19 | vertices = np.random.rand(*spatial, N, 3) 20 | L = np.random.randint(1, 1000) 21 | faces = np.random.randint(0, N, size=(*spatial, L, 3)) 22 | faces[..., 1] = (faces[..., 0] + 1) % N 23 | faces[..., 2] = (faces[..., 0] + 2) % N 24 | 25 | expected = utils3d.numpy.compute_face_normal(vertices, faces) 26 | 27 | device = [torch.device('cpu'), torch.device('cuda')][np.random.randint(2)] 28 | vertices = torch.tensor(vertices, device=device) 29 | faces = torch.tensor(faces, device=device) 30 | 31 | actual = utils3d.torch.compute_face_normal(vertices, faces).cpu().numpy() 32 | 33 | assert np.allclose(expected, actual), '\n' + \ 34 | 'Input:\n' + \ 35 | f'{faces}\n' + \ 36 | 'Actual:\n' + \ 37 | f'{actual}\n' + \ 38 | 'Expected:\n' + \ 39 | f'{expected}' 40 | -------------------------------------------------------------------------------- /test/torch_/mesh/compute_vertex_normal.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import torch 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | vertices = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0]], dtype=float) 13 | faces = np.array([[0, 1, 2]]) 14 | expected = np.array([[0, 0, 1], [0, 0, 1], [0, 0, 1]]) 15 | else: 16 | dim = np.random.randint(4) 17 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 18 | N = np.random.randint(100, 1000) 19 | vertices = np.random.rand(*spatial, N, 3) 20 | L = np.random.randint(1, 1000) 21 | faces = np.random.randint(0, N, size=(*spatial, L, 3)) 22 | faces[..., 1] = (faces[..., 0] + 1) % N 23 | faces[..., 2] = (faces[..., 0] + 2) % N 24 | 25 | expected = utils3d.numpy.compute_vertex_normal(vertices, faces) 26 | 27 | device = [torch.device('cpu'), torch.device('cuda')][np.random.randint(2)] 28 | vertices = torch.tensor(vertices, device=device) 29 | faces = torch.tensor(faces, device=device) 30 | 31 | actual = utils3d.torch.compute_vertex_normal(vertices, faces).cpu().numpy() 32 | 33 | assert np.allclose(expected, actual), '\n' + \ 34 | 'Input:\n' + \ 35 | f'{faces}\n' + \ 36 | 'Actual:\n' + \ 37 | f'{actual}\n' + \ 38 | 'Expected:\n' + \ 39 | f'{expected}' 40 | -------------------------------------------------------------------------------- /test/torch_/mesh/compute_vertex_normal_weighted.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import torch 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | vertices = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0]], dtype=float) 13 | faces = np.array([[0, 1, 2]]) 14 | expected = np.array([[0, 0, 1], [0, 0, 1], [0, 0, 1]]) 15 | else: 16 | dim = np.random.randint(4) 17 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 18 | N = np.random.randint(100, 1000) 19 | vertices = np.random.rand(*spatial, N, 3) 20 | L = np.random.randint(1, 1000) 21 | faces = np.random.randint(0, N, size=(*spatial, L, 3)) 22 | faces[..., 1] = (faces[..., 0] + 1) % N 23 | faces[..., 2] = (faces[..., 0] + 2) % N 24 | 25 | expected = utils3d.numpy.compute_vertex_normal_weighted(vertices, faces) 26 | 27 | device = [torch.device('cpu'), torch.device('cuda')][np.random.randint(2)] 28 | vertices = torch.tensor(vertices, device=device) 29 | faces = torch.tensor(faces, device=device) 30 | 31 | actual = utils3d.torch.compute_vertex_normal_weighted(vertices, faces).cpu().numpy() 32 | 33 | assert np.allclose(expected, actual), '\n' + \ 34 | 'Input:\n' + \ 35 | f'{faces}\n' + \ 36 | 'Actual:\n' + \ 37 | f'{actual}\n' + \ 38 | 'Expected:\n' + \ 39 | f'{expected}' 40 | -------------------------------------------------------------------------------- /test/torch_/mesh/merge_duplicate_vertices.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import torch 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | vertices = np.array([[0, 0, 0], [1, 0, 0], [1, 0, 0]], dtype=float) 13 | faces = np.array([[0, 1, 2]]) 14 | expected_vertices = np.array([[0, 0, 0], [1, 0, 0]]) 15 | expected_faces = np.array([[0, 1, 1]]) 16 | expected = expected_vertices[expected_faces] 17 | else: 18 | N = np.random.randint(100, 1000) 19 | vertices = np.random.rand(N, 3) 20 | L = np.random.randint(1, 1000) 21 | faces = np.random.randint(0, N, size=(L, 3)) 22 | faces[..., 1] = (faces[..., 0] + 1) % N 23 | faces[..., 2] = (faces[..., 0] + 2) % N 24 | vertices[-(N//2):] = vertices[:N//2] 25 | 26 | expected_vertices, expected_faces = utils3d.numpy.merge_duplicate_vertices(vertices, faces) 27 | expected = expected_vertices[expected_faces] 28 | 29 | device = [torch.device('cpu'), torch.device('cuda')][np.random.randint(2)] 30 | vertices = torch.tensor(vertices, device=device) 31 | faces = torch.tensor(faces, device=device) 32 | 33 | actual_vertices, actual_faces = utils3d.torch.merge_duplicate_vertices(vertices, faces) 34 | actual_vertices = actual_vertices.cpu().numpy() 35 | actual_faces = actual_faces.cpu().numpy() 36 | actual = actual_vertices[actual_faces] 37 | 38 | assert expected_vertices.shape == actual_vertices.shape and np.allclose(expected, actual), '\n' + \ 39 | 'Input:\n' + \ 40 | f'{faces}\n' + \ 41 | 'Actual:\n' + \ 42 | f'{actual}\n' + \ 43 | 'Expected:\n' + \ 44 | f'{expected}' 45 | -------------------------------------------------------------------------------- /test/torch_/mesh/remove_corrupted_faces.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import torch 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | faces = np.array([[0, 1, 2], [0, 2, 2], [0, 2, 3]]) 12 | expected = np.array([[0, 1, 2], [0, 2, 3]]) 13 | else: 14 | L = np.random.randint(1, 1000) 15 | N = np.random.randint(100, 1000) 16 | faces = np.random.randint(0, N, size=(L, 3)) 17 | faces[..., 1] = (faces[..., 0] + 1) % N 18 | faces[..., 2] = (faces[..., 0] + 2) % N 19 | 20 | expected = utils3d.numpy.remove_corrupted_faces(faces) 21 | 22 | device = [torch.device('cpu'), torch.device('cuda')][np.random.randint(2)] 23 | faces = torch.tensor(faces, device=device) 24 | 25 | actual = utils3d.torch.remove_corrupted_faces(faces).cpu().numpy() 26 | 27 | assert np.allclose(expected, actual), '\n' + \ 28 | 'Input:\n' + \ 29 | f'{faces}\n' + \ 30 | 'Actual:\n' + \ 31 | f'{actual}\n' + \ 32 | 'Expected:\n' + \ 33 | f'{expected}' 34 | -------------------------------------------------------------------------------- /test/torch_/mesh/triangulate.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import torch 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | L = 1 13 | N = 5 14 | faces = np.array([[0, 1, 2, 3, 4]]) 15 | expected = np.array([[0, 1, 2], [0, 2, 3], [0, 3, 4]]) 16 | else: 17 | dim = np.random.randint(4) 18 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 19 | L = np.random.randint(1, 1000) 20 | N = np.random.randint(3, 10) 21 | faces = np.random.randint(0, 10000, size=(*spatial, L, N)) 22 | 23 | expected = utils3d.numpy.triangulate(faces) 24 | 25 | device = [torch.device('cpu'), torch.device('cuda')][np.random.randint(2)] 26 | faces = torch.tensor(faces, device=device) 27 | 28 | actual = utils3d.torch.triangulate(faces).cpu().numpy() 29 | 30 | assert np.allclose(expected, actual), '\n' + \ 31 | 'Input:\n' + \ 32 | f'{faces}\n' + \ 33 | 'Actual:\n' + \ 34 | f'{actual}\n' + \ 35 | 'Expected:\n' + \ 36 | f'{expected}' 37 | -------------------------------------------------------------------------------- /test/torch_/rasterization/warp_image_by_depth.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import torch 7 | import imageio 8 | 9 | def run(): 10 | depth = torch.ones((1, 128, 128), dtype=torch.float32, device='cuda') * 2 11 | depth[:, 32:48, 32:48] = 1 12 | intrinsics = utils3d.torch.transforms.intrinsics(1.0, 1.0, 0.5, 0.5).to(depth) 13 | extrinsics_src = utils3d.torch.transforms.extrinsics_look_at([0., 0., 1.], [0., 0., 0.], [0., 1., 0.]).to(depth) 14 | extrinsics_tgt = utils3d.torch.transforms.extrinsics_look_at([1., 0., 1.], [0., 0., 0.], [0., 1., 0.]).to(depth) 15 | ctx = utils3d.torch.rasterization.RastContext(backend='gl', device='cuda') 16 | uv, _ = utils3d.torch.rasterization.warp_image_by_depth( 17 | ctx, 18 | depth, 19 | extrinsics_src=extrinsics_src, 20 | extrinsics_tgt=extrinsics_tgt, 21 | intrinsics_src=intrinsics, 22 | antialiasing=False, 23 | ) 24 | uv = torch.cat([uv, torch.zeros((1, 1, 128, 128)).to(uv)], dim=1) * 255 25 | uv = uv.permute(0, 2, 3, 1).squeeze().cpu().numpy().astype(np.uint8) 26 | 27 | imageio.imwrite(os.path.join(os.path.dirname(__file__), '..', '..', 'results_to_check', 'torch_warp_image_uv.png'), uv) 28 | 29 | if __name__ == '__main__': 30 | run() 31 | -------------------------------------------------------------------------------- /test/torch_/transforms/crop_intrinsic.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import torch 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | else: 13 | dim = np.random.randint(4) 14 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 15 | fov = np.random.uniform(5 / 180 * np.pi, 175 / 180 * np.pi, spatial) 16 | width = np.random.uniform(1, 10000, spatial) 17 | height = np.random.uniform(1, 10000, spatial) 18 | left = np.random.uniform(0, width, spatial) 19 | top = np.random.uniform(0, height, spatial) 20 | crop_width = np.random.uniform(0, width - left, spatial) 21 | crop_height = np.random.uniform(0, height - top, spatial) 22 | 23 | expected = utils3d.numpy.crop_intrinsics( 24 | utils3d.numpy.normalize_intrinsics( 25 | utils3d.numpy.intrinsics_from_fov(fov, width, height), 26 | width, height 27 | ), 28 | width, height, left, top, crop_width, crop_height 29 | ) 30 | 31 | device = [torch.device('cpu'), torch.device('cuda')][np.random.randint(2)] 32 | fov = torch.tensor(fov, device=device) 33 | width = torch.tensor(width, device=device) 34 | height = torch.tensor(height, device=device) 35 | left = torch.tensor(left, device=device) 36 | top = torch.tensor(top, device=device) 37 | crop_width = torch.tensor(crop_width, device=device) 38 | crop_height = torch.tensor(crop_height, device=device) 39 | 40 | actual = utils3d.torch.crop_intrinsics( 41 | utils3d.torch.normalize_intrinsics( 42 | utils3d.torch.intrinsics_from_fov(fov, width, height), 43 | width, height 44 | ), 45 | width, height, left, top, crop_width, crop_height 46 | ).cpu().numpy() 47 | 48 | assert np.allclose(expected, actual), '\n' + \ 49 | 'Input:\n' + \ 50 | f'\tfov: {fov}\n' + \ 51 | f'\twidth: {width}\n' + \ 52 | f'\theight: {height}\n' + \ 53 | f'\tleft: {left}\n' + \ 54 | f'\ttop: {top}\n' + \ 55 | f'\tcrop_width: {crop_width}\n' + \ 56 | f'\tcrop_height: {crop_height}\n' + \ 57 | 'Actual:\n' + \ 58 | f'{actual}\n' + \ 59 | 'Expected:\n' + \ 60 | f'{expected}' 61 | -------------------------------------------------------------------------------- /test/torch_/transforms/extrinsic_look_at.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import torch 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | else: 13 | dim = np.random.randint(4) 14 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 15 | eye = np.random.uniform(-10, 10, [*spatial, 3]).astype(np.float32) 16 | lookat = np.random.uniform(-10, 10, [*spatial, 3]).astype(np.float32) 17 | up = np.random.uniform(-10, 10, [*spatial, 3]).astype(np.float32) 18 | 19 | expected = utils3d.numpy.extrinsics_look_at(eye, lookat, up) 20 | 21 | device = [torch.device('cpu'), torch.device('cuda')][np.random.randint(2)] 22 | eye = torch.tensor(eye, device=device) 23 | lookat = torch.tensor(lookat, device=device) 24 | up = torch.tensor(up, device=device) 25 | 26 | actual = utils3d.torch.extrinsics_look_at(eye, lookat, up).cpu().numpy() 27 | 28 | assert np.allclose(expected, actual, 1e-5, 1e-5), '\n' + \ 29 | 'Input:\n' + \ 30 | f'eye: {eye}\n' + \ 31 | f'lookat: {lookat}\n' + \ 32 | f'up: {up}\n' + \ 33 | 'Actual:\n' + \ 34 | f'{actual}\n' + \ 35 | 'Expected:\n' + \ 36 | f'{expected}' 37 | -------------------------------------------------------------------------------- /test/torch_/transforms/extrinsic_to_view.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import torch 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | else: 13 | dim = np.random.randint(4) 14 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 15 | eye = np.random.uniform(-10, 10, [*spatial, 3]).astype(np.float32) 16 | lookat = np.random.uniform(-10, 10, [*spatial, 3]).astype(np.float32) 17 | up = np.random.uniform(-10, 10, [*spatial, 3]).astype(np.float32) 18 | 19 | expected = utils3d.numpy.view_to_extrinsics(utils3d.numpy.extrinsics_look_at(eye, lookat, up)) 20 | 21 | device = [torch.device('cpu'), torch.device('cuda')][np.random.randint(2)] 22 | eye = torch.tensor(eye, device=device) 23 | lookat = torch.tensor(lookat, device=device) 24 | up = torch.tensor(up, device=device) 25 | 26 | actual = utils3d.torch.view_to_extrinsics(utils3d.torch.extrinsics_look_at(eye, lookat, up)).cpu().numpy() 27 | 28 | assert np.allclose(expected, actual, 1e-5, 1e-5), '\n' + \ 29 | 'Input:\n' + \ 30 | f'eye: {eye}\n' + \ 31 | f'lookat: {lookat}\n' + \ 32 | f'up: {up}\n' + \ 33 | 'Actual:\n' + \ 34 | f'{actual}\n' + \ 35 | 'Expected:\n' + \ 36 | f'{expected}' 37 | -------------------------------------------------------------------------------- /test/torch_/transforms/intrinsic.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import torch 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | else: 13 | dim = np.random.randint(4) 14 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 15 | focal_x = np.random.uniform(1, 10000, spatial) 16 | focal_y = np.random.uniform(1, 10000, spatial) 17 | center_x = np.random.uniform(1, 10000, spatial) 18 | center_y = np.random.uniform(1, 10000, spatial) 19 | 20 | expected = utils3d.numpy.intrinsics(focal_x, focal_y, center_x, center_y) 21 | 22 | device = [torch.device('cpu'), torch.device('cuda')][np.random.randint(2)] 23 | focal_x = torch.tensor(focal_x, device=device) 24 | focal_y = torch.tensor(focal_y, device=device) 25 | center_x = torch.tensor(center_x, device=device) 26 | center_y = torch.tensor(center_y, device=device) 27 | 28 | actual = utils3d.torch.intrinsics(focal_x, focal_y, center_x, center_y).cpu().numpy() 29 | 30 | assert np.allclose(expected, actual), '\n' + \ 31 | 'Input:\n' + \ 32 | f'\tfocal_x: {focal_x}\n' + \ 33 | f'\tfocal_y: {focal_y}\n' + \ 34 | f'\tcenter_x: {center_x}\n' + \ 35 | f'\tcenter_y: {center_y}\n' + \ 36 | 'Actual:\n' + \ 37 | f'{actual}\n' + \ 38 | 'Expected:\n' + \ 39 | f'{expected}' 40 | -------------------------------------------------------------------------------- /test/torch_/transforms/intrinsic_from_fov.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import torch 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | else: 13 | dim = np.random.randint(4) 14 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 15 | fov = np.random.uniform(5 / 180 * np.pi, 175 / 180 * np.pi, spatial) 16 | width = np.random.uniform(1, 10000, spatial) 17 | height = np.random.uniform(1, 10000, spatial) 18 | 19 | expected = utils3d.numpy.intrinsics_from_fov(fov, width, height) 20 | 21 | device = [torch.device('cpu'), torch.device('cuda')][np.random.randint(2)] 22 | fov = torch.tensor(fov, device=device) 23 | width = torch.tensor(width, device=device) 24 | height = torch.tensor(height, device=device) 25 | 26 | actual = utils3d.torch.intrinsics_from_fov(fov, width, height).cpu().numpy() 27 | 28 | assert np.allclose(expected, actual), '\n' + \ 29 | 'Input:\n' + \ 30 | f'\tfov: {fov}\n' + \ 31 | f'\twidth: {width}\n' + \ 32 | f'\theight: {height}\n' + \ 33 | 'Actual:\n' + \ 34 | f'{actual}\n' + \ 35 | 'Expected:\n' + \ 36 | f'{expected}' 37 | -------------------------------------------------------------------------------- /test/torch_/transforms/intrinsic_from_fov_xy.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import torch 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | else: 13 | dim = np.random.randint(4) 14 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 15 | fov_x = np.random.uniform(5 / 180 * np.pi, 175 / 180 * np.pi, spatial) 16 | fov_y = np.random.uniform(5 / 180 * np.pi, 175 / 180 * np.pi, spatial) 17 | 18 | expected = utils3d.numpy.intrinsics_from_fov_xy(fov_x, fov_y) 19 | 20 | device = [torch.device('cpu'), torch.device('cuda')][np.random.randint(2)] 21 | fov_x = torch.tensor(fov_x, device=device) 22 | fov_y = torch.tensor(fov_y, device=device) 23 | 24 | actual = utils3d.torch.intrinsics_from_fov_xy(fov_x, fov_y).cpu().numpy() 25 | 26 | assert np.allclose(expected, actual), '\n' + \ 27 | 'Input:\n' + \ 28 | f'\tfov_x: {fov_x}\n' + \ 29 | f'\tfov_y: {fov_y}\n' + \ 30 | 'Actual:\n' + \ 31 | f'{actual}\n' + \ 32 | 'Expected:\n' + \ 33 | f'{expected}' 34 | -------------------------------------------------------------------------------- /test/torch_/transforms/intrinsic_to_perspective.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import torch 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | else: 13 | dim = np.random.randint(4) 14 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 15 | fov_x = np.random.uniform(5 / 180 * np.pi, 175 / 180 * np.pi, spatial) 16 | fov_y = np.random.uniform(5 / 180 * np.pi, 175 / 180 * np.pi, spatial) 17 | near = np.random.uniform(0.1, 100, spatial) 18 | far = np.random.uniform(near*2, 1000, spatial) 19 | 20 | expected = utils3d.numpy.intrinsics_to_perspective( 21 | utils3d.numpy.intrinsics_from_fov_xy(fov_x, fov_y), 22 | near, 23 | far 24 | ) 25 | 26 | device = [torch.device('cpu'), torch.device('cuda')][np.random.randint(2)] 27 | fov_x = torch.tensor(fov_x, device=device) 28 | fov_y = torch.tensor(fov_y, device=device) 29 | near = torch.tensor(near, device=device) 30 | far = torch.tensor(far, device=device) 31 | 32 | actual = utils3d.torch.intrinsics_to_perspective( 33 | utils3d.torch.intrinsics_from_fov_xy(fov_x, fov_y), 34 | near, 35 | far 36 | ).cpu().numpy() 37 | 38 | assert np.allclose(expected, actual), '\n' + \ 39 | 'Input:\n' + \ 40 | f'\tfov_x: {fov_x}\n' + \ 41 | f'\tfov_y: {fov_y}\n' + \ 42 | f'\tnear: {near}\n' + \ 43 | f'\tfar: {far}\n' + \ 44 | 'Actual:\n' + \ 45 | f'{actual}\n' + \ 46 | 'Expected:\n' + \ 47 | f'{expected}' 48 | -------------------------------------------------------------------------------- /test/torch_/transforms/linearize_depth.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import torch 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | else: 13 | dim = np.random.randint(4) 14 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 15 | near = np.random.uniform(0.1, 100, spatial) 16 | far = np.random.uniform(near*2, 1000, spatial) 17 | depth = np.random.uniform(near, far, spatial) 18 | 19 | expected = depth 20 | 21 | device = [torch.device('cpu'), torch.device('cuda')][np.random.randint(2)] 22 | near = torch.tensor(near, device=device) 23 | far = torch.tensor(far, device=device) 24 | depth = torch.tensor(depth, device=device) 25 | 26 | actual = utils3d.torch.depth_buffer_to_linear( 27 | utils3d.torch.project_depth(depth, near, far), 28 | near, far 29 | ).cpu().numpy() 30 | 31 | assert np.allclose(expected, actual), '\n' + \ 32 | 'Input:\n' + \ 33 | f'\tdepth: {depth}\n' + \ 34 | f'\tnear: {near}\n' + \ 35 | f'\tfar: {far}\n' + \ 36 | 'Actual:\n' + \ 37 | f'{actual}\n' + \ 38 | 'Expected:\n' + \ 39 | f'{expected}' 40 | -------------------------------------------------------------------------------- /test/torch_/transforms/matrix_to_quaternion.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | from scipy.spatial.transform import Rotation as R 4 | import utils3d 5 | 6 | 7 | def run(): 8 | for i in range(10): 9 | if i == 0: 10 | spatial = [] 11 | else: 12 | dim = np.random.randint(4) 13 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 14 | angle = np.random.uniform(-np.pi, np.pi, spatial) 15 | axis = np.random.uniform(-1, 1, spatial + [3]) 16 | axis = axis / np.linalg.norm(axis, axis=-1, keepdims=True) 17 | axis_angle = angle[..., None] * axis 18 | matrix = R.from_rotvec(axis_angle.reshape((-1, 3))).as_matrix().reshape(spatial + [3, 3]) 19 | # matrix = np.array([ 20 | # [1, 0, 0], 21 | # [0, 0, -1], 22 | # [0, 1, 0] 23 | # ]).astype(np.float32) 24 | # dim = 0 25 | # spatial = [] 26 | expected = R.from_matrix(matrix.reshape(-1, 3, 3)).as_quat().reshape(spatial + [4])[..., [3, 0, 1, 2]] 27 | actual = utils3d.torch.matrix_to_quaternion( 28 | torch.from_numpy(matrix) 29 | ).cpu().numpy() 30 | assert np.allclose(expected, actual), '\n' + \ 31 | 'Input:\n' + \ 32 | f'\tangle: {angle}\n' + \ 33 | f'\taxis: {axis}\n' + \ 34 | f'\tmatrix: {matrix}\n' + \ 35 | 'Actual:\n' + \ 36 | f'{actual}\n' + \ 37 | 'Expected:\n' + \ 38 | f'{expected}' 39 | 40 | if __name__ == '__main__': 41 | run() -------------------------------------------------------------------------------- /test/torch_/transforms/normalize_intrinsic.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import torch 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | else: 13 | dim = np.random.randint(4) 14 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 15 | fov = np.random.uniform(5 / 180 * np.pi, 175 / 180 * np.pi, spatial) 16 | width = np.random.uniform(1, 10000, spatial) 17 | height = np.random.uniform(1, 10000, spatial) 18 | fov_x = np.where(width >= height, fov, 2 * np.arctan(np.tan(fov / 2) * width / height)) 19 | fov_y = np.where(width >= height, 2 * np.arctan(np.tan(fov / 2) * height / width), fov) 20 | 21 | expected = utils3d.numpy.normalize_intrinsics(utils3d.numpy.intrinsics_from_fov(fov, width, height), width, height) 22 | 23 | device = [torch.device('cpu'), torch.device('cuda')][np.random.randint(2)] 24 | fov = torch.tensor(fov, device=device) 25 | width = torch.tensor(width, device=device) 26 | height = torch.tensor(height, device=device) 27 | 28 | actual = utils3d.torch.normalize_intrinsics(utils3d.torch.intrinsics_from_fov(fov, width, height), width, height).cpu().numpy() 29 | 30 | assert np.allclose(expected, actual), '\n' + \ 31 | 'Input:\n' + \ 32 | f'\tfov: {fov}\n' + \ 33 | f'\twidth: {width}\n' + \ 34 | f'\theight: {height}\n' + \ 35 | 'Actual:\n' + \ 36 | f'{actual}\n' + \ 37 | 'Expected:\n' + \ 38 | f'{expected}' 39 | -------------------------------------------------------------------------------- /test/torch_/transforms/perspective.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import torch 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | else: 13 | dim = np.random.randint(4) 14 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 15 | fovy = np.random.uniform(5 / 180 * np.pi, 175 / 180 * np.pi, spatial) 16 | aspect = np.random.uniform(0.01, 100, spatial) 17 | near = np.random.uniform(0.1, 100, spatial) 18 | far = np.random.uniform(near*2, 1000, spatial) 19 | 20 | expected = utils3d.numpy.perspective(fovy, aspect, near, far) 21 | 22 | device = [torch.device('cpu'), torch.device('cuda')][np.random.randint(2)] 23 | fovy = torch.tensor(fovy, device=device) 24 | aspect = torch.tensor(aspect, device=device) 25 | near = torch.tensor(near, device=device) 26 | far = torch.tensor(far, device=device) 27 | 28 | actual = utils3d.torch.perspective(fovy, aspect, near, far).cpu().numpy() 29 | 30 | assert np.allclose(expected, actual), '\n' + \ 31 | 'Input:\n' + \ 32 | f'\tfovy: {fovy}\n' + \ 33 | f'\taspect: {aspect}\n' + \ 34 | f'\tnear: {near}\n' + \ 35 | f'\tfar: {far}\n' + \ 36 | 'Actual:\n' + \ 37 | f'{actual}\n' + \ 38 | 'Expected:\n' + \ 39 | f'{expected}' 40 | -------------------------------------------------------------------------------- /test/torch_/transforms/perspective_from_fov.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import torch 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | else: 13 | dim = np.random.randint(4) 14 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 15 | fov = np.random.uniform(5 / 180 * np.pi, 175 / 180 * np.pi, spatial) 16 | width = np.random.uniform(1, 10000, spatial) 17 | height = np.random.uniform(1, 10000, spatial) 18 | near = np.random.uniform(0.1, 100, spatial) 19 | far = np.random.uniform(near*2, 1000, spatial) 20 | 21 | expected = utils3d.numpy.perspective_from_fov(fov, width, height, near, far) 22 | 23 | device = [torch.device('cpu'), torch.device('cuda')][np.random.randint(2)] 24 | fov = torch.tensor(fov, device=device) 25 | width = torch.tensor(width, device=device) 26 | height = torch.tensor(height, device=device) 27 | near = torch.tensor(near, device=device) 28 | far = torch.tensor(far, device=device) 29 | 30 | actual = utils3d.torch.perspective_from_fov(fov, width, height, near, far).cpu().numpy() 31 | 32 | assert np.allclose(expected, actual), '\n' + \ 33 | 'Input:\n' + \ 34 | f'\tfov: {fov}\n' + \ 35 | f'\twidth: {width}\n' + \ 36 | f'\theight: {height}\n' + \ 37 | f'\tnear: {near}\n' + \ 38 | f'\tfar: {far}\n' + \ 39 | 'Actual:\n' + \ 40 | f'{actual}\n' + \ 41 | 'Expected:\n' + \ 42 | f'{expected}' 43 | -------------------------------------------------------------------------------- /test/torch_/transforms/perspective_from_fov_xy.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import torch 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | else: 13 | dim = np.random.randint(4) 14 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 15 | fov_x = np.random.uniform(5 / 180 * np.pi, 175 / 180 * np.pi, spatial) 16 | fov_y = np.random.uniform(5 / 180 * np.pi, 175 / 180 * np.pi, spatial) 17 | near = np.random.uniform(0.1, 100, spatial) 18 | far = np.random.uniform(near*2, 1000, spatial) 19 | 20 | expected = utils3d.numpy.perspective_from_fov_xy(fov_x, fov_y, near, far) 21 | 22 | device = [torch.device('cpu'), torch.device('cuda')][np.random.randint(2)] 23 | fov_x = torch.tensor(fov_x, device=device) 24 | fov_y = torch.tensor(fov_y, device=device) 25 | near = torch.tensor(near, device=device) 26 | far = torch.tensor(far, device=device) 27 | 28 | actual = utils3d.torch.perspective_from_fov_xy(fov_x, fov_y, near, far).cpu().numpy() 29 | 30 | assert np.allclose(expected, actual), '\n' + \ 31 | 'Input:\n' + \ 32 | f'\tfov_x: {fov_x}\n' + \ 33 | f'\tfov_y: {fov_y}\n' + \ 34 | f'\tnear: {near}\n' + \ 35 | f'\tfar: {far}\n' + \ 36 | 'Actual:\n' + \ 37 | f'{actual}\n' + \ 38 | 'Expected:\n' + \ 39 | f'{expected}' 40 | -------------------------------------------------------------------------------- /test/torch_/transforms/perspective_to_intrinsic.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import torch 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | else: 13 | dim = np.random.randint(4) 14 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 15 | fov_x = np.random.uniform(5 / 180 * np.pi, 175 / 180 * np.pi, spatial) 16 | fov_y = np.random.uniform(5 / 180 * np.pi, 175 / 180 * np.pi, spatial) 17 | near = np.random.uniform(0.1, 100, spatial) 18 | far = np.random.uniform(near*2, 1000, spatial) 19 | 20 | expected = utils3d.numpy.perspective_to_intrinsics( 21 | utils3d.numpy.perspective_from_fov_xy(fov_x, fov_y, near, far) 22 | ) 23 | 24 | device = [torch.device('cpu'), torch.device('cuda')][np.random.randint(2)] 25 | fov_x = torch.tensor(fov_x, device=device) 26 | fov_y = torch.tensor(fov_y, device=device) 27 | near = torch.tensor(near, device=device) 28 | far = torch.tensor(far, device=device) 29 | 30 | actual = utils3d.torch.perspective_to_intrinsics( 31 | utils3d.torch.perspective_from_fov_xy(fov_x, fov_y, near, far) 32 | ).cpu().numpy() 33 | 34 | assert np.allclose(expected, actual), '\n' + \ 35 | 'Input:\n' + \ 36 | f'\tfov_x: {fov_x}\n' + \ 37 | f'\tfov_y: {fov_y}\n' + \ 38 | f'\tnear: {near}\n' + \ 39 | f'\tfar: {far}\n' + \ 40 | 'Actual:\n' + \ 41 | f'{actual}\n' + \ 42 | 'Expected:\n' + \ 43 | f'{expected}' 44 | -------------------------------------------------------------------------------- /test/torch_/transforms/pixel_to_ndc.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import torch 7 | 8 | def run(): 9 | for i in range(100): 10 | H = np.random.randint(1, 1000) 11 | W = np.random.randint(1, 1000) 12 | x, y = np.meshgrid(np.arange(W), np.arange(H), indexing='xy') 13 | pixel = np.stack([x, y], axis=-1) 14 | 15 | expected = utils3d.numpy.pixel_to_ndc(pixel, W, H) 16 | 17 | device = [torch.device('cpu'), torch.device('cuda')][np.random.randint(2)] 18 | pixel = torch.tensor(pixel, device=device) 19 | W = torch.tensor(W, device=device) 20 | H = torch.tensor(H, device=device) 21 | 22 | actual = utils3d.torch.pixel_to_ndc(pixel, W, H).cpu().numpy() 23 | 24 | assert np.allclose(expected, actual, atol=1e-6), '\n' + \ 25 | 'Input:\n' + \ 26 | f'\tH: {H}\n' + \ 27 | f'\tW: {W}\n' + \ 28 | 'Actual:\n' + \ 29 | f'{actual}\n' + \ 30 | 'Expected:\n' + \ 31 | f'{expected}' 32 | -------------------------------------------------------------------------------- /test/torch_/transforms/pixel_to_uv.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import torch 7 | 8 | def run(): 9 | for i in range(100): 10 | H = np.random.randint(1, 1000) 11 | W = np.random.randint(1, 1000) 12 | x, y = np.meshgrid(np.arange(W), np.arange(H), indexing='xy') 13 | pixel = np.stack([x, y], axis=-1) 14 | 15 | expected = utils3d.numpy.pixel_to_uv(pixel, W, H) 16 | 17 | device = [torch.device('cpu'), torch.device('cuda')][np.random.randint(2)] 18 | pixel = torch.tensor(pixel, device=device) 19 | W = torch.tensor(W, device=device) 20 | H = torch.tensor(H, device=device) 21 | 22 | actual = utils3d.torch.pixel_to_uv(pixel, W, H).cpu().numpy() 23 | 24 | assert np.allclose(expected, actual), '\n' + \ 25 | 'Input:\n' + \ 26 | f'\tH: {H}\n' + \ 27 | f'\tW: {W}\n' + \ 28 | 'Actual:\n' + \ 29 | f'{actual}\n' + \ 30 | 'Expected:\n' + \ 31 | f'{expected}' 32 | -------------------------------------------------------------------------------- /test/torch_/transforms/project_cv.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import torch 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | N = 1 13 | else: 14 | dim = np.random.randint(4) 15 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 16 | N = np.random.randint(1, 10) 17 | focal_x = np.random.uniform(0, 10, spatial) 18 | focal_y = np.random.uniform(0, 10, spatial) 19 | center_x = np.random.uniform(0, 1, spatial) 20 | center_y = np.random.uniform(0, 1, spatial) 21 | eye = np.random.uniform(-10, 10, [*spatial, 3]) 22 | lookat = np.random.uniform(-10, 10, [*spatial, 3]) 23 | up = np.random.uniform(-10, 10, [*spatial, 3]) 24 | points = np.random.uniform(-10, 10, [*spatial, N, 3]) 25 | 26 | expected, _ = utils3d.numpy.transforms.project_cv(points, 27 | utils3d.numpy.extrinsics_look_at(eye, lookat, up), 28 | utils3d.numpy.intrinsics(focal_x, focal_y, center_x, center_y)) 29 | 30 | device = [torch.device('cpu'), torch.device('cuda')][np.random.randint(2)] 31 | focal_x = torch.tensor(focal_x, device=device) 32 | focal_y = torch.tensor(focal_y, device=device) 33 | center_x = torch.tensor(center_x, device=device) 34 | center_y = torch.tensor(center_y, device=device) 35 | eye = torch.tensor(eye, device=device) 36 | lookat = torch.tensor(lookat, device=device) 37 | up = torch.tensor(up, device=device) 38 | points = torch.tensor(points, device=device) 39 | 40 | actual, _ = utils3d.torch.project_cv(points, 41 | utils3d.torch.extrinsics_look_at(eye, lookat, up), 42 | utils3d.torch.intrinsics(focal_x, focal_y, center_x, center_y) 43 | ) 44 | actual = actual.cpu().numpy() 45 | 46 | assert np.allclose(expected, actual), '\n' + \ 47 | 'Input:\n' + \ 48 | f'\tfocal_x: {focal_x}\n' + \ 49 | f'\tfocal_y: {focal_y}\n' + \ 50 | f'\tcenter_x: {center_x}\n' + \ 51 | f'\tcenter_y: {center_y}\n' + \ 52 | f'\teye: {eye}\n' + \ 53 | f'\tlookat: {lookat}\n' + \ 54 | f'\tup: {up}\n' + \ 55 | f'\tpoints: {points}\n' + \ 56 | 'Actual:\n' + \ 57 | f'{actual}\n' + \ 58 | 'Expected:\n' + \ 59 | f'{expected}' 60 | -------------------------------------------------------------------------------- /test/torch_/transforms/project_depth.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import torch 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | else: 13 | dim = np.random.randint(4) 14 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 15 | near = np.random.uniform(0.1, 100, spatial) 16 | far = np.random.uniform(near*2, 1000, spatial) 17 | depth = np.random.uniform(near, far, spatial) 18 | 19 | expected = utils3d.numpy.project_depth(depth, near, far) 20 | 21 | device = [torch.device('cpu'), torch.device('cuda')][np.random.randint(2)] 22 | near = torch.tensor(near, device=device) 23 | far = torch.tensor(far, device=device) 24 | depth = torch.tensor(depth, device=device) 25 | 26 | actual = utils3d.torch.project_depth(depth, near, far).cpu().numpy() 27 | 28 | assert np.allclose(expected, actual), '\n' + \ 29 | 'Input:\n' + \ 30 | f'\tdepth: {depth}\n' + \ 31 | f'\tnear: {near}\n' + \ 32 | f'\tfar: {far}\n' + \ 33 | 'Actual:\n' + \ 34 | f'{actual}\n' + \ 35 | 'Expected:\n' + \ 36 | f'{expected}' 37 | -------------------------------------------------------------------------------- /test/torch_/transforms/project_gl.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import torch 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | N = 1 13 | else: 14 | dim = np.random.randint(4) 15 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 16 | N = np.random.randint(1, 10) 17 | fovy = np.random.uniform(5 / 180 * np.pi, 175 / 180 * np.pi, spatial) 18 | aspect = np.random.uniform(0.01, 100, spatial) 19 | near = np.random.uniform(0.1, 100, spatial) 20 | far = np.random.uniform(near*2, 1000, spatial) 21 | eye = np.random.uniform(-10, 10, [*spatial, 3]) 22 | lookat = np.random.uniform(-10, 10, [*spatial, 3]) 23 | up = np.random.uniform(-10, 10, [*spatial, 3]) 24 | points = np.random.uniform(-10, 10, [*spatial, N, 3]) 25 | 26 | expected, _ = utils3d.numpy.transforms.project_gl(points, None, 27 | utils3d.numpy.view_look_at(eye, lookat, up), 28 | utils3d.numpy.perspective(fovy, aspect, near, far)) 29 | 30 | device = [torch.device('cpu'), torch.device('cuda')][np.random.randint(2)] 31 | fovy = torch.tensor(fovy, device=device) 32 | aspect = torch.tensor(aspect, device=device) 33 | near = torch.tensor(near, device=device) 34 | far = torch.tensor(far, device=device) 35 | eye = torch.tensor(eye, device=device) 36 | lookat = torch.tensor(lookat, device=device) 37 | up = torch.tensor(up, device=device) 38 | points = torch.tensor(points, device=device) 39 | 40 | actual, _ = utils3d.torch.project_gl(points, None, 41 | utils3d.torch.view_look_at(eye, lookat, up), 42 | utils3d.torch.perspective(fovy, aspect, near, far)) 43 | actual = actual.cpu().numpy() 44 | 45 | assert np.allclose(expected, actual), '\n' + \ 46 | 'Input:\n' + \ 47 | f'\tfovy: {fovy}\n' + \ 48 | f'\taspect: {aspect}\n' + \ 49 | f'\tnear: {near}\n' + \ 50 | f'\tfar: {far}\n' + \ 51 | f'\teye: {eye}\n' + \ 52 | f'\tlookat: {lookat}\n' + \ 53 | f'\tup: {up}\n' + \ 54 | f'\tpoints: {points}\n' + \ 55 | 'Actual:\n' + \ 56 | f'{actual}\n' + \ 57 | 'Expected:\n' + \ 58 | f'{expected}' 59 | -------------------------------------------------------------------------------- /test/torch_/transforms/project_gl_cv.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import torch 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | N = 1 13 | else: 14 | dim = np.random.randint(4) 15 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 16 | N = np.random.randint(1, 10) 17 | fovy = np.random.uniform(5 / 180 * np.pi, 175 / 180 * np.pi, spatial) 18 | aspect = np.random.uniform(0.01, 100, spatial) 19 | focal_x = 0.5 / (np.tan(fovy / 2) * aspect) 20 | focal_y = 0.5 / np.tan(fovy / 2) 21 | near = np.random.uniform(0.1, 100, spatial) 22 | far = np.random.uniform(near*2, 1000, spatial) 23 | eye = np.random.uniform(-10, 10, [*spatial, 3]) 24 | lookat = np.random.uniform(-10, 10, [*spatial, 3]) 25 | up = np.random.uniform(-10, 10, [*spatial, 3]) 26 | points = np.random.uniform(-10, 10, [*spatial, N, 3]) 27 | 28 | device = [torch.device('cpu'), torch.device('cuda')][np.random.randint(2)] 29 | fovy = torch.tensor(fovy, device=device) 30 | aspect = torch.tensor(aspect, device=device) 31 | focal_x = torch.tensor(focal_x, device=device) 32 | focal_y = torch.tensor(focal_y, device=device) 33 | near = torch.tensor(near, device=device) 34 | far = torch.tensor(far, device=device) 35 | eye = torch.tensor(eye, device=device) 36 | lookat = torch.tensor(lookat, device=device) 37 | up = torch.tensor(up, device=device) 38 | points = torch.tensor(points, device=device) 39 | 40 | gl = utils3d.torch.project_gl(points, None, 41 | utils3d.torch.view_look_at(eye, lookat, up), 42 | utils3d.torch.perspective(fovy, aspect, near, far)) 43 | gl_uv = gl[0][..., :2].cpu().numpy() 44 | gl_uv[..., 1] = 1 - gl_uv[..., 1] 45 | gl_depth = gl[1].cpu().numpy() 46 | 47 | cv = utils3d.torch.project_cv(points, 48 | utils3d.torch.extrinsics_look_at(eye, lookat, up), 49 | utils3d.torch.intrinsics(focal_x, focal_y, 0.5, 0.5)) 50 | cv_uv = cv[0][..., :2].cpu().numpy() 51 | cv_depth = cv[1].cpu().numpy() 52 | 53 | assert np.allclose(gl_uv, cv_uv) and np.allclose(gl_depth, cv_depth), '\n' + \ 54 | 'Input:\n' + \ 55 | f'\tfovy: {fovy}\n' + \ 56 | f'\taspect: {aspect}\n' + \ 57 | f'\teye: {eye}\n' + \ 58 | f'\tlookat: {lookat}\n' + \ 59 | f'\tup: {up}\n' + \ 60 | f'\tpoints: {points}\n' + \ 61 | 'GL:\n' + \ 62 | f'{gl}\n' + \ 63 | 'CV:\n' + \ 64 | f'{cv}' 65 | -------------------------------------------------------------------------------- /test/torch_/transforms/quaternion_to_matrix.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | from scipy.spatial.transform import Rotation as R 4 | import utils3d 5 | 6 | 7 | def run(): 8 | for i in range(10): 9 | if i == 0: 10 | spatial = [] 11 | else: 12 | dim = np.random.randint(4) 13 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 14 | angle = np.random.uniform(-np.pi, np.pi, spatial).astype(np.float32) 15 | axis = np.random.uniform(-1, 1, spatial + [3]).astype(np.float32) 16 | axis = axis / np.linalg.norm(axis, axis=-1, keepdims=True) 17 | axis_angle = angle[..., None] * axis 18 | quat = R.from_rotvec(axis_angle.reshape((-1, 3))).as_quat().reshape(spatial + [4]) 19 | expected = R.from_quat(quat.reshape(-1, 4)).as_matrix().reshape(spatial + [3, 3]) 20 | actual = utils3d.torch.quaternion_to_matrix( 21 | torch.from_numpy(quat[..., [3, 0, 1, 2]]) 22 | ).cpu().numpy() 23 | assert np.allclose(expected, actual), '\n' + \ 24 | 'Input:\n' + \ 25 | f'\tangle: {angle}\n' + \ 26 | f'\taxis: {axis}\n' + \ 27 | 'Actual:\n' + \ 28 | f'{actual}\n' + \ 29 | 'Expected:\n' + \ 30 | f'{expected}' 31 | 32 | if __name__ == '__main__': 33 | run() -------------------------------------------------------------------------------- /test/torch_/transforms/slerp.py: -------------------------------------------------------------------------------- 1 | from scipy.spatial.transform import Rotation, Slerp 2 | import numpy as np 3 | import torch 4 | import utils3d 5 | 6 | 7 | def run(): 8 | for i in range(100): 9 | quat_1 = np.random.rand(4) # [w, x, y, z] 10 | quat_2 = np.random.rand(4) 11 | t = np.array(0) 12 | expected = Slerp([0, 1], Rotation.from_quat([quat_1[[1, 2, 3, 0]], quat_2[[1, 2, 3, 0]]]))(t).as_matrix() 13 | matrix_1 = Rotation.from_quat(quat_1[[1, 2, 3, 0]]).as_matrix() 14 | matrix_2 = Rotation.from_quat(quat_2[[1, 2, 3, 0]]).as_matrix() 15 | actual = utils3d.torch.slerp( 16 | torch.from_numpy(matrix_1), 17 | torch.from_numpy(matrix_2), 18 | torch.from_numpy(t) 19 | ).numpy() 20 | assert np.allclose(actual, expected) 21 | 22 | if __name__ == '__main__': 23 | run() -------------------------------------------------------------------------------- /test/torch_/transforms/unproject_cv.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import torch 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | N = 1 13 | else: 14 | dim = np.random.randint(4) 15 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 16 | N = np.random.randint(1, 10) 17 | focal_x = np.random.uniform(0, 10, spatial) 18 | focal_y = np.random.uniform(0, 10, spatial) 19 | center_x = np.random.uniform(0, 1, spatial) 20 | center_y = np.random.uniform(0, 1, spatial) 21 | eye = np.random.uniform(-10, 10, [*spatial, 3]) 22 | lookat = np.random.uniform(-10, 10, [*spatial, 3]) 23 | up = np.random.uniform(-10, 10, [*spatial, 3]) 24 | points = np.random.uniform(-10, 10, [*spatial, N, 3]) 25 | 26 | expected = utils3d.numpy.transforms.unproject_cv( 27 | *utils3d.numpy.transforms.project_cv(points, 28 | utils3d.numpy.extrinsics_look_at(eye, lookat, up), 29 | utils3d.numpy.intrinsics(focal_x, focal_y, center_x, center_y)), 30 | utils3d.numpy.extrinsics_look_at(eye, lookat, up), 31 | utils3d.numpy.intrinsics(focal_x, focal_y, center_x, center_y) 32 | ) 33 | 34 | device = [torch.device('cpu'), torch.device('cuda')][np.random.randint(2)] 35 | focal_x = torch.tensor(focal_x, device=device) 36 | focal_y = torch.tensor(focal_y, device=device) 37 | center_x = torch.tensor(center_x, device=device) 38 | center_y = torch.tensor(center_y, device=device) 39 | eye = torch.tensor(eye, device=device) 40 | lookat = torch.tensor(lookat, device=device) 41 | up = torch.tensor(up, device=device) 42 | points = torch.tensor(points, device=device) 43 | 44 | actual = utils3d.torch.unproject_cv( 45 | *utils3d.torch.project_cv(points, 46 | utils3d.torch.extrinsics_look_at(eye, lookat, up), 47 | utils3d.torch.intrinsics(focal_x, focal_y, center_x, center_y)), 48 | utils3d.torch.extrinsics_look_at(eye, lookat, up), 49 | utils3d.torch.intrinsics(focal_x, focal_y, center_x, center_y) 50 | ) 51 | actual = actual.cpu().numpy() 52 | 53 | assert np.allclose(expected, actual), '\n' + \ 54 | 'Input:\n' + \ 55 | f'\tfocal_x: {focal_x}\n' + \ 56 | f'\tfocal_y: {focal_y}\n' + \ 57 | f'\tcenter_x: {center_x}\n' + \ 58 | f'\tcenter_y: {center_y}\n' + \ 59 | f'\teye: {eye}\n' + \ 60 | f'\tlookat: {lookat}\n' + \ 61 | f'\tup: {up}\n' + \ 62 | f'\tpoints: {points}\n' + \ 63 | 'Actual:\n' + \ 64 | f'{actual}\n' + \ 65 | 'Expected:\n' + \ 66 | f'{expected}' 67 | -------------------------------------------------------------------------------- /test/torch_/transforms/unproject_gl.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import torch 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | N = 1 13 | else: 14 | dim = np.random.randint(4) 15 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 16 | N = np.random.randint(1, 10) 17 | fovy = np.random.uniform(5 / 180 * np.pi, 175 / 180 * np.pi, spatial) 18 | aspect = np.random.uniform(0.01, 100, spatial) 19 | near = np.random.uniform(0.1, 100, spatial) 20 | far = np.random.uniform(near*2, 1000, spatial) 21 | eye = np.random.uniform(-10, 10, [*spatial, 3]) 22 | lookat = np.random.uniform(-10, 10, [*spatial, 3]) 23 | up = np.random.uniform(-10, 10, [*spatial, 3]) 24 | points = np.random.uniform(-10, 10, [*spatial, N, 3]) 25 | 26 | expected = utils3d.numpy.transforms.unproject_gl( 27 | utils3d.numpy.transforms.project_gl(points, None, 28 | utils3d.numpy.view_look_at(eye, lookat, up), 29 | utils3d.numpy.perspective(fovy, aspect, near, far))[0], 30 | None, 31 | utils3d.numpy.view_look_at(eye, lookat, up), 32 | utils3d.numpy.perspective(fovy, aspect, near, far) 33 | ) 34 | 35 | device = [torch.device('cpu'), torch.device('cuda')][np.random.randint(2)] 36 | fovy = torch.tensor(fovy, device=device) 37 | aspect = torch.tensor(aspect, device=device) 38 | near = torch.tensor(near, device=device) 39 | far = torch.tensor(far, device=device) 40 | eye = torch.tensor(eye, device=device) 41 | lookat = torch.tensor(lookat, device=device) 42 | up = torch.tensor(up, device=device) 43 | points = torch.tensor(points, device=device) 44 | 45 | actual = utils3d.torch.unproject_gl( 46 | utils3d.torch.project_gl(points, None, 47 | utils3d.torch.view_look_at(eye, lookat, up), 48 | utils3d.torch.perspective(fovy, aspect, near, far))[0], 49 | None, 50 | utils3d.torch.view_look_at(eye, lookat, up), 51 | utils3d.torch.perspective(fovy, aspect, near, far) 52 | ) 53 | actual = actual.cpu().numpy() 54 | 55 | assert np.allclose(expected, actual), '\n' + \ 56 | 'Input:\n' + \ 57 | f'\tfovy: {fovy}\n' + \ 58 | f'\taspect: {aspect}\n' + \ 59 | f'\tnear: {near}\n' + \ 60 | f'\tfar: {far}\n' + \ 61 | f'\teye: {eye}\n' + \ 62 | f'\tlookat: {lookat}\n' + \ 63 | f'\tup: {up}\n' + \ 64 | f'\tpoints: {points}\n' + \ 65 | 'Actual:\n' + \ 66 | f'{actual}\n' + \ 67 | 'Expected:\n' + \ 68 | f'{expected}' 69 | -------------------------------------------------------------------------------- /test/torch_/transforms/view_look_at.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import torch 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | else: 13 | dim = np.random.randint(4) 14 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 15 | eye = np.random.uniform(-10, 10, [*spatial, 3]).astype(np.float32) 16 | lookat = np.random.uniform(-10, 10, [*spatial, 3]).astype(np.float32) 17 | up = np.random.uniform(-10, 10, [*spatial, 3]).astype(np.float32) 18 | 19 | expected = utils3d.numpy.view_look_at(eye, lookat, up) 20 | 21 | device = [torch.device('cpu'), torch.device('cuda')][np.random.randint(2)] 22 | eye = torch.tensor(eye, device=device) 23 | lookat = torch.tensor(lookat, device=device) 24 | up = torch.tensor(up, device=device) 25 | 26 | actual = utils3d.torch.view_look_at(eye, lookat, up).cpu().numpy() 27 | 28 | assert np.allclose(expected, actual, 1e-5, 1e-5), '\n' + \ 29 | 'Input:\n' + \ 30 | f'eye: {eye}\n' + \ 31 | f'lookat: {lookat}\n' + \ 32 | f'up: {up}\n' + \ 33 | 'Actual:\n' + \ 34 | f'{actual}\n' + \ 35 | 'Expected:\n' + \ 36 | f'{expected}' 37 | -------------------------------------------------------------------------------- /test/torch_/transforms/view_to_extrinsic.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import torch 7 | 8 | def run(): 9 | for i in range(100): 10 | if i == 0: 11 | spatial = [] 12 | else: 13 | dim = np.random.randint(4) 14 | spatial = [np.random.randint(1, 10) for _ in range(dim)] 15 | eye = np.random.uniform(-10, 10, [*spatial, 3]).astype(np.float32) 16 | lookat = np.random.uniform(-10, 10, [*spatial, 3]).astype(np.float32) 17 | up = np.random.uniform(-10, 10, [*spatial, 3]).astype(np.float32) 18 | 19 | expected = utils3d.numpy.view_to_extrinsics(utils3d.numpy.view_look_at(eye, lookat, up)) 20 | 21 | device = [torch.device('cpu'), torch.device('cuda')][np.random.randint(2)] 22 | eye = torch.tensor(eye, device=device) 23 | lookat = torch.tensor(lookat, device=device) 24 | up = torch.tensor(up, device=device) 25 | 26 | actual = utils3d.torch.view_to_extrinsics(utils3d.torch.view_look_at(eye, lookat, up)).cpu().numpy() 27 | 28 | assert np.allclose(expected, actual, 1e-5, 1e-5), '\n' + \ 29 | 'Input:\n' + \ 30 | f'eye: {eye}\n' + \ 31 | f'lookat: {lookat}\n' + \ 32 | f'up: {up}\n' + \ 33 | 'Actual:\n' + \ 34 | f'{actual}\n' + \ 35 | 'Expected:\n' + \ 36 | f'{expected}' 37 | -------------------------------------------------------------------------------- /test/torch_/utils/image_mesh.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))) 4 | import utils3d 5 | import numpy as np 6 | import torch 7 | 8 | def run(): 9 | args = [ 10 | {'W':2, 'H':2, 'backslash': torch.tensor([False])}, 11 | {'W':2, 'H':2, 'backslash': torch.tensor([True])}, 12 | {'H':2, 'W':3, 'backslash': torch.tensor([True, False])}, 13 | ] 14 | 15 | expected = [ 16 | np.array([[0, 2, 1], [1, 2, 3]]), 17 | np.array([[0, 2, 3], [0, 3, 1]]), 18 | np.array([[0, 3, 4], [0, 4, 1], [1, 4, 2], [2, 4, 5]]), 19 | ] 20 | 21 | for args, expected in zip(args, expected): 22 | actual = utils3d.torch.triangulate( 23 | utils3d.torch.image_mesh(args['H'], args['W'])[1], 24 | backslash=args.get('backslash', None), 25 | ).cpu().numpy() 26 | 27 | assert np.allclose(expected, actual), '\n' + \ 28 | 'Input:\n' + \ 29 | f'{args}\n' + \ 30 | 'Actual:\n' + \ 31 | f'{actual}\n' + \ 32 | 'Expected:\n' + \ 33 | f'{expected}' 34 | -------------------------------------------------------------------------------- /utils3d/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | A package for common utility functions in 3D computer graphics and vision. 3 | Providing NumPy utilities in `utils3d.numpy`, PyTorch utilities in `utils3d.torch`, and IO utilities in `utils3d.io`. 4 | """ 5 | import importlib 6 | from typing import TYPE_CHECKING 7 | 8 | try: 9 | from ._unified import * 10 | except ImportError: 11 | pass 12 | 13 | __all__ = ['numpy', 'torch', 'io'] 14 | 15 | def __getattr__(name: str): 16 | return globals().get(name, importlib.import_module(f'.{name}', __package__)) 17 | 18 | if TYPE_CHECKING: 19 | from . import torch 20 | from . import numpy 21 | from . import io -------------------------------------------------------------------------------- /utils3d/_helpers.py: -------------------------------------------------------------------------------- 1 | from functools import wraps 2 | import warnings 3 | 4 | 5 | def suppress_traceback(fn): 6 | @wraps(fn) 7 | def wrapper(*args, **kwargs): 8 | try: 9 | return fn(*args, **kwargs) 10 | except Exception as e: 11 | e.__traceback__ = e.__traceback__.tb_next.tb_next 12 | raise 13 | return wrapper 14 | 15 | 16 | class no_warnings: 17 | def __init__(self, action: str = 'ignore', **kwargs): 18 | self.action = action 19 | self.filter_kwargs = kwargs 20 | 21 | def __call__(self, fn): 22 | @wraps(fn) 23 | def wrapper(*args, **kwargs): 24 | with warnings.catch_warnings(): 25 | warnings.simplefilter(self.action, **self.filter_kwargs) 26 | return fn(*args, **kwargs) 27 | return wrapper 28 | 29 | def __enter__(self): 30 | self.warnings_manager = warnings.catch_warnings() 31 | self.warnings_manager.__enter__() 32 | warnings.simplefilter(self.action, **self.filter_kwargs) 33 | 34 | def __exit__(self, exc_type, exc_val, exc_tb): 35 | self.warnings_manager.__exit__(exc_type, exc_val, exc_tb) 36 | -------------------------------------------------------------------------------- /utils3d/io/__init__.py: -------------------------------------------------------------------------------- 1 | from .obj import * 2 | from .colmap import * 3 | from .ply import * 4 | -------------------------------------------------------------------------------- /utils3d/io/colmap.py: -------------------------------------------------------------------------------- 1 | from typing import * 2 | from pathlib import Path 3 | 4 | import numpy as np 5 | from scipy.spatial.transform import Rotation 6 | 7 | 8 | __all__ = ['read_extrinsics_from_colmap', 'read_intrinsics_from_colmap', 'write_extrinsics_as_colmap', 'write_intrinsics_as_colmap'] 9 | 10 | 11 | def write_extrinsics_as_colmap(file: Union[str, Path], extrinsics: np.ndarray, image_names: Union[str, List[str]] = 'image_{i:04d}.png', camera_ids: List[int] = None): 12 | """ 13 | Write extrinsics to colmap `images.txt` file. 14 | Args: 15 | file: Path to `images.txt` file. 16 | extrinsics: (N, 4, 4) array of extrinsics. 17 | image_names: str or List of str, image names. Length is N. 18 | If str, it should be a format string with `i` as the index. (i starts from 1, in correspondence with IMAGE_ID in colmap) 19 | camera_ids: List of int, camera ids. Length is N. 20 | If None, it will be set to [1, 2, ..., N]. 21 | """ 22 | assert extrinsics.shape[1:] == (4, 4) and extrinsics.ndim == 3 or extrinsics.shape == (4, 4) 23 | if extrinsics.ndim == 2: 24 | extrinsics = extrinsics[np.newaxis, ...] 25 | quats = Rotation.from_matrix(extrinsics[:, :3, :3]).as_quat() 26 | trans = extrinsics[:, :3, 3] 27 | if camera_ids is None: 28 | camera_ids = list(range(1, len(extrinsics) + 1)) 29 | if isinstance(image_names, str): 30 | image_names = [image_names.format(i=i) for i in range(1, len(extrinsics) + 1)] 31 | assert len(extrinsics) == len(image_names) == len(camera_ids), \ 32 | f'Number of extrinsics ({len(extrinsics)}), image_names ({len(image_names)}), and camera_ids ({len(camera_ids)}) must be the same' 33 | with open(file, 'w') as fp: 34 | print("# IMAGE_ID, QW, QX, QY, QZ, TX, TY, TZ, CAMERA_ID, NAME", file=fp) 35 | for i, (quat, t, name, camera_id) in enumerate(zip(quats.tolist(), trans.tolist(), image_names, camera_ids)): 36 | # Colmap has wxyz order while scipy.spatial.transform.Rotation has xyzw order. 37 | qx, qy, qz, qw = quat 38 | tx, ty, tz = t 39 | print(f'{i + 1} {qw:f} {qx:f} {qy:f} {qz:f} {tx:f} {ty:f} {tz:f} {camera_id:d} {name}', file=fp) 40 | print() 41 | 42 | 43 | def write_intrinsics_as_colmap(file: Union[str, Path], intrinsics: np.ndarray, width: int, height: int, normalized: bool = False): 44 | """ 45 | Write intrinsics to colmap `cameras.txt` file. Currently only support PINHOLE model (no distortion) 46 | Args: 47 | file: Path to `cameras.txt` file. 48 | intrinsics: (N, 3, 3) array of intrinsics. 49 | width: Image width. 50 | height: Image height. 51 | normalized: Whether the intrinsics are normalized. If True, the intrinsics will unnormalized for writing. 52 | """ 53 | assert intrinsics.shape[1:] == (3, 3) and intrinsics.ndim == 3 or intrinsics.shape == (3, 3) 54 | if intrinsics.ndim == 2: 55 | intrinsics = intrinsics[np.newaxis, ...] 56 | if normalized: 57 | intrinsics = intrinsics * np.array([width, height, 1])[:, None] 58 | with open(file, 'w') as fp: 59 | print("# CAMERA_ID, MODEL, WIDTH, HEIGHT, PARAMS[]", file=fp) 60 | for i, intr in enumerate(intrinsics): 61 | fx, fy, cx, cy = intr[0, 0], intr[1, 1], intr[0, 2], intr[1, 2] 62 | print(f'{i + 1} PINHOLE {width:d} {height:d} {fx:f} {fy:f} {cx:f} {cy:f}', file=fp) 63 | 64 | 65 | def read_extrinsics_from_colmap(file: Union[str, Path]) -> Union[np.ndarray, List[int], List[str]]: 66 | """ 67 | Read extrinsics from colmap `images.txt` file. 68 | Args: 69 | file: Path to `images.txt` file. 70 | Returns: 71 | extrinsics: (N, 4, 4) array of extrinsics. 72 | camera_ids: List of int, camera ids. Length is N. Note that camera ids in colmap typically starts from 1. 73 | image_names: List of str, image names. Length is N. 74 | """ 75 | with open(file) as fp: 76 | lines = fp.readlines() 77 | image_names, quats, trans, camera_ids = [], [], [], [] 78 | i_line = 0 79 | for line in lines: 80 | line = line.strip() 81 | if line.startswith('#'): 82 | continue 83 | i_line += 1 84 | if i_line % 2 == 0: 85 | continue 86 | image_id, qw, qx, qy, qz, tx, ty, tz, camera_id, name = line.split() 87 | quats.append([float(qx), float(qy), float(qz), float(qw)]) 88 | trans.append([float(tx), float(ty), float(tz)]) 89 | camera_ids.append(int(camera_id)) 90 | image_names.append(name) 91 | 92 | quats = np.array(quats, dtype=np.float32) 93 | trans = np.array(trans, dtype=np.float32) 94 | rotation = Rotation.from_quat(quats).as_matrix() 95 | extrinsics = np.concatenate([ 96 | np.concatenate([rotation, trans[..., None]], axis=-1), 97 | np.array([0, 0, 0, 1], dtype=np.float32)[None, None, :].repeat(len(quats), axis=0) 98 | ], axis=-2) 99 | 100 | return extrinsics, camera_ids, image_names 101 | 102 | 103 | def read_intrinsics_from_colmap(file: Union[str, Path], normalize: bool = False) -> Tuple[List[int], np.ndarray, np.ndarray]: 104 | """ 105 | Read intrinsics from colmap `cameras.txt` file. 106 | Args: 107 | file: Path to `cameras.txt` file. 108 | normalize: Whether to normalize the intrinsics. If True, the intrinsics will be normalized. (mapping coordinates to [0, 1] range) 109 | Returns: 110 | camera_ids: List of int, camera ids. Length is N. Note that camera ids in colmap typically starts from 1. 111 | intrinsics: (N, 3, 3) array of intrinsics. 112 | distortions: (N, 5) array of distortions. 113 | """ 114 | with open(file) as fp: 115 | lines = fp.readlines() 116 | intrinsics, distortions, camera_ids = [], [], [] 117 | for line in lines: 118 | line = line.strip() 119 | if not line or line.startswith('#'): 120 | continue 121 | camera_id, model, width, height, *params = line.split() 122 | camera_id, width, height = int(camera_id), int(width), int(height) 123 | if model == 'PINHOLE': 124 | fx, fy, cx, cy = map(float, params[:4]) 125 | k1 = k2 = k3 = p1 = p2 = 0.0 126 | elif model == 'OPENCV': 127 | fx, fy, cx, cy, k1, k2, p1, p2, k3 = *map(float, params[:8]), 0.0 128 | elif model == 'SIMPLE_RADIAL': 129 | f, cx, cy, k = map(float, params[:4]) 130 | fx = fy = f 131 | k1, k2, p1, p2, k3 = k, 0.0, 0.0, 0.0, 0.0 132 | camera_ids.append(camera_id) 133 | if normalize: 134 | fx, fy, cx, cy = fx / width, fy / height, cx / width, cy / height 135 | intrinsics.append([[fx, 0, cx], [0, fy, cy], [0, 0, 1]]) 136 | distortions.append([k1, k2, p1, p2, k3]) 137 | intrinsics = np.array(intrinsics, dtype=np.float32) 138 | distortions = np.array(distortions, dtype=np.float32) 139 | return camera_ids, intrinsics, distortions 140 | -------------------------------------------------------------------------------- /utils3d/io/obj.py: -------------------------------------------------------------------------------- 1 | from io import TextIOWrapper 2 | from typing import Dict, Any, Union, Iterable 3 | import numpy as np 4 | from pathlib import Path 5 | 6 | __all__ = [ 7 | 'read_obj', 8 | 'write_obj', 9 | 'simple_write_obj' 10 | ] 11 | 12 | def read_obj( 13 | file : Union[str, Path, TextIOWrapper], 14 | encoding: Union[str, None] = None, 15 | ignore_unknown: bool = False 16 | ): 17 | """ 18 | Read wavefront .obj file, without preprocessing. 19 | 20 | Why bothering having this read_obj() while we already have other libraries like `trimesh`? 21 | This function read the raw format from .obj file and keeps the order of vertices and faces, 22 | while trimesh which involves modification like merge/split vertices, which could break the orders of vertices and faces, 23 | Those libraries are commonly aiming at geometry processing and rendering supporting various formats. 24 | If you want mesh geometry processing, you may turn to `trimesh` for more features. 25 | 26 | ### Parameters 27 | `file` (str, Path, TextIOWrapper): filepath or file object 28 | encoding (str, optional): 29 | 30 | ### Returns 31 | obj (dict): A dict containing .obj components 32 | { 33 | 'mtllib': [], 34 | 'v': [[0,1, 0.2, 1.0], [1.2, 0.0, 0.0], ...], 35 | 'vt': [[0.5, 0.5], ...], 36 | 'vn': [[0., 0.7, 0.7], [0., -0.7, 0.7], ...], 37 | 'f': [[0, 1, 2], [2, 3, 4],...], 38 | 'usemtl': [{'name': 'mtl1', 'f': 7}] 39 | } 40 | """ 41 | if hasattr(file,'read'): 42 | lines = file.read().splitlines() 43 | else: 44 | with open(file, 'r', encoding=encoding) as fp: 45 | lines = fp.read().splitlines() 46 | mtllib = [] 47 | v, vt, vn, vp = [], [], [], [] # Vertex coordinates, Vertex texture coordinate, Vertex normal, Vertex parameter 48 | f, ft, fn = [], [], [] # Face indices, Face texture indices, Face normal indices 49 | o = [] 50 | s = [] 51 | usemtl = [] 52 | 53 | def pad(l: list, n: Any): 54 | return l + [n] * (3 - len(l)) 55 | 56 | for i, line in enumerate(lines): 57 | sq = line.strip().split() 58 | if len(sq) == 0: 59 | continue 60 | if sq[0] == 'v': 61 | assert 4 <= len(sq) <= 5, f'Invalid format of line {i}: {line}' 62 | v.append([float(e) for e in sq[1:]][:3]) 63 | elif sq[0] == 'vt': 64 | assert 3 <= len(sq) <= 4, f'Invalid format of line {i}: {line}' 65 | vt.append([float(e) for e in sq[1:]][:2]) 66 | elif sq[0] == 'vn': 67 | assert len(sq) == 4, f'Invalid format of line {i}: {line}' 68 | vn.append([float(e) for e in sq[1:]]) 69 | elif sq[0] == 'vp': 70 | assert 2 <= len(sq) <= 4, f'Invalid format of line {i}: {line}' 71 | vp.append(pad([float(e) for e in sq[1:]], 0)) 72 | elif sq[0] == 'f': 73 | spliting = [pad([int(j) - 1 for j in e.split('/')], -1) for e in sq[1:]] 74 | f.append([e[0] for e in spliting]) 75 | ft.append([e[1] for e in spliting]) 76 | fn.append([e[2] for e in spliting]) 77 | elif sq[0] == 'usemtl': 78 | assert len(sq) == 2 79 | usemtl.append((sq[1], len(f))) 80 | elif sq[0] == 'o': 81 | assert len(sq) == 2 82 | o.append((sq[1], len(f))) 83 | elif sq[0] == 's': 84 | s.append((sq[1], len(f))) 85 | elif sq[0] == 'mtllib': 86 | assert len(sq) == 2 87 | mtllib.append(sq[1]) 88 | elif sq[0][0] == '#': 89 | continue 90 | else: 91 | if not ignore_unknown: 92 | raise Exception(f'Unknown keyword {sq[0]}') 93 | 94 | min_poly_vertices = min(len(f) for f in f) 95 | max_poly_vertices = max(len(f) for f in f) 96 | 97 | return { 98 | 'mtllib': mtllib, 99 | 'v': np.array(v, dtype=np.float32), 100 | 'vt': np.array(vt, dtype=np.float32), 101 | 'vn': np.array(vn, dtype=np.float32), 102 | 'vp': np.array(vp, dtype=np.float32), 103 | 'f': np.array(f, dtype=np.int32) if min_poly_vertices == max_poly_vertices else f, 104 | 'ft': np.array(ft, dtype=np.int32) if min_poly_vertices == max_poly_vertices else ft, 105 | 'fn': np.array(fn, dtype=np.int32) if min_poly_vertices == max_poly_vertices else fn, 106 | 'o': o, 107 | 's': s, 108 | 'usemtl': usemtl, 109 | } 110 | 111 | 112 | def write_obj( 113 | file: Union[str, Path], 114 | obj: Dict[str, Any], 115 | encoding: Union[str, None] = None 116 | ): 117 | with open(file, 'w', encoding=encoding) as fp: 118 | for k in ['v', 'vt', 'vn', 'vp']: 119 | if k not in obj: 120 | continue 121 | for v in obj[k]: 122 | print(k, *map(float, v), file=fp) 123 | for f in obj['f']: 124 | print('f', *((str('/').join(map(int, i)) if isinstance(int(i), Iterable) else i) for i in f), file=fp) 125 | 126 | 127 | def simple_write_obj( 128 | file: Union[str, Path], 129 | vertices: np.ndarray, 130 | faces: np.ndarray, 131 | encoding: Union[str, None] = None 132 | ): 133 | """ 134 | Write wavefront .obj file, without preprocessing. 135 | 136 | Args: 137 | vertices (np.ndarray): [N, 3] 138 | faces (np.ndarray): [T, 3] 139 | file (Any): filepath 140 | encoding (str, optional): 141 | """ 142 | with open(file, 'w', encoding=encoding) as fp: 143 | for v in vertices: 144 | print('v', *map(float, v), file=fp) 145 | for f in faces: 146 | print('f', *map(int, f + 1), file=fp) 147 | -------------------------------------------------------------------------------- /utils3d/io/ply.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from typing import * 4 | from pathlib import Path 5 | 6 | 7 | def read_ply( 8 | file: Union[str, Path], 9 | encoding: Union[str, None] = None, 10 | ignore_unknown: bool = False 11 | ) -> Tuple[np.ndarray, np.ndarray]: 12 | """ 13 | Read .ply file, without preprocessing. 14 | 15 | Args: 16 | file (Any): filepath 17 | encoding (str, optional): 18 | 19 | Returns: 20 | Tuple[np.ndarray, np.ndarray]: vertices, faces 21 | """ 22 | import plyfile 23 | plydata = plyfile.PlyData.read(file) 24 | vertices = np.stack([plydata['vertex'][k] for k in ['x', 'y', 'z']], axis=-1) 25 | if 'face' in plydata: 26 | faces = np.array(plydata['face']['vertex_indices'].tolist()) 27 | else: 28 | faces = None 29 | return vertices, faces 30 | 31 | 32 | def write_ply( 33 | file: Union[str, Path], 34 | vertices: np.ndarray, 35 | faces: np.ndarray = None, 36 | edges: np.ndarray = None, 37 | vertex_colors: np.ndarray = None, 38 | edge_colors: np.ndarray = None, 39 | text: bool = False 40 | ): 41 | """ 42 | Write .ply file, without preprocessing. 43 | 44 | Args: 45 | file (Any): filepath 46 | vertices (np.ndarray): [N, 3] 47 | faces (np.ndarray): [T, E] 48 | edges (np.ndarray): [E, 2] 49 | vertex_colors (np.ndarray, optional): [N, 3]. Defaults to None. 50 | edge_colors (np.ndarray, optional): [E, 3]. Defaults to None. 51 | text (bool, optional): save data in text format. Defaults to False. 52 | """ 53 | import plyfile 54 | assert vertices.ndim == 2 and vertices.shape[1] == 3 55 | vertices = vertices.astype(np.float32) 56 | if faces is not None: 57 | assert faces.ndim == 2 58 | faces = faces.astype(np.int32) 59 | if edges is not None: 60 | assert edges.ndim == 2 and edges.shape[1] == 2 61 | edges = edges.astype(np.int32) 62 | 63 | if vertex_colors is not None: 64 | assert vertex_colors.ndim == 2 and vertex_colors.shape[1] == 3 65 | if vertex_colors.dtype in [np.float32, np.float64]: 66 | vertex_colors = vertex_colors * 255 67 | vertex_colors = np.clip(vertex_colors, 0, 255).astype(np.uint8) 68 | vertices_data = np.zeros(len(vertices), dtype=[('x', 'f4'), ('y', 'f4'), ('z', 'f4'), ('red', 'u1'), ('green', 'u1'), ('blue', 'u1')]) 69 | vertices_data['x'] = vertices[:, 0] 70 | vertices_data['y'] = vertices[:, 1] 71 | vertices_data['z'] = vertices[:, 2] 72 | vertices_data['red'] = vertex_colors[:, 0] 73 | vertices_data['green'] = vertex_colors[:, 1] 74 | vertices_data['blue'] = vertex_colors[:, 2] 75 | else: 76 | vertices_data = np.array([tuple(v) for v in vertices], dtype=[('x', 'f4'), ('y', 'f4'), ('z', 'f4')]) 77 | 78 | if faces is not None: 79 | faces_data = np.zeros(len(faces), dtype=[('vertex_indices', 'i4', (faces.shape[1],))]) 80 | faces_data['vertex_indices'] = faces 81 | 82 | if edges is not None: 83 | if edge_colors is not None: 84 | assert edge_colors.ndim == 2 and edge_colors.shape[1] == 3 85 | if edge_colors.dtype in [np.float32, np.float64]: 86 | edge_colors = edge_colors * 255 87 | edge_colors = np.clip(edge_colors, 0, 255).astype(np.uint8) 88 | edges_data = np.zeros(len(edges), dtype=[('vertex1', 'i4'), ('vertex2', 'i4'), ('red', 'u1'), ('green', 'u1'), ('blue', 'u1')]) 89 | edges_data['vertex1'] = edges[:, 0] 90 | edges_data['vertex2'] = edges[:, 1] 91 | edges_data['red'] = edge_colors[:, 0] 92 | edges_data['green'] = edge_colors[:, 1] 93 | edges_data['blue'] = edge_colors[:, 2] 94 | else: 95 | edges_data = np.array([tuple(e) for e in edges], dtype=[('vertex1', 'i4'), ('vertex2', 'i4')]) 96 | 97 | ply_data = [plyfile.PlyElement.describe(vertices_data, 'vertex')] 98 | if faces is not None: 99 | ply_data.append(plyfile.PlyElement.describe(faces_data, 'face')) 100 | if edges is not None: 101 | ply_data.append(plyfile.PlyElement.describe(edges_data, 'edge')) 102 | 103 | plyfile.PlyData(ply_data, text=text).write(file) 104 | -------------------------------------------------------------------------------- /utils3d/numpy/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | 3D utility functions workings with NumPy. 3 | """ 4 | import importlib 5 | import itertools 6 | import numpy 7 | from typing import TYPE_CHECKING 8 | 9 | 10 | __modules_all__ = { 11 | 'mesh':[ 12 | 'triangulate', 13 | 'compute_face_normal', 14 | 'compute_face_angle', 15 | 'compute_vertex_normal', 16 | 'compute_vertex_normal_weighted', 17 | 'remove_corrupted_faces', 18 | 'merge_duplicate_vertices', 19 | 'remove_unreferenced_vertices', 20 | 'subdivide_mesh_simple', 21 | 'mesh_relations', 22 | 'flatten_mesh_indices' 23 | ], 24 | 'quadmesh': [ 25 | 'calc_quad_candidates', 26 | 'calc_quad_distortion', 27 | 'calc_quad_direction', 28 | 'calc_quad_smoothness', 29 | 'sovle_quad', 30 | 'sovle_quad_qp', 31 | 'tri_to_quad' 32 | ], 33 | 'utils': [ 34 | 'sliding_window_1d', 35 | 'sliding_window_nd', 36 | 'sliding_window_2d', 37 | 'max_pool_1d', 38 | 'max_pool_2d', 39 | 'max_pool_nd', 40 | 'depth_edge', 41 | 'normals_edge', 42 | 'depth_aliasing', 43 | 'interpolate', 44 | 'image_scrcoord', 45 | 'image_uv', 46 | 'image_pixel_center', 47 | 'image_pixel', 48 | 'image_mesh', 49 | 'image_mesh_from_depth', 50 | 'depth_to_normals', 51 | 'points_to_normals', 52 | 'depth_to_points', 53 | 'chessboard', 54 | 'cube', 55 | 'icosahedron', 56 | 'square', 57 | 'camera_frustum', 58 | ], 59 | 'transforms': [ 60 | 'perspective', 61 | 'perspective_from_fov', 62 | 'perspective_from_fov_xy', 63 | 'intrinsics_from_focal_center', 64 | 'intrinsics_from_fov', 65 | 'fov_to_focal', 66 | 'focal_to_fov', 67 | 'intrinsics_to_fov', 68 | 'view_look_at', 69 | 'extrinsics_look_at', 70 | 'perspective_to_intrinsics', 71 | 'perspective_to_near_far', 72 | 'intrinsics_to_perspective', 73 | 'extrinsics_to_view', 74 | 'view_to_extrinsics', 75 | 'normalize_intrinsics', 76 | 'crop_intrinsics', 77 | 'pixel_to_uv', 78 | 'pixel_to_ndc', 79 | 'uv_to_pixel', 80 | 'project_depth', 81 | 'depth_buffer_to_linear', 82 | 'unproject_cv', 83 | 'unproject_gl', 84 | 'project_cv', 85 | 'project_gl', 86 | 'quaternion_to_matrix', 87 | 'axis_angle_to_matrix', 88 | 'matrix_to_quaternion', 89 | 'extrinsics_to_essential', 90 | 'euler_axis_angle_rotation', 91 | 'euler_angles_to_matrix', 92 | 'skew_symmetric', 93 | 'rotation_matrix_from_vectors', 94 | 'ray_intersection', 95 | 'se3_matrix', 96 | 'slerp_quaternion', 97 | 'slerp_vector', 98 | 'lerp', 99 | 'lerp_se3_matrix', 100 | 'piecewise_lerp', 101 | 'piecewise_lerp_se3_matrix', 102 | 'apply_transform', 103 | 'angle_diff_vec3' 104 | ], 105 | 'spline': [ 106 | 'linear_spline_interpolate', 107 | ], 108 | 'rasterization': [ 109 | 'RastContext', 110 | 'rasterize_triangle_faces', 111 | 'rasterize_edges', 112 | 'texture', 113 | 'warp_image_by_depth', 114 | 'test_rasterization' 115 | ], 116 | } 117 | 118 | 119 | __all__ = list(itertools.chain(*__modules_all__.values())) 120 | 121 | def __getattr__(name): 122 | try: 123 | return globals()[name] 124 | except KeyError: 125 | pass 126 | 127 | try: 128 | module_name = next(m for m in __modules_all__ if name in __modules_all__[m]) 129 | except StopIteration: 130 | raise AttributeError(f"module '{__name__}' has no attribute '{name}'") 131 | module = importlib.import_module(f'.{module_name}', __name__) 132 | for key in __modules_all__[module_name]: 133 | globals()[key] = getattr(module, key) 134 | 135 | return globals()[name] 136 | 137 | 138 | if TYPE_CHECKING: 139 | from .quadmesh import * 140 | from .transforms import * 141 | from .mesh import * 142 | from .utils import * 143 | from .rasterization import * 144 | from .spline import * -------------------------------------------------------------------------------- /utils3d/numpy/_helpers.py: -------------------------------------------------------------------------------- 1 | # decorator 2 | import numpy as np 3 | from numbers import Number 4 | import inspect 5 | from functools import wraps 6 | from typing import * 7 | from .._helpers import suppress_traceback 8 | 9 | 10 | def get_args_order(func, args, kwargs): 11 | """ 12 | Get the order of the arguments of a function. 13 | """ 14 | names = inspect.getfullargspec(func).args 15 | names_idx = {name: i for i, name in enumerate(names)} 16 | args_order = [] 17 | kwargs_order = {} 18 | for name, arg in kwargs.items(): 19 | if name in names: 20 | kwargs_order[name] = names_idx[name] 21 | names.remove(name) 22 | for i, arg in enumerate(args): 23 | if i < len(names): 24 | args_order.append(names_idx[names[i]]) 25 | return args_order, kwargs_order 26 | 27 | 28 | def broadcast_args(args, kwargs, args_dim, kwargs_dim): 29 | spatial = [] 30 | for arg, arg_dim in zip(args + list(kwargs.values()), args_dim + list(kwargs_dim.values())): 31 | if isinstance(arg, np.ndarray) and arg_dim is not None: 32 | arg_spatial = arg.shape[:arg.ndim-arg_dim] 33 | if len(arg_spatial) > len(spatial): 34 | spatial = [1] * (len(arg_spatial) - len(spatial)) + spatial 35 | for j in range(len(arg_spatial)): 36 | if spatial[-j] < arg_spatial[-j]: 37 | if spatial[-j] == 1: 38 | spatial[-j] = arg_spatial[-j] 39 | else: 40 | raise ValueError("Cannot broadcast arguments.") 41 | for i, arg in enumerate(args): 42 | if isinstance(arg, np.ndarray) and args_dim[i] is not None: 43 | args[i] = np.broadcast_to(arg, [*spatial, *arg.shape[arg.ndim-args_dim[i]:]]) 44 | for key, arg in kwargs.items(): 45 | if isinstance(arg, np.ndarray) and kwargs_dim[key] is not None: 46 | kwargs[key] = np.broadcast_to(arg, [*spatial, *arg.shape[arg.ndim-kwargs_dim[key]:]]) 47 | return args, kwargs, spatial 48 | 49 | 50 | def batched(*dims): 51 | """ 52 | Decorator that allows a function to be called with batched arguments. 53 | """ 54 | def decorator(func): 55 | @wraps(func) 56 | @suppress_traceback 57 | def wrapper(*args, **kwargs): 58 | args = list(args) 59 | # get arguments dimensions 60 | args_order, kwargs_order = get_args_order(func, args, kwargs) 61 | args_dim = [dims[i] for i in args_order] 62 | kwargs_dim = {key: dims[i] for key, i in kwargs_order.items()} 63 | # convert to numpy array 64 | for i, arg in enumerate(args): 65 | if isinstance(arg, (Number, list, tuple)) and args_dim[i] is not None: 66 | args[i] = np.array(arg) 67 | for key, arg in kwargs.items(): 68 | if isinstance(arg, (Number, list, tuple)) and kwargs_dim[key] is not None: 69 | kwargs[key] = np.array(arg) 70 | # broadcast arguments 71 | args, kwargs, spatial = broadcast_args(args, kwargs, args_dim, kwargs_dim) 72 | for i, (arg, arg_dim) in enumerate(zip(args, args_dim)): 73 | if isinstance(arg, np.ndarray) and arg_dim is not None: 74 | args[i] = arg.reshape([-1, *arg.shape[arg.ndim-arg_dim:]]) 75 | for key, arg in kwargs.items(): 76 | if isinstance(arg, np.ndarray) and kwargs_dim[key] is not None: 77 | kwargs[key] = arg.reshape([-1, *arg.shape[arg.ndim-kwargs_dim[key]:]]) 78 | # call function 79 | results = func(*args, **kwargs) 80 | type_results = type(results) 81 | results = list(results) if isinstance(results, (tuple, list)) else [results] 82 | # restore spatial dimensions 83 | for i, result in enumerate(results): 84 | results[i] = result.reshape([*spatial, *result.shape[1:]]) 85 | if type_results == tuple: 86 | results = tuple(results) 87 | elif type_results == list: 88 | results = list(results) 89 | else: 90 | results = results[0] 91 | return results 92 | return wrapper 93 | return decorator 94 | -------------------------------------------------------------------------------- /utils3d/numpy/mesh.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from typing import * 3 | from ._helpers import batched 4 | 5 | 6 | __all__ = [ 7 | 'triangulate', 8 | 'compute_face_normal', 9 | 'compute_face_angle', 10 | 'compute_vertex_normal', 11 | 'compute_vertex_normal_weighted', 12 | 'remove_corrupted_faces', 13 | 'merge_duplicate_vertices', 14 | 'remove_unreferenced_vertices', 15 | 'subdivide_mesh_simple', 16 | 'mesh_relations', 17 | 'flatten_mesh_indices' 18 | ] 19 | 20 | 21 | def triangulate( 22 | faces: np.ndarray, 23 | vertices: np.ndarray = None, 24 | backslash: np.ndarray = None 25 | ) -> np.ndarray: 26 | """ 27 | Triangulate a polygonal mesh. 28 | 29 | Args: 30 | faces (np.ndarray): [L, P] polygonal faces 31 | vertices (np.ndarray, optional): [N, 3] 3-dimensional vertices. 32 | If given, the triangulation is performed according to the distance 33 | between vertices. Defaults to None. 34 | backslash (np.ndarray, optional): [L] boolean array indicating 35 | how to triangulate the quad faces. Defaults to None. 36 | 37 | Returns: 38 | (np.ndarray): [L * (P - 2), 3] triangular faces 39 | """ 40 | if faces.shape[-1] == 3: 41 | return faces 42 | P = faces.shape[-1] 43 | if vertices is not None: 44 | assert faces.shape[-1] == 4, "now only support quad mesh" 45 | if backslash is None: 46 | backslash = np.linalg.norm(vertices[faces[:, 0]] - vertices[faces[:, 2]], axis=-1) < \ 47 | np.linalg.norm(vertices[faces[:, 1]] - vertices[faces[:, 3]], axis=-1) 48 | if backslash is None: 49 | loop_indice = np.stack([ 50 | np.zeros(P - 2, dtype=int), 51 | np.arange(1, P - 1, 1, dtype=int), 52 | np.arange(2, P, 1, dtype=int) 53 | ], axis=1) 54 | return faces[:, loop_indice].reshape((-1, 3)) 55 | else: 56 | assert faces.shape[-1] == 4, "now only support quad mesh" 57 | faces = np.where( 58 | backslash[:, None], 59 | faces[:, [0, 1, 2, 0, 2, 3]], 60 | faces[:, [0, 1, 3, 3, 1, 2]] 61 | ).reshape((-1, 3)) 62 | return faces 63 | 64 | 65 | @batched(2, None) 66 | def compute_face_normal( 67 | vertices: np.ndarray, 68 | faces: np.ndarray 69 | ) -> np.ndarray: 70 | """ 71 | Compute face normals of a triangular mesh 72 | 73 | Args: 74 | vertices (np.ndarray): [..., N, 3] 3-dimensional vertices 75 | faces (np.ndarray): [T, 3] triangular face indices 76 | 77 | Returns: 78 | normals (np.ndarray): [..., T, 3] face normals 79 | """ 80 | normal = np.cross( 81 | vertices[..., faces[:, 1], :] - vertices[..., faces[:, 0], :], 82 | vertices[..., faces[:, 2], :] - vertices[..., faces[:, 0], :] 83 | ) 84 | normal_norm = np.linalg.norm(normal, axis=-1, keepdims=True) 85 | normal_norm[normal_norm == 0] = 1 86 | normal /= normal_norm 87 | return normal 88 | 89 | 90 | @batched(2, None) 91 | def compute_face_angle( 92 | vertices: np.ndarray, 93 | faces: np.ndarray, 94 | eps: float = 1e-12 95 | ) -> np.ndarray: 96 | """ 97 | Compute face angles of a triangular mesh 98 | 99 | Args: 100 | vertices (np.ndarray): [..., N, 3] 3-dimensional vertices 101 | faces (np.ndarray): [T, 3] triangular face indices 102 | 103 | Returns: 104 | angles (np.ndarray): [..., T, 3] face angles 105 | """ 106 | face_angle = np.zeros_like(faces, dtype=vertices.dtype) 107 | for i in range(3): 108 | edge1 = vertices[..., faces[:, (i + 1) % 3], :] - vertices[..., faces[:, i], :] 109 | edge2 = vertices[..., faces[:, (i + 2) % 3], :] - vertices[..., faces[:, i], :] 110 | face_angle[..., i] = np.arccos(np.sum( 111 | edge1 / np.clip(np.linalg.norm(edge1, axis=-1, keepdims=True), eps, None) * 112 | edge2 / np.clip(np.linalg.norm(edge2, axis=-1, keepdims=True), eps, None), 113 | axis=-1 114 | )) 115 | return face_angle 116 | 117 | 118 | @batched(2, None, 2) 119 | def compute_vertex_normal( 120 | vertices: np.ndarray, 121 | faces: np.ndarray, 122 | face_normal: np.ndarray = None 123 | ) -> np.ndarray: 124 | """ 125 | Compute vertex normals of a triangular mesh by averaging neightboring face normals 126 | TODO: can be improved. 127 | 128 | Args: 129 | vertices (np.ndarray): [..., N, 3] 3-dimensional vertices 130 | faces (np.ndarray): [T, 3] triangular face indices 131 | face_normal (np.ndarray, optional): [..., T, 3] face normals. 132 | None to compute face normals from vertices and faces. Defaults to None. 133 | 134 | Returns: 135 | normals (np.ndarray): [..., N, 3] vertex normals 136 | """ 137 | if face_normal is None: 138 | face_normal = compute_face_normal(vertices, faces) 139 | vertex_normal = np.zeros_like(vertices, dtype=vertices.dtype) 140 | for n in range(vertices.shape[0]): 141 | for i in range(3): 142 | vertex_normal[n, :, 0] += np.bincount(faces[:, i], weights=face_normal[n, :, 0], minlength=vertices.shape[1]) 143 | vertex_normal[n, :, 1] += np.bincount(faces[:, i], weights=face_normal[n, :, 1], minlength=vertices.shape[1]) 144 | vertex_normal[n, :, 2] += np.bincount(faces[:, i], weights=face_normal[n, :, 2], minlength=vertices.shape[1]) 145 | vertex_normal_norm = np.linalg.norm(vertex_normal, axis=-1, keepdims=True) 146 | vertex_normal_norm[vertex_normal_norm == 0] = 1 147 | vertex_normal /= vertex_normal_norm 148 | return vertex_normal 149 | 150 | 151 | @batched(2, None, 2) 152 | def compute_vertex_normal_weighted( 153 | vertices: np.ndarray, 154 | faces: np.ndarray, 155 | face_normal: np.ndarray = None 156 | ) -> np.ndarray: 157 | """ 158 | Compute vertex normals of a triangular mesh by weighted sum of neightboring face normals 159 | according to the angles 160 | 161 | Args: 162 | vertices (np.ndarray): [..., N, 3] 3-dimensional vertices 163 | faces (np.ndarray): [..., T, 3] triangular face indices 164 | face_normal (np.ndarray, optional): [..., T, 3] face normals. 165 | None to compute face normals from vertices and faces. Defaults to None. 166 | 167 | Returns: 168 | normals (np.ndarray): [..., N, 3] vertex normals 169 | """ 170 | if face_normal is None: 171 | face_normal = compute_face_normal(vertices, faces) 172 | face_angle = compute_face_angle(vertices, faces) 173 | vertex_normal = np.zeros_like(vertices) 174 | for n in range(vertices.shape[0]): 175 | for i in range(3): 176 | vertex_normal[n, :, 0] += np.bincount(faces[n, :, i], weights=face_normal[n, :, 0] * face_angle[n, :, i], minlength=vertices.shape[1]) 177 | vertex_normal[n, :, 1] += np.bincount(faces[n, :, i], weights=face_normal[n, :, 1] * face_angle[n, :, i], minlength=vertices.shape[1]) 178 | vertex_normal[n, :, 2] += np.bincount(faces[n, :, i], weights=face_normal[n, :, 2] * face_angle[n, :, i], minlength=vertices.shape[1]) 179 | vertex_normal_norm = np.linalg.norm(vertex_normal, axis=-1, keepdims=True) 180 | vertex_normal_norm[vertex_normal_norm == 0] = 1 181 | vertex_normal /= vertex_normal_norm 182 | return vertex_normal 183 | 184 | 185 | def remove_corrupted_faces( 186 | faces: np.ndarray 187 | ) -> np.ndarray: 188 | """ 189 | Remove corrupted faces (faces with duplicated vertices) 190 | 191 | Args: 192 | faces (np.ndarray): [T, 3] triangular face indices 193 | 194 | Returns: 195 | np.ndarray: [T_, 3] triangular face indices 196 | """ 197 | corrupted = (faces[:, 0] == faces[:, 1]) | (faces[:, 1] == faces[:, 2]) | (faces[:, 2] == faces[:, 0]) 198 | return faces[~corrupted] 199 | 200 | 201 | def merge_duplicate_vertices( 202 | vertices: np.ndarray, 203 | faces: np.ndarray, 204 | tol: float = 1e-6 205 | ) -> Tuple[np.ndarray, np.ndarray]: 206 | """ 207 | Merge duplicate vertices of a triangular mesh. 208 | Duplicate vertices are merged by selecte one of them, and the face indices are updated accordingly. 209 | 210 | Args: 211 | vertices (np.ndarray): [N, 3] 3-dimensional vertices 212 | faces (np.ndarray): [T, 3] triangular face indices 213 | tol (float, optional): tolerance for merging. Defaults to 1e-6. 214 | 215 | Returns: 216 | vertices (np.ndarray): [N_, 3] 3-dimensional vertices 217 | faces (np.ndarray): [T, 3] triangular face indices 218 | """ 219 | vertices_round = np.round(vertices / tol) 220 | _, uni_i, uni_inv = np.unique(vertices_round, return_index=True, return_inverse=True, axis=0) 221 | vertices = vertices[uni_i] 222 | faces = uni_inv[faces] 223 | return vertices, faces 224 | 225 | 226 | def remove_unreferenced_vertices( 227 | faces: np.ndarray, 228 | *vertice_attrs, 229 | return_indices: bool = False 230 | ) -> Tuple[np.ndarray, ...]: 231 | """ 232 | Remove unreferenced vertices of a mesh. 233 | Unreferenced vertices are removed, and the face indices are updated accordingly. 234 | 235 | Args: 236 | faces (np.ndarray): [T, P] face indices 237 | *vertice_attrs: vertex attributes 238 | 239 | Returns: 240 | faces (np.ndarray): [T, P] face indices 241 | *vertice_attrs: vertex attributes 242 | indices (np.ndarray, optional): [N] indices of vertices that are kept. Defaults to None. 243 | """ 244 | P = faces.shape[-1] 245 | fewer_indices, inv_map = np.unique(faces, return_inverse=True) 246 | faces = inv_map.astype(np.int32).reshape(-1, P) 247 | ret = [faces] 248 | for attr in vertice_attrs: 249 | ret.append(attr[fewer_indices]) 250 | if return_indices: 251 | ret.append(fewer_indices) 252 | return tuple(ret) 253 | 254 | 255 | def subdivide_mesh_simple( 256 | vertices: np.ndarray, 257 | faces: np.ndarray, 258 | n: int = 1 259 | ) -> Tuple[np.ndarray, np.ndarray]: 260 | """ 261 | Subdivide a triangular mesh by splitting each triangle into 4 smaller triangles. 262 | NOTE: All original vertices are kept, and new vertices are appended to the end of the vertex list. 263 | 264 | Args: 265 | vertices (np.ndarray): [N, 3] 3-dimensional vertices 266 | faces (np.ndarray): [T, 3] triangular face indices 267 | n (int, optional): number of subdivisions. Defaults to 1. 268 | 269 | Returns: 270 | vertices (np.ndarray): [N_, 3] subdivided 3-dimensional vertices 271 | faces (np.ndarray): [4 * T, 3] subdivided triangular face indices 272 | """ 273 | for _ in range(n): 274 | edges = np.stack([faces[:, [0, 1]], faces[:, [1, 2]], faces[:, [2, 0]]], axis=0) 275 | edges = np.sort(edges, axis=2) 276 | uni_edges, uni_inv = np.unique(edges.reshape(-1, 2), return_inverse=True, axis=0) 277 | uni_inv = uni_inv.reshape(3, -1) 278 | midpoints = (vertices[uni_edges[:, 0]] + vertices[uni_edges[:, 1]]) / 2 279 | 280 | n_vertices = vertices.shape[0] 281 | vertices = np.concatenate([vertices, midpoints], axis=0) 282 | faces = np.concatenate([ 283 | np.stack([faces[:, 0], n_vertices + uni_inv[0], n_vertices + uni_inv[2]], axis=1), 284 | np.stack([faces[:, 1], n_vertices + uni_inv[1], n_vertices + uni_inv[0]], axis=1), 285 | np.stack([faces[:, 2], n_vertices + uni_inv[2], n_vertices + uni_inv[1]], axis=1), 286 | np.stack([n_vertices + uni_inv[0], n_vertices + uni_inv[1], n_vertices + uni_inv[2]], axis=1), 287 | ], axis=0) 288 | return vertices, faces 289 | 290 | 291 | def mesh_relations( 292 | faces: np.ndarray, 293 | ) -> Tuple[np.ndarray, np.ndarray]: 294 | """ 295 | Calculate the relation between vertices and faces. 296 | NOTE: The input mesh must be a manifold triangle mesh. 297 | 298 | Args: 299 | faces (np.ndarray): [T, 3] triangular face indices 300 | 301 | Returns: 302 | edges (np.ndarray): [E, 2] edge indices 303 | edge2face (np.ndarray): [E, 2] edge to face relation. The second column is -1 if the edge is boundary. 304 | face2edge (np.ndarray): [T, 3] face to edge relation 305 | face2face (np.ndarray): [T, 3] face to face relation 306 | """ 307 | T = faces.shape[0] 308 | edges = np.stack([faces[:, [0, 1]], faces[:, [1, 2]], faces[:, [2, 0]]], axis=1).reshape(-1, 2) # [3T, 2] 309 | edges = np.sort(edges, axis=1) # [3T, 2] 310 | edges, face2edge, occurence = np.unique(edges, axis=0, return_inverse=True, return_counts=True) # [E, 2], [3T], [E] 311 | E = edges.shape[0] 312 | assert np.all(occurence <= 2), "The input mesh is not a manifold mesh." 313 | 314 | # Edge to face relation 315 | padding = np.arange(E, dtype=np.int32)[occurence == 1] 316 | padded_face2edge = np.concatenate([face2edge, padding], axis=0) # [2E] 317 | edge2face = np.argsort(padded_face2edge, kind='stable').reshape(-1, 2) // 3 # [E, 2] 318 | edge2face_valid = edge2face[:, 1] < T # [E] 319 | edge2face[~edge2face_valid, 1] = -1 320 | 321 | # Face to edge relation 322 | face2edge = face2edge.reshape(-1, 3) # [T, 3] 323 | 324 | # Face to face relation 325 | face2face = edge2face[face2edge] # [T, 3, 2] 326 | face2face = face2face[face2face != np.arange(T)[:, None, None]].reshape(T, 3) # [T, 3] 327 | 328 | return edges, edge2face, face2edge, face2face 329 | 330 | 331 | @overload 332 | def flatten_mesh_indices(faces1: np.ndarray, attr1: np.ndarray, *other_faces_attrs_pairs: np.ndarray) -> Tuple[np.ndarray, ...]: 333 | """ 334 | Rearrange the indices of a mesh to a flattened version. Vertices will be no longer shared. 335 | 336 | ### Parameters: 337 | - `faces1`: [T, P] face indices of the first attribute 338 | - `attr1`: [N1, ...] attributes of the first mesh 339 | - ... 340 | 341 | ### Returns: 342 | - `faces`: [T, P] flattened face indices, contigous from 0 to T * P - 1 343 | - `attr1`: [T * P, ...] attributes of the first mesh, where every P values correspond to a face 344 | _ ... 345 | """ 346 | def flatten_mesh_indices(*args: np.ndarray) -> Tuple[np.ndarray, ...]: 347 | assert len(args) % 2 == 0, "The number of arguments must be even." 348 | T, P = args[0].shape 349 | assert all(arg.shape[0] == T and arg.shape[1] == P for arg in args[::2]), "The faces must have the same shape." 350 | attr_flat = [] 351 | for faces_, attr_ in zip(args[::2], args[1::2]): 352 | attr_flat_ = attr_[faces_].reshape(-1, *attr_.shape[1:]) 353 | attr_flat.append(attr_flat_) 354 | faces_flat = np.arange(T * P, dtype=np.int32).reshape(T, P) 355 | return faces_flat, *attr_flat -------------------------------------------------------------------------------- /utils3d/numpy/rasterization.py: -------------------------------------------------------------------------------- 1 | import os 2 | from typing import * 3 | 4 | import numpy as np 5 | import moderngl 6 | 7 | from . import transforms, utils, mesh 8 | 9 | 10 | __all__ = [ 11 | 'RastContext', 12 | 'rasterize_triangle_faces', 13 | 'rasterize_edges', 14 | 'texture', 15 | 'test_rasterization', 16 | 'warp_image_by_depth', 17 | ] 18 | 19 | 20 | def map_np_dtype(dtype) -> str: 21 | if dtype == int: 22 | return 'i4' 23 | elif dtype == np.uint8: 24 | return 'u1' 25 | elif dtype == np.uint32: 26 | return 'u2' 27 | elif dtype == np.float16: 28 | return 'f2' 29 | elif dtype == np.float32: 30 | return 'f4' 31 | 32 | 33 | def one_value(dtype): 34 | if dtype == 'u1': 35 | return 255 36 | elif dtype == 'u2': 37 | return 65535 38 | else: 39 | return 1 40 | 41 | 42 | class RastContext: 43 | def __init__(self, *args, **kwargs): 44 | """ 45 | Create a moderngl context. 46 | 47 | Args: 48 | See moderngl.create_context 49 | """ 50 | if len(args) == 1 and isinstance(args[0], moderngl.Context): 51 | self.mgl_ctx = args[0] 52 | else: 53 | self.mgl_ctx = moderngl.create_context(*args, **kwargs) 54 | self.__prog_src = {} 55 | self.__prog = {} 56 | 57 | def program_vertex_attribute(self, n: int) -> moderngl.Program: 58 | assert n in [1, 2, 3, 4], 'vertex attribute only supports channels 1, 2, 3, 4' 59 | 60 | if 'vertex_attribute_vsh' not in self.__prog_src: 61 | with open(os.path.join(os.path.dirname(__file__), 'shaders', 'vertex_attribute.vsh'), 'r') as f: 62 | self.__prog_src['vertex_attribute_vsh'] = f.read() 63 | if 'vertex_attribute_fsh' not in self.__prog_src: 64 | with open(os.path.join(os.path.dirname(__file__), 'shaders', 'vertex_attribute.fsh'), 'r') as f: 65 | self.__prog_src['vertex_attribute_fsh'] = f.read() 66 | 67 | if f'vertex_attribute_{n}' not in self.__prog: 68 | vsh = self.__prog_src['vertex_attribute_vsh'].replace('vecN', f'vec{n}') 69 | fsh = self.__prog_src['vertex_attribute_fsh'].replace('vecN', f'vec{n}') 70 | self.__prog[f'vertex_attribute_{n}'] = self.mgl_ctx.program(vertex_shader=vsh, fragment_shader=fsh) 71 | 72 | return self.__prog[f'vertex_attribute_{n}'] 73 | 74 | def program_texture(self, n: int) -> moderngl.Program: 75 | assert n in [1, 2, 3, 4], 'texture only supports channels 1, 2, 3, 4' 76 | 77 | if 'texture_vsh' not in self.__prog_src: 78 | with open(os.path.join(os.path.dirname(__file__), 'shaders', 'texture.vsh'), 'r') as f: 79 | self.__prog_src['texture_vsh'] = f.read() 80 | if 'texture_fsh' not in self.__prog_src: 81 | with open(os.path.join(os.path.dirname(__file__), 'shaders', 'texture.fsh'), 'r') as f: 82 | self.__prog_src['texture_fsh'] = f.read() 83 | 84 | if f'texture_{n}' not in self.__prog: 85 | vsh = self.__prog_src['texture_vsh'].replace('vecN', f'vec{n}') 86 | fsh = self.__prog_src['texture_fsh'].replace('vecN', f'vec{n}') 87 | self.__prog[f'texture_{n}'] = self.mgl_ctx.program(vertex_shader=vsh, fragment_shader=fsh) 88 | self.__prog[f'texture_{n}']['tex'] = 0 89 | self.__prog[f'texture_{n}']['uv'] = 1 90 | 91 | return self.__prog[f'texture_{n}'] 92 | 93 | 94 | def rasterize_triangle_faces( 95 | ctx: RastContext, 96 | vertices: np.ndarray, 97 | faces: np.ndarray, 98 | attr: np.ndarray, 99 | width: int, 100 | height: int, 101 | transform: np.ndarray = None, 102 | cull_backface: bool = True, 103 | return_depth: bool = False, 104 | image: np.ndarray = None, 105 | depth: np.ndarray = None 106 | ) -> Tuple[np.ndarray, np.ndarray]: 107 | """ 108 | Rasterize vertex attribute. 109 | 110 | Args: 111 | vertices (np.ndarray): [N, 3] 112 | faces (np.ndarray): [T, 3] 113 | attr (np.ndarray): [N, C] 114 | width (int): width of rendered image 115 | height (int): height of rendered image 116 | transform (np.ndarray): [4, 4] model-view-projection transformation matrix. 117 | cull_backface (bool): whether to cull backface 118 | image: (np.ndarray): [H, W, C] background image 119 | depth: (np.ndarray): [H, W] background depth 120 | 121 | Returns: 122 | image (np.ndarray): [H, W, C] rendered image 123 | depth (np.ndarray): [H, W] screen space depth, ranging from 0 to 1. If return_depth is False, it is None. 124 | """ 125 | assert vertices.ndim == 2 and vertices.shape[1] == 3 126 | assert faces.ndim == 2 and faces.shape[1] == 3, f"Faces should be a 2D array with shape (T, 3), but got {faces.shape}" 127 | assert attr.ndim == 2 and attr.shape[1] in [1, 2, 3, 4], f'Vertex attribute only supports channels 1, 2, 3, 4, but got {attr.shape}' 128 | assert vertices.shape[0] == attr.shape[0] 129 | assert vertices.dtype == np.float32 130 | assert faces.dtype == np.uint32 or faces.dtype == np.int32 131 | assert attr.dtype == np.float32, "Attribute should be float32" 132 | assert transform is None or transform.shape == (4, 4), f"Transform should be a 4x4 matrix, but got {transform.shape}" 133 | assert transform is None or transform.dtype == np.float32, f"Transform should be float32, but got {transform.dtype}" 134 | if image is not None: 135 | assert image.ndim == 3 and image.shape == (height, width, attr.shape[1]), f"Image should be a 3D array with shape (H, W, {attr.shape[1]}), but got {image.shape}" 136 | assert image.dtype == np.float32, f"Image should be float32, but got {image.dtype}" 137 | if depth is not None: 138 | assert depth.ndim == 2 and depth.shape == (height, width), f"Depth should be a 2D array with shape (H, W), but got {depth.shape}" 139 | assert depth.dtype == np.float32, f"Depth should be float32, but got {depth.dtype}" 140 | 141 | C = attr.shape[1] 142 | prog = ctx.program_vertex_attribute(C) 143 | 144 | transform = np.eye(4, np.float32) if transform is None else transform 145 | 146 | # Create buffers 147 | ibo = ctx.mgl_ctx.buffer(np.ascontiguousarray(faces, dtype='i4')) 148 | vbo_vertices = ctx.mgl_ctx.buffer(np.ascontiguousarray(vertices, dtype='f4')) 149 | vbo_attr = ctx.mgl_ctx.buffer(np.ascontiguousarray(attr, dtype='f4')) 150 | vao = ctx.mgl_ctx.vertex_array( 151 | prog, 152 | [ 153 | (vbo_vertices, '3f', 'i_position'), 154 | (vbo_attr, f'{C}f', 'i_attr'), 155 | ], 156 | ibo, 157 | mode=moderngl.TRIANGLES, 158 | ) 159 | 160 | # Create framebuffer 161 | image_tex = ctx.mgl_ctx.texture((width, height), C, dtype='f4', data=np.ascontiguousarray(image[::-1, :, :]) if image is not None else None) 162 | depth_tex = ctx.mgl_ctx.depth_texture((width, height), data=np.ascontiguousarray(depth[::-1, :]) if depth is not None else None) 163 | fbo = ctx.mgl_ctx.framebuffer( 164 | color_attachments=[image_tex], 165 | depth_attachment=depth_tex, 166 | ) 167 | 168 | # Render 169 | prog['u_mvp'].write(transform.transpose().copy().astype('f4')) 170 | fbo.use() 171 | fbo.viewport = (0, 0, width, height) 172 | ctx.mgl_ctx.depth_func = '<' 173 | if depth is None: 174 | ctx.mgl_ctx.clear(depth=1.0) 175 | ctx.mgl_ctx.enable(ctx.mgl_ctx.DEPTH_TEST) 176 | if cull_backface: 177 | ctx.mgl_ctx.enable(ctx.mgl_ctx.CULL_FACE) 178 | else: 179 | ctx.mgl_ctx.disable(ctx.mgl_ctx.CULL_FACE) 180 | vao.render() 181 | ctx.mgl_ctx.disable(ctx.mgl_ctx.DEPTH_TEST) 182 | 183 | # Read 184 | image = np.zeros((height, width, C), dtype='f4') 185 | image_tex.read_into(image) 186 | image = image[::-1, :, :] 187 | if return_depth: 188 | depth = np.zeros((height, width), dtype='f4') 189 | depth_tex.read_into(depth) 190 | depth = depth[::-1, :] 191 | else: 192 | depth = None 193 | 194 | # Release 195 | vao.release() 196 | ibo.release() 197 | vbo_vertices.release() 198 | vbo_attr.release() 199 | fbo.release() 200 | image_tex.release() 201 | depth_tex.release() 202 | 203 | return image, depth 204 | 205 | 206 | def rasterize_edges( 207 | ctx: RastContext, 208 | vertices: np.ndarray, 209 | edges: np.ndarray, 210 | attr: np.ndarray, 211 | width: int, 212 | height: int, 213 | transform: np.ndarray = None, 214 | line_width: float = 1.0, 215 | return_depth: bool = False, 216 | image: np.ndarray = None, 217 | depth: np.ndarray = None 218 | ) -> Tuple[np.ndarray, ...]: 219 | """ 220 | Rasterize vertex attribute. 221 | 222 | Args: 223 | vertices (np.ndarray): [N, 3] 224 | faces (np.ndarray): [T, 3] 225 | attr (np.ndarray): [N, C] 226 | width (int): width of rendered image 227 | height (int): height of rendered image 228 | transform (np.ndarray): [4, 4] model-view-projection matrix 229 | line_width (float): width of line. Defaults to 1.0. NOTE: Values other than 1.0 may not work across all platforms. 230 | cull_backface (bool): whether to cull backface 231 | 232 | Returns: 233 | image (np.ndarray): [H, W, C] rendered image 234 | depth (np.ndarray): [H, W] screen space depth, ranging from 0 to 1. If return_depth is False, it is None. 235 | """ 236 | assert vertices.ndim == 2 and vertices.shape[1] == 3 237 | assert edges.ndim == 2 and edges.shape[1] == 2, f"Edges should be a 2D array with shape (T, 2), but got {edges.shape}" 238 | assert attr.ndim == 2 and attr.shape[1] in [1, 2, 3, 4], f'Vertex attribute only supports channels 1, 2, 3, 4, but got {attr.shape}' 239 | assert vertices.shape[0] == attr.shape[0] 240 | assert vertices.dtype == np.float32 241 | assert edges.dtype == np.uint32 or edges.dtype == np.int32 242 | assert attr.dtype == np.float32, "Attribute should be float32" 243 | 244 | C = attr.shape[1] 245 | prog = ctx.program_vertex_attribute(C) 246 | 247 | transform = transform if transform is not None else np.eye(4, np.float32) 248 | 249 | # Create buffers 250 | ibo = ctx.mgl_ctx.buffer(np.ascontiguousarray(edges, dtype='i4')) 251 | vbo_vertices = ctx.mgl_ctx.buffer(np.ascontiguousarray(vertices, dtype='f4')) 252 | vbo_attr = ctx.mgl_ctx.buffer(np.ascontiguousarray(attr, dtype='f4')) 253 | vao = ctx.mgl_ctx.vertex_array( 254 | prog, 255 | [ 256 | (vbo_vertices, '3f', 'i_position'), 257 | (vbo_attr, f'{C}f', 'i_attr'), 258 | ], 259 | ibo, 260 | mode=moderngl.LINES, 261 | ) 262 | 263 | # Create framebuffer 264 | image_tex = ctx.mgl_ctx.texture((width, height), C, dtype='f4', data=np.ascontiguousarray(image[::-1, :, :]) if image is not None else None) 265 | depth_tex = ctx.mgl_ctx.depth_texture((width, height), data=np.ascontiguousarray(depth[::-1, :]) if depth is not None else None) 266 | fbo = ctx.mgl_ctx.framebuffer( 267 | color_attachments=[image_tex], 268 | depth_attachment=depth_tex, 269 | ) 270 | 271 | # Render 272 | prog['u_mvp'].write(transform.transpose().copy().astype('f4')) 273 | fbo.use() 274 | fbo.viewport = (0, 0, width, height) 275 | if depth is None: 276 | ctx.mgl_ctx.clear(depth=1.0) 277 | ctx.mgl_ctx.depth_func = '<' 278 | ctx.mgl_ctx.enable(ctx.mgl_ctx.DEPTH_TEST) 279 | ctx.mgl_ctx.line_width = line_width 280 | vao.render() 281 | ctx.mgl_ctx.disable(ctx.mgl_ctx.DEPTH_TEST) 282 | 283 | # Read 284 | image = np.zeros((height, width, C), dtype='f4') 285 | image_tex.read_into(image) 286 | image = image[::-1, :, :] 287 | if return_depth: 288 | depth = np.zeros((height, width), dtype='f4') 289 | depth_tex.read_into(depth) 290 | depth = depth[::-1, :] 291 | else: 292 | depth = None 293 | 294 | # Release 295 | vao.release() 296 | ibo.release() 297 | vbo_vertices.release() 298 | vbo_attr.release() 299 | fbo.release() 300 | image_tex.release() 301 | depth_tex.release() 302 | 303 | return image, depth 304 | 305 | 306 | def texture( 307 | ctx: RastContext, 308 | uv: np.ndarray, 309 | texture: np.ndarray, 310 | interpolation: str= 'linear', 311 | wrap: str = 'clamp' 312 | ) -> np.ndarray: 313 | """ 314 | Given an UV image, texturing from the texture map 315 | """ 316 | assert len(texture.shape) == 3 and 1 <= texture.shape[2] <= 4 317 | assert uv.shape[2] == 2 318 | height, width = uv.shape[:2] 319 | texture_dtype = map_np_dtype(texture.dtype) 320 | 321 | # Create VAO 322 | screen_quad_vbo = ctx.mgl_ctx.buffer(np.array([[-1, -1], [1, -1], [1, 1], [-1, 1]], dtype='f4')) 323 | screen_quad_ibo = ctx.mgl_ctx.buffer(np.array([0, 1, 2, 0, 2, 3], dtype=np.int32)) 324 | screen_quad_vao = ctx.mgl_ctx.vertex_array(ctx.program_texture(texture.shape[2]), [(screen_quad_vbo, '2f4', 'in_vert')], index_buffer=screen_quad_ibo, index_element_size=4) 325 | 326 | # Create texture, set filter and bind. TODO: min mag filter, mipmap 327 | texture_tex = ctx.mgl_ctx.texture((texture.shape[1], texture.shape[0]), texture.shape[2], dtype=texture_dtype, data=np.ascontiguousarray(texture)) 328 | if interpolation == 'linear': 329 | texture_tex.filter = (moderngl.LINEAR, moderngl.LINEAR) 330 | elif interpolation == 'nearest': 331 | texture_tex.filter = (moderngl.NEAREST, moderngl.NEAREST) 332 | texture_tex.use(location=0) 333 | texture_uv = ctx.mgl_ctx.texture((width, height), 2, dtype='f4', data=np.ascontiguousarray(uv.astype('f4', copy=False))) 334 | texture_uv.filter = (moderngl.NEAREST, moderngl.NEAREST) 335 | texture_uv.use(location=1) 336 | 337 | # Create render buffer and frame buffer 338 | rb = ctx.mgl_ctx.renderbuffer((uv.shape[1], uv.shape[0]), texture.shape[2], dtype=texture_dtype) 339 | fbo = ctx.mgl_ctx.framebuffer(color_attachments=[rb]) 340 | 341 | # Render 342 | fbo.use() 343 | fbo.viewport = (0, 0, width, height) 344 | ctx.mgl_ctx.disable(ctx.mgl_ctx.BLEND) 345 | screen_quad_vao.render() 346 | 347 | # Read buffer 348 | image_buffer = np.frombuffer(fbo.read(components=texture.shape[2], attachment=0, dtype=texture_dtype), dtype=texture_dtype).reshape((height, width, texture.shape[2])) 349 | 350 | # Release 351 | texture_tex.release() 352 | rb.release() 353 | fbo.release() 354 | 355 | return image_buffer 356 | 357 | 358 | def warp_image_by_depth( 359 | ctx: RastContext, 360 | src_depth: np.ndarray, 361 | src_image: np.ndarray = None, 362 | width: int = None, 363 | height: int = None, 364 | *, 365 | extrinsics_src: np.ndarray = None, 366 | extrinsics_tgt: np.ndarray = None, 367 | intrinsics_src: np.ndarray = None, 368 | intrinsics_tgt: np.ndarray = None, 369 | near: float = 0.1, 370 | far: float = 100.0, 371 | cull_backface: bool = True, 372 | ssaa: int = 1, 373 | return_depth: bool = False, 374 | ) -> Tuple[np.ndarray, ...]: 375 | """ 376 | Warp image by depth map. 377 | 378 | Args: 379 | ctx (RastContext): rasterizer context 380 | src_depth (np.ndarray): [H, W] 381 | src_image (np.ndarray, optional): [H, W, C]. The image to warp. Defaults to None (use uv coordinates). 382 | width (int, optional): width of the output image. None to use depth map width. Defaults to None. 383 | height (int, optional): height of the output image. None to use depth map height. Defaults to None. 384 | extrinsics_src (np.ndarray, optional): extrinsics matrix of the source camera. Defaults to None (identity). 385 | extrinsics_tgt (np.ndarray, optional): extrinsics matrix of the target camera. Defaults to None (identity). 386 | intrinsics_src (np.ndarray, optional): intrinsics matrix of the source camera. Defaults to None (use the same as intrinsics_tgt). 387 | intrinsics_tgt (np.ndarray, optional): intrinsics matrix of the target camera. Defaults to None (use the same as intrinsics_src). 388 | cull_backface (bool, optional): whether to cull backface. Defaults to True. 389 | ssaa (int, optional): super sampling anti-aliasing. Defaults to 1. 390 | 391 | Returns: 392 | tgt_image (np.ndarray): [H, W, C] warped image (or uv coordinates if image is None). 393 | tgt_depth (np.ndarray): [H, W] screen space depth, ranging from 0 to 1. If return_depth is False, it is None. 394 | """ 395 | assert src_depth.ndim == 2 396 | 397 | if width is None: 398 | width = src_depth.shape[1] 399 | if height is None: 400 | height = src_depth.shape[0] 401 | if src_image is not None: 402 | assert src_image.shape[-2:] == src_depth.shape[-2:], f'Shape of source image {src_image.shape} does not match shape of source depth {src_depth.shape}' 403 | 404 | # set up default camera parameters 405 | extrinsics_src = np.eye(4) if extrinsics_src is None else extrinsics_src 406 | extrinsics_tgt = np.eye(4) if extrinsics_tgt is None else extrinsics_tgt 407 | intrinsics_src = intrinsics_tgt if intrinsics_src is None else intrinsics_src 408 | intrinsics_tgt = intrinsics_src if intrinsics_tgt is None else intrinsics_tgt 409 | 410 | assert all(x is not None for x in [extrinsics_src, extrinsics_tgt, intrinsics_src, intrinsics_tgt]), "Make sure you have provided all the necessary camera parameters." 411 | 412 | # check shapes 413 | assert extrinsics_src.shape == (4, 4) and extrinsics_tgt.shape == (4, 4) 414 | assert intrinsics_src.shape == (3, 3) and intrinsics_tgt.shape == (3, 3) 415 | 416 | # convert to view and perspective matrices 417 | view_tgt = transforms.extrinsics_to_view(extrinsics_tgt) 418 | perspective_tgt = transforms.intrinsics_to_perspective(intrinsics_tgt, near=near, far=far) 419 | 420 | # unproject depth map 421 | uv, faces = utils.image_mesh(*src_depth.shape[-2:]) 422 | pts = transforms.unproject_cv(uv, src_depth.reshape(-1), extrinsics_src, intrinsics_src) 423 | faces = mesh.triangulate(faces, vertices=pts) 424 | 425 | # rasterize attributes 426 | if src_image is not None: 427 | attr = src_image.reshape(-1, src_image.shape[-1]) 428 | else: 429 | attr = uv 430 | 431 | tgt_image, tgt_depth = rasterize_triangle_faces( 432 | ctx, 433 | pts, 434 | faces, 435 | attr, 436 | width * ssaa, 437 | height * ssaa, 438 | transform=perspective_tgt @ view_tgt, 439 | cull_backface=cull_backface, 440 | return_depth=return_depth, 441 | ) 442 | 443 | if ssaa > 1: 444 | tgt_image = tgt_image.reshape(height, ssaa, width, ssaa, -1).mean(axis=(1, 3)) 445 | tgt_depth = tgt_depth.reshape(height, ssaa, width, ssaa, -1).mean(axis=(1, 3)) if return_depth else None 446 | 447 | return tgt_image, tgt_depth 448 | 449 | def test_rasterization(ctx: RastContext): 450 | """ 451 | Test if rasterization works. It will render a cube with random colors and save it as a CHECKME.png file. 452 | """ 453 | vertices, faces = utils.cube(tri=True) 454 | attr = np.random.rand(len(vertices), 3).astype(np.float32) 455 | perspective = transforms.perspective(np.deg2rad(60), 1, 0.01, 100) 456 | view = transforms.view_look_at(np.array([2, 2, 2]), np.array([0, 0, 0]), np.array([0, 1, 0])) 457 | image, depth = rasterize_triangle_faces( 458 | ctx, 459 | vertices, 460 | faces, 461 | attr, 462 | 512, 512, 463 | transform=(perspective @ view).astype(np.float32), 464 | cull_backface=False, 465 | return_depth=True, 466 | ) 467 | import cv2 468 | cv2.imwrite('CHECKME.png', cv2.cvtColor((image.clip(0, 1) * 255).astype(np.uint8), cv2.COLOR_RGB2BGR)) 469 | -------------------------------------------------------------------------------- /utils3d/numpy/shaders/texture.fsh: -------------------------------------------------------------------------------- 1 | #version 330 2 | 3 | uniform sampler2D tex; 4 | uniform sampler2D uv; 5 | 6 | in vec2 scr_coord; 7 | out vecN tex_color; 8 | 9 | void main() { 10 | tex_color = vecN(texture(tex, texture(uv, scr_coord).xy)); 11 | } -------------------------------------------------------------------------------- /utils3d/numpy/shaders/texture.vsh: -------------------------------------------------------------------------------- 1 | #version 330 core 2 | 3 | in vec2 in_vert; 4 | out vec2 scr_coord; 5 | 6 | void main() { 7 | scr_coord = in_vert * 0.5 + 0.5; 8 | gl_Position = vec4(in_vert, 0., 1.); 9 | } -------------------------------------------------------------------------------- /utils3d/numpy/shaders/vertex_attribute.fsh: -------------------------------------------------------------------------------- 1 | #version 330 2 | 3 | in vecN v_attr; 4 | 5 | out vecN f_attr; 6 | 7 | void main() { 8 | f_attr = v_attr; 9 | } 10 | -------------------------------------------------------------------------------- /utils3d/numpy/shaders/vertex_attribute.vsh: -------------------------------------------------------------------------------- 1 | #version 330 2 | 3 | uniform mat4 u_mvp; 4 | 5 | in vec3 i_position; 6 | in vecN i_attr; 7 | 8 | out vecN v_attr; 9 | 10 | void main() { 11 | gl_Position = u_mvp * vec4(i_position, 1.0); 12 | v_attr = i_attr; 13 | } 14 | -------------------------------------------------------------------------------- /utils3d/numpy/spline.py: -------------------------------------------------------------------------------- 1 | from typing import * 2 | 3 | import numpy as np 4 | 5 | 6 | __all__ = ['linear_spline_interpolate'] 7 | 8 | 9 | def linear_spline_interpolate(x: np.ndarray, t: np.ndarray, s: np.ndarray, extrapolation_mode: Literal['constant', 'linear'] = 'constant') -> np.ndarray: 10 | """ 11 | Linear spline interpolation. 12 | 13 | ### Parameters: 14 | - `x`: np.ndarray, shape (n, d): the values of data points. 15 | - `t`: np.ndarray, shape (n,): the times of the data points. 16 | - `s`: np.ndarray, shape (m,): the times to be interpolated. 17 | - `extrapolation_mode`: str, the mode of extrapolation. 'constant' means extrapolate the boundary values, 'linear' means extrapolate linearly. 18 | 19 | ### Returns: 20 | - `y`: np.ndarray, shape (..., m, d): the interpolated values. 21 | """ 22 | i = np.searchsorted(t, s, side='left') 23 | if extrapolation_mode == 'constant': 24 | prev = np.clip(i - 1, 0, len(t) - 1) 25 | suc = np.clip(i, 0, len(t) - 1) 26 | elif extrapolation_mode == 'linear': 27 | prev = np.clip(i - 1, 0, len(t) - 2) 28 | suc = np.clip(i, 1, len(t) - 1) 29 | else: 30 | raise ValueError(f'Invalid extrapolation_mode: {extrapolation_mode}') 31 | 32 | u = (s - t[prev]) / np.maximum(t[suc] - t[prev], 1e-12) 33 | y = u * x[suc] + (1 - u) * x[prev] 34 | 35 | return y 36 | 37 | 38 | 39 | def _solve_tridiagonal(a: np.ndarray, b: np.ndarray, c: np.ndarray, d: np.ndarray) -> np.ndarray: 40 | n = b.shape[-1] 41 | cc = np.zeros_like(b) 42 | dd = np.zeros_like(b) 43 | cc[..., 0] = c[..., 0] / b[..., 0] 44 | dd[..., 0] = d[..., 0] / b[..., 0] 45 | for i in range(1, n): 46 | cc[..., i] = c[..., i] / (b[..., i] - a[..., i - 1] * cc[..., i - 1]) 47 | dd[..., i] = (d[..., i] - a[..., i - 1] * dd[..., i - 1]) / (b[..., i] - a[..., i - 1] * cc[..., i - 1]) 48 | x = np.zeros_like(b) 49 | x[..., -1] = dd[..., -1] 50 | for i in range(n - 2, -1, -1): 51 | x[..., i] = dd[..., i] - cc[..., i] * x[..., i + 1] 52 | return x 53 | 54 | 55 | def cubic_spline_interpolate(x: np.ndarray, t: np.ndarray, s: np.ndarray, v0: np.ndarray = None, vn: np.ndarray = None) -> np.ndarray: 56 | """ 57 | Cubic spline interpolation. 58 | 59 | ### Parameters: 60 | - `x`: np.ndarray, shape (..., n,): the x-coordinates of the data points. 61 | - `t`: np.ndarray, shape (n,): the knot vector. NOTE: t must be sorted in ascending order. 62 | - `s`: np.ndarray, shape (..., m,): the y-coordinates of the data points. 63 | - `v0`: np.ndarray, shape (...,): the value of the derivative at the first knot, as the boundary condition. If None, it is set to zero. 64 | - `vn`: np.ndarray, shape (...,): the value of the derivative at the last knot, as the boundary condition. If None, it is set to zero. 65 | 66 | ### Returns: 67 | - `y`: np.ndarray, shape (..., m): the interpolated values. 68 | """ 69 | h = t[..., 1:] - t[..., :-1] 70 | mu = h[..., :-1] / (h[..., :-1] + h[..., 1:]) 71 | la = 1 - mu 72 | d = (x[..., 1:] - x[..., :-1]) / h 73 | d = 6 * (d[..., 1:] - d[..., :-1]) / (t[..., 2:] - t[..., :-2]) 74 | 75 | mu = np.concatenate([mu, np.ones_like(mu[..., :1])], axis=-1) 76 | la = np.concatenate([np.ones_like(la[..., :1]), la], axis=-1) 77 | d = np.concatenate([(((x[..., 1] - x[..., 0]) / h[0] - v0) / h[0])[..., None], d, ((vn - (x[..., -1] - x[..., -2]) / h[-1]) / h[-1])[..., None]], axis=-1) 78 | 79 | M = _solve_tridiagonal(mu, np.full_like(d, fill_value=2), la, d) 80 | 81 | i = np.searchsorted(t, s, side='left') 82 | 83 | -------------------------------------------------------------------------------- /utils3d/torch/__init__.py: -------------------------------------------------------------------------------- 1 | import importlib 2 | import itertools 3 | import torch 4 | from typing import TYPE_CHECKING 5 | 6 | __modules_all__ = { 7 | 'mesh': [ 8 | 'triangulate', 9 | 'compute_face_normal', 10 | 'compute_face_angles', 11 | 'compute_vertex_normal', 12 | 'compute_vertex_normal_weighted', 13 | 'compute_edges', 14 | 'compute_connected_components', 15 | 'compute_edge_connected_components', 16 | 'compute_boundarys', 17 | 'compute_dual_graph', 18 | 'remove_unreferenced_vertices', 19 | 'remove_corrupted_faces', 20 | 'remove_isolated_pieces', 21 | 'merge_duplicate_vertices', 22 | 'subdivide_mesh_simple', 23 | 'compute_face_tbn', 24 | 'compute_vertex_tbn', 25 | 'laplacian', 26 | 'laplacian_smooth_mesh', 27 | 'taubin_smooth_mesh', 28 | 'laplacian_hc_smooth_mesh', 29 | ], 30 | 'nerf': [ 31 | 'get_rays', 32 | 'get_image_rays', 33 | 'get_mipnerf_cones', 34 | 'volume_rendering', 35 | 'bin_sample', 36 | 'importance_sample', 37 | 'nerf_render_rays', 38 | 'mipnerf_render_rays', 39 | 'nerf_render_view', 40 | 'mipnerf_render_view', 41 | 'InstantNGP', 42 | ], 43 | 'utils': [ 44 | 'sliding_window_1d', 45 | 'sliding_window_2d', 46 | 'sliding_window_nd', 47 | 'image_uv', 48 | 'image_pixel_center', 49 | 'image_mesh', 50 | 'chessboard', 51 | 'depth_edge', 52 | 'depth_aliasing', 53 | 'image_mesh_from_depth', 54 | 'points_to_normals', 55 | 'depth_to_points', 56 | 'depth_to_normals', 57 | 'masked_min', 58 | 'masked_max', 59 | 'bounding_rect' 60 | ], 61 | 'transforms': [ 62 | 'perspective', 63 | 'perspective_from_fov', 64 | 'perspective_from_fov_xy', 65 | 'intrinsics_from_focal_center', 66 | 'intrinsics_from_fov', 67 | 'intrinsics_from_fov_xy', 68 | 'focal_to_fov', 69 | 'fov_to_focal', 70 | 'intrinsics_to_fov', 71 | 'view_look_at', 72 | 'extrinsics_look_at', 73 | 'perspective_to_intrinsics', 74 | 'intrinsics_to_perspective', 75 | 'extrinsics_to_view', 76 | 'view_to_extrinsics', 77 | 'normalize_intrinsics', 78 | 'crop_intrinsics', 79 | 'pixel_to_uv', 80 | 'pixel_to_ndc', 81 | 'uv_to_pixel', 82 | 'project_depth', 83 | 'depth_buffer_to_linear', 84 | 'project_gl', 85 | 'project_cv', 86 | 'unproject_gl', 87 | 'unproject_cv', 88 | 'skew_symmetric', 89 | 'rotation_matrix_from_vectors', 90 | 'euler_axis_angle_rotation', 91 | 'euler_angles_to_matrix', 92 | 'matrix_to_euler_angles', 93 | 'matrix_to_quaternion', 94 | 'quaternion_to_matrix', 95 | 'matrix_to_axis_angle', 96 | 'axis_angle_to_matrix', 97 | 'axis_angle_to_quaternion', 98 | 'quaternion_to_axis_angle', 99 | 'slerp', 100 | 'interpolate_extrinsics', 101 | 'interpolate_view', 102 | 'extrinsics_to_essential', 103 | 'to4x4', 104 | 'rotation_matrix_2d', 105 | 'rotate_2d', 106 | 'translate_2d', 107 | 'scale_2d', 108 | 'apply_2d', 109 | ], 110 | 'rasterization': [ 111 | 'RastContext', 112 | 'rasterize_triangle_faces', 113 | 'rasterize_triangle_faces_depth_peeling', 114 | 'texture', 115 | 'texture_composite', 116 | 'warp_image_by_depth', 117 | 'warp_image_by_forward_flow', 118 | ], 119 | } 120 | 121 | 122 | __all__ = list(itertools.chain(*__modules_all__.values())) 123 | 124 | def __getattr__(name): 125 | try: 126 | return globals()[name] 127 | except KeyError: 128 | pass 129 | 130 | try: 131 | module_name = next(m for m in __modules_all__ if name in __modules_all__[m]) 132 | except StopIteration: 133 | raise AttributeError(f"module '{__name__}' has no attribute '{name}'") 134 | module = importlib.import_module(f'.{module_name}', __name__) 135 | for key in __modules_all__[module_name]: 136 | globals()[key] = getattr(module, key) 137 | 138 | return globals()[name] 139 | 140 | 141 | if TYPE_CHECKING: 142 | from .transforms import * 143 | from .mesh import * 144 | from .utils import * 145 | from .nerf import * 146 | from .rasterization import * -------------------------------------------------------------------------------- /utils3d/torch/_helpers.py: -------------------------------------------------------------------------------- 1 | # decorator 2 | import torch 3 | from numbers import Number 4 | import inspect 5 | from functools import wraps 6 | from .._helpers import suppress_traceback 7 | 8 | 9 | def get_device(args, kwargs): 10 | device = None 11 | for arg in (list(args) + list(kwargs.values())): 12 | if isinstance(arg, torch.Tensor): 13 | if device is None: 14 | device = arg.device 15 | elif device != arg.device: 16 | raise ValueError("All tensors must be on the same device.") 17 | return device 18 | 19 | 20 | def get_args_order(func, args, kwargs): 21 | """ 22 | Get the order of the arguments of a function. 23 | """ 24 | names = inspect.getfullargspec(func).args 25 | names_idx = {name: i for i, name in enumerate(names)} 26 | args_order = [] 27 | kwargs_order = {} 28 | for name, arg in kwargs.items(): 29 | if name in names: 30 | kwargs_order[name] = names_idx[name] 31 | names.remove(name) 32 | for i, arg in enumerate(args): 33 | if i < len(names): 34 | args_order.append(names_idx[names[i]]) 35 | return args_order, kwargs_order 36 | 37 | 38 | def broadcast_args(args, kwargs, args_dim, kwargs_dim): 39 | spatial = [] 40 | for arg, arg_dim in zip(args + list(kwargs.values()), args_dim + list(kwargs_dim.values())): 41 | if isinstance(arg, torch.Tensor) and arg_dim is not None: 42 | arg_spatial = arg.shape[:arg.ndim-arg_dim] 43 | if len(arg_spatial) > len(spatial): 44 | spatial = [1] * (len(arg_spatial) - len(spatial)) + spatial 45 | for j in range(len(arg_spatial)): 46 | if spatial[-j] < arg_spatial[-j]: 47 | if spatial[-j] == 1: 48 | spatial[-j] = arg_spatial[-j] 49 | else: 50 | raise ValueError("Cannot broadcast arguments.") 51 | for i, arg in enumerate(args): 52 | if isinstance(arg, torch.Tensor) and args_dim[i] is not None: 53 | args[i] = torch.broadcast_to(arg, [*spatial, *arg.shape[arg.ndim-args_dim[i]:]]) 54 | for key, arg in kwargs.items(): 55 | if isinstance(arg, torch.Tensor) and kwargs_dim[key] is not None: 56 | kwargs[key] = torch.broadcast_to(arg, [*spatial, *arg.shape[arg.ndim-kwargs_dim[key]:]]) 57 | return args, kwargs, spatial 58 | 59 | @suppress_traceback 60 | def batched(*dims): 61 | """ 62 | Decorator that allows a function to be called with batched arguments. 63 | """ 64 | def decorator(func): 65 | @wraps(func) 66 | def wrapper(*args, device=torch.device('cpu'), **kwargs): 67 | args = list(args) 68 | # get arguments dimensions 69 | args_order, kwargs_order = get_args_order(func, args, kwargs) 70 | args_dim = [dims[i] for i in args_order] 71 | kwargs_dim = {key: dims[i] for key, i in kwargs_order.items()} 72 | # convert to torch tensor 73 | device = get_device(args, kwargs) or device 74 | for i, arg in enumerate(args): 75 | if isinstance(arg, (Number, list, tuple)) and args_dim[i] is not None: 76 | args[i] = torch.tensor(arg, device=device) 77 | for key, arg in kwargs.items(): 78 | if isinstance(arg, (Number, list, tuple)) and kwargs_dim[key] is not None: 79 | kwargs[key] = torch.tensor(arg, device=device) 80 | # broadcast arguments 81 | args, kwargs, spatial = broadcast_args(args, kwargs, args_dim, kwargs_dim) 82 | for i, (arg, arg_dim) in enumerate(zip(args, args_dim)): 83 | if isinstance(arg, torch.Tensor) and arg_dim is not None: 84 | args[i] = arg.reshape([-1, *arg.shape[arg.ndim-arg_dim:]]) 85 | for key, arg in kwargs.items(): 86 | if isinstance(arg, torch.Tensor) and kwargs_dim[key] is not None: 87 | kwargs[key] = arg.reshape([-1, *arg.shape[arg.ndim-kwargs_dim[key]:]]) 88 | # call function 89 | results = func(*args, **kwargs) 90 | type_results = type(results) 91 | results = list(results) if isinstance(results, (tuple, list)) else [results] 92 | # restore spatial dimensions 93 | for i, result in enumerate(results): 94 | results[i] = result.reshape([*spatial, *result.shape[1:]]) 95 | if type_results == tuple: 96 | results = tuple(results) 97 | elif type_results == list: 98 | results = list(results) 99 | else: 100 | results = results[0] 101 | return results 102 | return wrapper 103 | return decorator -------------------------------------------------------------------------------- /utils3d/torch/utils.py: -------------------------------------------------------------------------------- 1 | from typing import * 2 | 3 | import torch 4 | import torch.nn.functional as F 5 | 6 | from . import transforms 7 | from . import mesh 8 | from ._helpers import batched 9 | from .._helpers import no_warnings 10 | 11 | 12 | __all__ = [ 13 | 'sliding_window_1d', 14 | 'sliding_window_2d', 15 | 'sliding_window_nd', 16 | 'image_uv', 17 | 'image_pixel_center', 18 | 'image_mesh', 19 | 'chessboard', 20 | 'depth_edge', 21 | 'depth_aliasing', 22 | 'image_mesh_from_depth', 23 | 'points_to_normals', 24 | 'depth_to_points', 25 | 'depth_to_normals', 26 | 'masked_min', 27 | 'masked_max', 28 | 'bounding_rect' 29 | ] 30 | 31 | 32 | def sliding_window_1d(x: torch.Tensor, window_size: int, stride: int = 1, dim: int = -1) -> torch.Tensor: 33 | """ 34 | Sliding window view of the input tensor. The dimension of the sliding window is appended to the end of the input tensor's shape. 35 | NOTE: Since Pytorch has `unfold` function, 1D sliding window view is just a wrapper of it. 36 | """ 37 | return x.unfold(dim, window_size, stride) 38 | 39 | 40 | def sliding_window_nd(x: torch.Tensor, window_size: Tuple[int, ...], stride: Tuple[int, ...], dim: Tuple[int, ...]) -> torch.Tensor: 41 | dim = [dim[i] % x.ndim for i in range(len(dim))] 42 | assert len(window_size) == len(stride) == len(dim) 43 | for i in range(len(window_size)): 44 | x = sliding_window_1d(x, window_size[i], stride[i], dim[i]) 45 | return x 46 | 47 | 48 | def sliding_window_2d(x: torch.Tensor, window_size: Union[int, Tuple[int, int]], stride: Union[int, Tuple[int, int]], dim: Union[int, Tuple[int, int]] = (-2, -1)) -> torch.Tensor: 49 | if isinstance(window_size, int): 50 | window_size = (window_size, window_size) 51 | if isinstance(stride, int): 52 | stride = (stride, stride) 53 | return sliding_window_nd(x, window_size, stride, dim) 54 | 55 | 56 | def image_uv(height: int, width: int, left: int = None, top: int = None, right: int = None, bottom: int = None, device: torch.device = None, dtype: torch.dtype = None) -> torch.Tensor: 57 | """ 58 | Get image space UV grid, ranging in [0, 1]. 59 | 60 | >>> image_uv(10, 10): 61 | [[[0.05, 0.05], [0.15, 0.05], ..., [0.95, 0.05]], 62 | [[0.05, 0.15], [0.15, 0.15], ..., [0.95, 0.15]], 63 | ... ... ... 64 | [[0.05, 0.95], [0.15, 0.95], ..., [0.95, 0.95]]] 65 | 66 | Args: 67 | width (int): image width 68 | height (int): image height 69 | 70 | Returns: 71 | torch.Tensor: shape (height, width, 2) 72 | """ 73 | if left is None: left = 0 74 | if top is None: top = 0 75 | if right is None: right = width 76 | if bottom is None: bottom = height 77 | u = torch.linspace((left + 0.5) / width, (right - 0.5) / width, right - left, device=device, dtype=dtype) 78 | v = torch.linspace((top + 0.5) / height, (bottom - 0.5) / height, bottom - top, device=device, dtype=dtype) 79 | u, v = torch.meshgrid(u, v, indexing='xy') 80 | uv = torch.stack([u, v], dim=-1) 81 | return uv 82 | 83 | 84 | def image_pixel_center( 85 | height: int, 86 | width: int, 87 | left: int = None, 88 | top: int = None, 89 | right: int = None, 90 | bottom: int = None, 91 | dtype: torch.dtype = None, 92 | device: torch.device = None 93 | ) -> torch.Tensor: 94 | """ 95 | Get image pixel center coordinates, ranging in [0, width] and [0, height]. 96 | `image[i, j]` has pixel center coordinates `(j + 0.5, i + 0.5)`. 97 | 98 | >>> image_pixel_center(10, 10): 99 | [[[0.5, 0.5], [1.5, 0.5], ..., [9.5, 0.5]], 100 | [[0.5, 1.5], [1.5, 1.5], ..., [9.5, 1.5]], 101 | ... ... ... 102 | [[0.5, 9.5], [1.5, 9.5], ..., [9.5, 9.5]]] 103 | 104 | Args: 105 | width (int): image width 106 | height (int): image height 107 | 108 | Returns: 109 | torch.Tensor: shape (height, width, 2) 110 | """ 111 | if left is None: left = 0 112 | if top is None: top = 0 113 | if right is None: right = width 114 | if bottom is None: bottom = height 115 | u = torch.linspace(left + 0.5, right - 0.5, right - left, dtype=dtype, device=device) 116 | v = torch.linspace(top + 0.5, bottom - 0.5, bottom - top, dtype=dtype, device=device) 117 | u, v = torch.meshgrid(u, v, indexing='xy') 118 | return torch.stack([u, v], dim=2) 119 | 120 | 121 | def image_mesh(height: int, width: int, mask: torch.Tensor = None, device: torch.device = None, dtype: torch.dtype = None) -> Tuple[torch.Tensor, torch.Tensor]: 122 | """ 123 | Get a quad mesh regarding image pixel uv coordinates as vertices and image grid as faces. 124 | 125 | Args: 126 | width (int): image width 127 | height (int): image height 128 | mask (torch.Tensor, optional): binary mask of shape (height, width), dtype=bool. Defaults to None. 129 | 130 | Returns: 131 | uv (torch.Tensor): uv corresponding to pixels as described in image_uv() 132 | faces (torch.Tensor): quad faces connecting neighboring pixels 133 | indices (torch.Tensor, optional): indices of vertices in the original mesh 134 | """ 135 | if device is None and mask is not None: 136 | device = mask.device 137 | if mask is not None: 138 | assert mask.shape[0] == height and mask.shape[1] == width 139 | assert mask.dtype == torch.bool 140 | uv = image_uv(height, width, device=device, dtype=dtype).reshape((-1, 2)) 141 | row_faces = torch.stack([ 142 | torch.arange(0, width - 1, dtype=torch.int32, device=device), 143 | torch.arange(width, 2 * width - 1, dtype=torch.int32, device=device), 144 | torch.arange(1 + width, 2 * width, dtype=torch.int32, device=device), 145 | torch.arange(1, width, dtype=torch.int32, device=device) 146 | ], dim=1) 147 | faces = (torch.arange(0, (height - 1) * width, width, device=device, dtype=torch.int32)[:, None, None] + row_faces[None, :, :]).reshape((-1, 4)) 148 | if mask is not None: 149 | quad_mask = (mask[:-1, :-1] & mask[1:, :-1] & mask[1:, 1:] & mask[:-1, 1:]).ravel() 150 | faces = faces[quad_mask] 151 | faces, uv, indices = mesh.remove_unreferenced_vertices(faces, uv, return_indices=True) 152 | return uv, faces, indices 153 | return uv, faces 154 | 155 | 156 | def depth_edge(depth: torch.Tensor, atol: float = None, rtol: float = None, kernel_size: int = 3, mask: torch.Tensor = None) -> torch.BoolTensor: 157 | """ 158 | Compute the edge mask of a depth map. The edge is defined as the pixels whose neighbors have a large difference in depth. 159 | 160 | Args: 161 | depth (torch.Tensor): shape (..., height, width), linear depth map 162 | atol (float): absolute tolerance 163 | rtol (float): relative tolerance 164 | 165 | Returns: 166 | edge (torch.Tensor): shape (..., height, width) of dtype torch.bool 167 | """ 168 | shape = depth.shape 169 | depth = depth.reshape(-1, 1, *shape[-2:]) 170 | if mask is not None: 171 | mask = mask.reshape(-1, 1, *shape[-2:]) 172 | 173 | if mask is None: 174 | diff = (F.max_pool2d(depth, kernel_size, stride=1, padding=kernel_size // 2) + F.max_pool2d(-depth, kernel_size, stride=1, padding=kernel_size // 2)) 175 | else: 176 | diff = (F.max_pool2d(torch.where(mask, depth, -torch.inf), kernel_size, stride=1, padding=kernel_size // 2) + F.max_pool2d(torch.where(mask, -depth, -torch.inf), kernel_size, stride=1, padding=kernel_size // 2)) 177 | 178 | edge = torch.zeros_like(depth, dtype=torch.bool) 179 | if atol is not None: 180 | edge |= diff > atol 181 | if rtol is not None: 182 | edge |= (diff / depth).nan_to_num_() > rtol 183 | edge = edge.reshape(*shape) 184 | return edge 185 | 186 | 187 | def depth_aliasing(depth: torch.Tensor, atol: float = None, rtol: float = None, kernel_size: int = 3, mask: torch.Tensor = None) -> torch.BoolTensor: 188 | """ 189 | Compute the map that indicates the aliasing of a depth map. The aliasing is defined as the pixels which neither close to the maximum nor the minimum of its neighbors. 190 | Args: 191 | depth (torch.Tensor): shape (..., height, width), linear depth map 192 | atol (float): absolute tolerance 193 | rtol (float): relative tolerance 194 | 195 | Returns: 196 | edge (torch.Tensor): shape (..., height, width) of dtype torch.bool 197 | """ 198 | shape = depth.shape 199 | depth = depth.reshape(-1, 1, *shape[-2:]) 200 | if mask is not None: 201 | mask = mask.reshape(-1, 1, *shape[-2:]) 202 | 203 | if mask is None: 204 | diff_max = F.max_pool2d(depth, kernel_size, stride=1, padding=kernel_size // 2) - depth 205 | diff_min = F.max_pool2d(-depth, kernel_size, stride=1, padding=kernel_size // 2) + depth 206 | else: 207 | diff_max = F.max_pool2d(torch.where(mask, depth, -torch.inf), kernel_size, stride=1, padding=kernel_size // 2) - depth 208 | diff_min = F.max_pool2d(torch.where(mask, -depth, -torch.inf), kernel_size, stride=1, padding=kernel_size // 2) + depth 209 | diff = torch.minimum(diff_max, diff_min) 210 | 211 | edge = torch.zeros_like(depth, dtype=torch.bool) 212 | if atol is not None: 213 | edge |= diff > atol 214 | if rtol is not None: 215 | edge |= (diff / depth).nan_to_num_() > rtol 216 | edge = edge.reshape(*shape) 217 | return edge 218 | 219 | 220 | def image_mesh_from_depth( 221 | depth: torch.Tensor, 222 | extrinsics: torch.Tensor = None, 223 | intrinsics: torch.Tensor = None 224 | ) -> Tuple[torch.Tensor, torch.Tensor]: 225 | height, width = depth.shape 226 | uv, faces = image_mesh(height, width) 227 | faces = faces.reshape(-1, 4) 228 | depth = depth.reshape(-1) 229 | pts = transforms.unproject_cv(image_uv, depth, extrinsics, intrinsics) 230 | faces = mesh.triangulate(faces, vertices=pts) 231 | return pts, faces 232 | 233 | 234 | @batched(3, 2, 2) 235 | def points_to_normals(point: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor: 236 | """ 237 | Calculate normal map from point map. Value range is [-1, 1]. Normal direction in OpenGL identity camera's coordinate system. 238 | 239 | Args: 240 | point (torch.Tensor): shape (..., height, width, 3), point map 241 | Returns: 242 | normal (torch.Tensor): shape (..., height, width, 3), normal map. 243 | """ 244 | has_mask = mask is not None 245 | 246 | if mask is None: 247 | mask = torch.ones_like(point[..., 0], dtype=torch.bool) 248 | mask = F.pad(mask, (1, 1, 1, 1), mode='constant', value=0) 249 | 250 | pts = F.pad(point.permute(0, 3, 1, 2), (1, 1, 1, 1), mode='constant', value=1).permute(0, 2, 3, 1) 251 | up = pts[:, :-2, 1:-1, :] - pts[:, 1:-1, 1:-1, :] 252 | left = pts[:, 1:-1, :-2, :] - pts[:, 1:-1, 1:-1, :] 253 | down = pts[:, 2:, 1:-1, :] - pts[:, 1:-1, 1:-1, :] 254 | right = pts[:, 1:-1, 2:, :] - pts[:, 1:-1, 1:-1, :] 255 | normal = torch.stack([ 256 | torch.cross(up, left, dim=-1), 257 | torch.cross(left, down, dim=-1), 258 | torch.cross(down, right, dim=-1), 259 | torch.cross(right, up, dim=-1), 260 | ]) 261 | normal = F.normalize(normal, dim=-1) 262 | valid = torch.stack([ 263 | mask[:, :-2, 1:-1] & mask[:, 1:-1, :-2], 264 | mask[:, 1:-1, :-2] & mask[:, 2:, 1:-1], 265 | mask[:, 2:, 1:-1] & mask[:, 1:-1, 2:], 266 | mask[:, 1:-1, 2:] & mask[:, :-2, 1:-1], 267 | ]) & mask[None, :, 1:-1, 1:-1] 268 | normal = (normal * valid[..., None]).sum(dim=0) 269 | normal = F.normalize(normal, dim=-1) 270 | 271 | if has_mask: 272 | return normal, valid.any(dim=0) 273 | else: 274 | return normal 275 | 276 | 277 | def depth_to_normals(depth: torch.Tensor, intrinsics: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor: 278 | """ 279 | Calculate normal map from depth map. Value range is [-1, 1]. Normal direction in OpenGL identity camera's coordinate system. 280 | 281 | Args: 282 | depth (torch.Tensor): shape (..., height, width), linear depth map 283 | intrinsics (torch.Tensor): shape (..., 3, 3), intrinsics matrix 284 | Returns: 285 | normal (torch.Tensor): shape (..., 3, height, width), normal map. 286 | """ 287 | pts = depth_to_points(depth, intrinsics) 288 | return points_to_normals(pts, mask) 289 | 290 | 291 | def depth_to_points(depth: torch.Tensor, intrinsics: torch.Tensor, extrinsics: torch.Tensor = None): 292 | height, width = depth.shape[-2:] 293 | uv = image_uv(width=width, height=height, dtype=depth.dtype, device=depth.device) 294 | pts = transforms.unproject_cv(uv, depth, intrinsics=intrinsics[..., None, :, :], extrinsics=extrinsics[..., None, :, :] if extrinsics is not None else None) 295 | return pts 296 | 297 | 298 | def masked_min(input: torch.Tensor, mask: torch.BoolTensor, dim: int = None, keepdim: bool = False) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: 299 | """Similar to torch.min, but with mask 300 | """ 301 | if dim is None: 302 | return torch.where(mask, input, torch.tensor(torch.inf, dtype=input.dtype, device=input.device)).min() 303 | else: 304 | return torch.where(mask, input, torch.tensor(torch.inf, dtype=input.dtype, device=input.device)).min(dim=dim, keepdim=keepdim) 305 | 306 | 307 | def masked_max(input: torch.Tensor, mask: torch.BoolTensor, dim: int = None, keepdim: bool = False) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: 308 | """Similar to torch.max, but with mask 309 | """ 310 | if dim is None: 311 | return torch.where(mask, input, torch.tensor(-torch.inf, dtype=input.dtype, device=input.device)).max() 312 | else: 313 | return torch.where(mask, input, torch.tensor(-torch.inf, dtype=input.dtype, device=input.device)).max(dim=dim, keepdim=keepdim) 314 | 315 | 316 | def bounding_rect(mask: torch.BoolTensor): 317 | """get bounding rectangle of a mask 318 | 319 | Args: 320 | mask (torch.Tensor): shape (..., height, width), mask 321 | 322 | Returns: 323 | rect (torch.Tensor): shape (..., 4), bounding rectangle (left, top, right, bottom) 324 | """ 325 | height, width = mask.shape[-2:] 326 | mask = mask.flatten(-2).unsqueeze(-1) 327 | uv = image_uv(height, width).to(mask.device).reshape(-1, 2) 328 | left_top = masked_min(uv, mask, dim=-2)[0] 329 | right_bottom = masked_max(uv, mask, dim=-2)[0] 330 | return torch.cat([left_top, right_bottom], dim=-1) 331 | 332 | 333 | def chessboard(width: int, height: int, grid_size: int, color_a: torch.Tensor, color_b: torch.Tensor) -> torch.Tensor: 334 | """get a chessboard image 335 | 336 | Args: 337 | width (int): image width 338 | height (int): image height 339 | grid_size (int): size of chessboard grid 340 | color_a (torch.Tensor): shape (chanenls,), color of the grid at the top-left corner 341 | color_b (torch.Tensor): shape (chanenls,), color in complementary grids 342 | 343 | Returns: 344 | image (torch.Tensor): shape (height, width, channels), chessboard image 345 | """ 346 | x = torch.div(torch.arange(width), grid_size, rounding_mode='floor') 347 | y = torch.div(torch.arange(height), grid_size, rounding_mode='floor') 348 | mask = ((x[None, :] + y[:, None]) % 2).to(color_a) 349 | image = (1 - mask[..., None]) * color_a + mask[..., None] * color_b 350 | return image --------------------------------------------------------------------------------