gaussian_renderer 0.2.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- gaussian_renderer-0.2.0/LICENSE +21 -0
- gaussian_renderer-0.2.0/PKG-INFO +127 -0
- gaussian_renderer-0.2.0/README.md +96 -0
- gaussian_renderer-0.2.0/pyproject.toml +76 -0
- gaussian_renderer-0.2.0/setup.cfg +4 -0
- gaussian_renderer-0.2.0/src/gaussian_renderer/__init__.py +35 -0
- gaussian_renderer-0.2.0/src/gaussian_renderer/batch_splat.py +221 -0
- gaussian_renderer-0.2.0/src/gaussian_renderer/core/__init__.py +28 -0
- gaussian_renderer-0.2.0/src/gaussian_renderer/core/batch_rasterization.py +350 -0
- gaussian_renderer-0.2.0/src/gaussian_renderer/core/gaussiandata.py +74 -0
- gaussian_renderer-0.2.0/src/gaussian_renderer/core/gs_renderer.py +207 -0
- gaussian_renderer-0.2.0/src/gaussian_renderer/core/super_splat_loader.py +678 -0
- gaussian_renderer-0.2.0/src/gaussian_renderer/core/util_gau.py +299 -0
- gaussian_renderer-0.2.0/src/gaussian_renderer/gs_renderer_motrixsim.py +134 -0
- gaussian_renderer-0.2.0/src/gaussian_renderer/gs_renderer_mujoco.py +162 -0
- gaussian_renderer-0.2.0/src/gaussian_renderer/simple_viewer.py +287 -0
- gaussian_renderer-0.2.0/src/gaussian_renderer/supersplat_compress.py +395 -0
- gaussian_renderer-0.2.0/src/gaussian_renderer/transform_gs_model.py +253 -0
- gaussian_renderer-0.2.0/src/gaussian_renderer.egg-info/PKG-INFO +127 -0
- gaussian_renderer-0.2.0/src/gaussian_renderer.egg-info/SOURCES.txt +27 -0
- gaussian_renderer-0.2.0/src/gaussian_renderer.egg-info/dependency_links.txt +1 -0
- gaussian_renderer-0.2.0/src/gaussian_renderer.egg-info/entry_points.txt +4 -0
- gaussian_renderer-0.2.0/src/gaussian_renderer.egg-info/requires.txt +27 -0
- gaussian_renderer-0.2.0/src/gaussian_renderer.egg-info/top_level.txt +1 -0
- gaussian_renderer-0.2.0/tests/test_gaussiandata.py +35 -0
- gaussian_renderer-0.2.0/tests/test_motrixsim_renderer.py +133 -0
- gaussian_renderer-0.2.0/tests/test_mujoco_renderer.py +113 -0
- gaussian_renderer-0.2.0/tests/test_ply_io.py +45 -0
- gaussian_renderer-0.2.0/tests/test_utils.py +36 -0
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2025 Yufei Jia
|
|
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.
|
|
@@ -0,0 +1,127 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: gaussian_renderer
|
|
3
|
+
Version: 0.2.0
|
|
4
|
+
Summary: A Gaussian Splatting Renderer and Tools
|
|
5
|
+
Requires-Python: >=3.10
|
|
6
|
+
Description-Content-Type: text/markdown
|
|
7
|
+
License-File: LICENSE
|
|
8
|
+
Requires-Dist: numpy
|
|
9
|
+
Requires-Dist: torch
|
|
10
|
+
Requires-Dist: scipy
|
|
11
|
+
Requires-Dist: plyfile
|
|
12
|
+
Requires-Dist: trimesh
|
|
13
|
+
Requires-Dist: gsplat
|
|
14
|
+
Provides-Extra: viewer
|
|
15
|
+
Requires-Dist: glfw; extra == "viewer"
|
|
16
|
+
Requires-Dist: PyOpenGL; extra == "viewer"
|
|
17
|
+
Provides-Extra: motrix
|
|
18
|
+
Requires-Dist: motrixsim>=0.5.0; extra == "motrix"
|
|
19
|
+
Provides-Extra: mujoco
|
|
20
|
+
Requires-Dist: mujoco; extra == "mujoco"
|
|
21
|
+
Provides-Extra: shs
|
|
22
|
+
Requires-Dist: einops; extra == "shs"
|
|
23
|
+
Requires-Dist: e3nn; extra == "shs"
|
|
24
|
+
Provides-Extra: dev
|
|
25
|
+
Requires-Dist: ruff; extra == "dev"
|
|
26
|
+
Requires-Dist: mypy; extra == "dev"
|
|
27
|
+
Requires-Dist: pytest; extra == "dev"
|
|
28
|
+
Requires-Dist: einops; extra == "dev"
|
|
29
|
+
Requires-Dist: e3nn; extra == "dev"
|
|
30
|
+
Dynamic: license-file
|
|
31
|
+
|
|
32
|
+
# Gaussian Renderer
|
|
33
|
+
|
|
34
|
+
A Gaussian Splatting Renderer and Tools package.
|
|
35
|
+
|
|
36
|
+
[中文文档](README_zh.md)
|
|
37
|
+
|
|
38
|
+
This repository primarily provides rendering capabilities for Gaussian Splatting models. It is developed as a component of the **DISCOVERSE** project.
|
|
39
|
+
|
|
40
|
+
For detailed usage within the simulation environment, see: [https://github.com/TATP-233/DISCOVERSE](https://github.com/TATP-233/DISCOVERSE)
|
|
41
|
+
|
|
42
|
+
## Requirements
|
|
43
|
+
|
|
44
|
+
Python >= 3.10
|
|
45
|
+
|
|
46
|
+
## Installation
|
|
47
|
+
|
|
48
|
+
```bash
|
|
49
|
+
uv add gaussian-renderer
|
|
50
|
+
# or: pip install gaussian-renderer
|
|
51
|
+
```
|
|
52
|
+
|
|
53
|
+
From source:
|
|
54
|
+
```bash
|
|
55
|
+
git clone https://github.com/TATP-233/GaussainRenderer.git
|
|
56
|
+
cd GaussainRenderer
|
|
57
|
+
uv pip install .
|
|
58
|
+
# or: pip install .
|
|
59
|
+
```
|
|
60
|
+
|
|
61
|
+
### Optional extras
|
|
62
|
+
|
|
63
|
+
```bash
|
|
64
|
+
uv add "gaussian-renderer[viewer]" # OpenGL viewer (glfw, PyOpenGL)
|
|
65
|
+
uv add "gaussian-renderer[mujoco]" # MuJoCo integration
|
|
66
|
+
uv add "gaussian-renderer[motrix]" # MotrixSim integration
|
|
67
|
+
|
|
68
|
+
# Combine as needed
|
|
69
|
+
uv add "gaussian-renderer[viewer,mujoco]"
|
|
70
|
+
# or: pip install ".[viewer,mujoco]"
|
|
71
|
+
```
|
|
72
|
+
|
|
73
|
+
## Usage
|
|
74
|
+
|
|
75
|
+
### Command-line tools
|
|
76
|
+
|
|
77
|
+
**`gs-viewer`** — OpenGL viewer for `.ply` models
|
|
78
|
+
```bash
|
|
79
|
+
gs-viewer path/to/model.ply
|
|
80
|
+
```
|
|
81
|
+
Controls: Left mouse = rotate, Right/Middle = pan, Scroll = zoom, Up/Down = SH degree, Drag & drop = load file
|
|
82
|
+
|
|
83
|
+
**`gs-compress`** — Compress 3DGS PLY to SuperSplat format
|
|
84
|
+
```bash
|
|
85
|
+
gs-compress input.ply
|
|
86
|
+
gs-compress input.ply -o output.ply
|
|
87
|
+
gs-compress models/ # batch
|
|
88
|
+
```
|
|
89
|
+
|
|
90
|
+
**`gs-transform`** — Apply translation/rotation/scale to a model
|
|
91
|
+
```bash
|
|
92
|
+
gs-transform input.ply -o output.ply -t 0 1 0 -s 2.0
|
|
93
|
+
gs-transform input.ply -r 0 0 0 1 # rotation quaternion xyzw
|
|
94
|
+
# --compress: save as compressed PLY
|
|
95
|
+
```
|
|
96
|
+
|
|
97
|
+
### Python API
|
|
98
|
+
|
|
99
|
+
```bash
|
|
100
|
+
uv run python -m gaussian_renderer.simple_viewer path/to/model.ply
|
|
101
|
+
uv run python -m gaussian_renderer.supersplat_compress input.ply
|
|
102
|
+
uv run python -m gaussian_renderer.transform_gs_model input.ply
|
|
103
|
+
```
|
|
104
|
+
|
|
105
|
+
## Development
|
|
106
|
+
|
|
107
|
+
```bash
|
|
108
|
+
uv pip install ".[dev]"
|
|
109
|
+
# or: pip install ".[dev]"
|
|
110
|
+
make lint # ruff check
|
|
111
|
+
make format # ruff format
|
|
112
|
+
make typecheck # mypy
|
|
113
|
+
make test # pytest
|
|
114
|
+
make ci # all of the above
|
|
115
|
+
```
|
|
116
|
+
|
|
117
|
+
## Citation
|
|
118
|
+
|
|
119
|
+
```bibtex
|
|
120
|
+
@article{jia2025discoverse,
|
|
121
|
+
title={DISCOVERSE: Efficient Robot Simulation in Complex High-Fidelity Environments},
|
|
122
|
+
author={Yufei Jia and Guangyu Wang and Yuhang Dong and Junzhe Wu and Yupei Zeng and Haonan Lin and Zifan Wang and Haizhou Ge and Weibin Gu and Chuxuan Li and Ziming Wang and Yunjie Cheng and Wei Sui and Ruqi Huang and Guyue Zhou},
|
|
123
|
+
journal={arXiv preprint arXiv:2507.21981},
|
|
124
|
+
year={2025},
|
|
125
|
+
url={https://arxiv.org/abs/2507.21981}
|
|
126
|
+
}
|
|
127
|
+
```
|
|
@@ -0,0 +1,96 @@
|
|
|
1
|
+
# Gaussian Renderer
|
|
2
|
+
|
|
3
|
+
A Gaussian Splatting Renderer and Tools package.
|
|
4
|
+
|
|
5
|
+
[中文文档](README_zh.md)
|
|
6
|
+
|
|
7
|
+
This repository primarily provides rendering capabilities for Gaussian Splatting models. It is developed as a component of the **DISCOVERSE** project.
|
|
8
|
+
|
|
9
|
+
For detailed usage within the simulation environment, see: [https://github.com/TATP-233/DISCOVERSE](https://github.com/TATP-233/DISCOVERSE)
|
|
10
|
+
|
|
11
|
+
## Requirements
|
|
12
|
+
|
|
13
|
+
Python >= 3.10
|
|
14
|
+
|
|
15
|
+
## Installation
|
|
16
|
+
|
|
17
|
+
```bash
|
|
18
|
+
uv add gaussian-renderer
|
|
19
|
+
# or: pip install gaussian-renderer
|
|
20
|
+
```
|
|
21
|
+
|
|
22
|
+
From source:
|
|
23
|
+
```bash
|
|
24
|
+
git clone https://github.com/TATP-233/GaussainRenderer.git
|
|
25
|
+
cd GaussainRenderer
|
|
26
|
+
uv pip install .
|
|
27
|
+
# or: pip install .
|
|
28
|
+
```
|
|
29
|
+
|
|
30
|
+
### Optional extras
|
|
31
|
+
|
|
32
|
+
```bash
|
|
33
|
+
uv add "gaussian-renderer[viewer]" # OpenGL viewer (glfw, PyOpenGL)
|
|
34
|
+
uv add "gaussian-renderer[mujoco]" # MuJoCo integration
|
|
35
|
+
uv add "gaussian-renderer[motrix]" # MotrixSim integration
|
|
36
|
+
|
|
37
|
+
# Combine as needed
|
|
38
|
+
uv add "gaussian-renderer[viewer,mujoco]"
|
|
39
|
+
# or: pip install ".[viewer,mujoco]"
|
|
40
|
+
```
|
|
41
|
+
|
|
42
|
+
## Usage
|
|
43
|
+
|
|
44
|
+
### Command-line tools
|
|
45
|
+
|
|
46
|
+
**`gs-viewer`** — OpenGL viewer for `.ply` models
|
|
47
|
+
```bash
|
|
48
|
+
gs-viewer path/to/model.ply
|
|
49
|
+
```
|
|
50
|
+
Controls: Left mouse = rotate, Right/Middle = pan, Scroll = zoom, Up/Down = SH degree, Drag & drop = load file
|
|
51
|
+
|
|
52
|
+
**`gs-compress`** — Compress 3DGS PLY to SuperSplat format
|
|
53
|
+
```bash
|
|
54
|
+
gs-compress input.ply
|
|
55
|
+
gs-compress input.ply -o output.ply
|
|
56
|
+
gs-compress models/ # batch
|
|
57
|
+
```
|
|
58
|
+
|
|
59
|
+
**`gs-transform`** — Apply translation/rotation/scale to a model
|
|
60
|
+
```bash
|
|
61
|
+
gs-transform input.ply -o output.ply -t 0 1 0 -s 2.0
|
|
62
|
+
gs-transform input.ply -r 0 0 0 1 # rotation quaternion xyzw
|
|
63
|
+
# --compress: save as compressed PLY
|
|
64
|
+
```
|
|
65
|
+
|
|
66
|
+
### Python API
|
|
67
|
+
|
|
68
|
+
```bash
|
|
69
|
+
uv run python -m gaussian_renderer.simple_viewer path/to/model.ply
|
|
70
|
+
uv run python -m gaussian_renderer.supersplat_compress input.ply
|
|
71
|
+
uv run python -m gaussian_renderer.transform_gs_model input.ply
|
|
72
|
+
```
|
|
73
|
+
|
|
74
|
+
## Development
|
|
75
|
+
|
|
76
|
+
```bash
|
|
77
|
+
uv pip install ".[dev]"
|
|
78
|
+
# or: pip install ".[dev]"
|
|
79
|
+
make lint # ruff check
|
|
80
|
+
make format # ruff format
|
|
81
|
+
make typecheck # mypy
|
|
82
|
+
make test # pytest
|
|
83
|
+
make ci # all of the above
|
|
84
|
+
```
|
|
85
|
+
|
|
86
|
+
## Citation
|
|
87
|
+
|
|
88
|
+
```bibtex
|
|
89
|
+
@article{jia2025discoverse,
|
|
90
|
+
title={DISCOVERSE: Efficient Robot Simulation in Complex High-Fidelity Environments},
|
|
91
|
+
author={Yufei Jia and Guangyu Wang and Yuhang Dong and Junzhe Wu and Yupei Zeng and Haonan Lin and Zifan Wang and Haizhou Ge and Weibin Gu and Chuxuan Li and Ziming Wang and Yunjie Cheng and Wei Sui and Ruqi Huang and Guyue Zhou},
|
|
92
|
+
journal={arXiv preprint arXiv:2507.21981},
|
|
93
|
+
year={2025},
|
|
94
|
+
url={https://arxiv.org/abs/2507.21981}
|
|
95
|
+
}
|
|
96
|
+
```
|
|
@@ -0,0 +1,76 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["setuptools>=61.0"]
|
|
3
|
+
build-backend = "setuptools.build_meta"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "gaussian_renderer"
|
|
7
|
+
dynamic = ["version"]
|
|
8
|
+
description = "A Gaussian Splatting Renderer and Tools"
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
requires-python = ">=3.10"
|
|
11
|
+
dependencies = [
|
|
12
|
+
"numpy",
|
|
13
|
+
"torch",
|
|
14
|
+
"scipy",
|
|
15
|
+
"plyfile",
|
|
16
|
+
"trimesh",
|
|
17
|
+
"gsplat",
|
|
18
|
+
]
|
|
19
|
+
|
|
20
|
+
[project.optional-dependencies]
|
|
21
|
+
viewer = [
|
|
22
|
+
"glfw",
|
|
23
|
+
"PyOpenGL",
|
|
24
|
+
]
|
|
25
|
+
|
|
26
|
+
motrix = [
|
|
27
|
+
"motrixsim >= 0.5.0",
|
|
28
|
+
]
|
|
29
|
+
|
|
30
|
+
mujoco = [
|
|
31
|
+
"mujoco",
|
|
32
|
+
]
|
|
33
|
+
|
|
34
|
+
shs = [
|
|
35
|
+
"einops",
|
|
36
|
+
"e3nn",
|
|
37
|
+
]
|
|
38
|
+
|
|
39
|
+
dev = [
|
|
40
|
+
"ruff",
|
|
41
|
+
"mypy",
|
|
42
|
+
"pytest",
|
|
43
|
+
"einops",
|
|
44
|
+
"e3nn",
|
|
45
|
+
]
|
|
46
|
+
|
|
47
|
+
[project.scripts]
|
|
48
|
+
gs-viewer = "gaussian_renderer.simple_viewer:main"
|
|
49
|
+
gs-compress = "gaussian_renderer.supersplat_compress:main"
|
|
50
|
+
gs-transform = "gaussian_renderer.transform_gs_model:main"
|
|
51
|
+
|
|
52
|
+
[tool.setuptools.dynamic]
|
|
53
|
+
version = {attr = "gaussian_renderer.__version__"}
|
|
54
|
+
|
|
55
|
+
[tool.setuptools.packages.find]
|
|
56
|
+
where = ["src"]
|
|
57
|
+
include = ["gaussian_renderer*"]
|
|
58
|
+
|
|
59
|
+
[tool.ruff]
|
|
60
|
+
line-length = 120
|
|
61
|
+
target-version = "py310"
|
|
62
|
+
|
|
63
|
+
[tool.ruff.lint]
|
|
64
|
+
select = ["E", "F", "I"]
|
|
65
|
+
ignore = ["E501", "F401", "E402", "F841"]
|
|
66
|
+
|
|
67
|
+
[tool.mypy]
|
|
68
|
+
python_version = "3.10"
|
|
69
|
+
ignore_missing_imports = true
|
|
70
|
+
warn_unused_ignores = false
|
|
71
|
+
disable_error_code = ["assignment", "call-overload", "union-attr", "arg-type", "dict-item", "index", "return-value", "operator", "override", "type-arg", "attr-defined"]
|
|
72
|
+
|
|
73
|
+
[tool.pytest.ini_options]
|
|
74
|
+
testpaths = ["tests"]
|
|
75
|
+
python_files = ["test_*.py"]
|
|
76
|
+
pythonpath = ["src"]
|
|
@@ -0,0 +1,35 @@
|
|
|
1
|
+
# SPDX-License-Identifier: MIT
|
|
2
|
+
#
|
|
3
|
+
# MIT License
|
|
4
|
+
#
|
|
5
|
+
# Copyright (c) 2025 Yufei Jia
|
|
6
|
+
#
|
|
7
|
+
# Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
8
|
+
# of this software and associated documentation files (the "Software"), to deal
|
|
9
|
+
# in the Software without restriction, including without limitation the rights
|
|
10
|
+
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
11
|
+
# copies of the Software, and to permit persons to whom the Software is
|
|
12
|
+
# furnished to do so, subject to the following conditions:
|
|
13
|
+
#
|
|
14
|
+
# The above copyright notice and this permission notice shall be included in all
|
|
15
|
+
# copies or substantial portions of the Software.
|
|
16
|
+
#
|
|
17
|
+
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
18
|
+
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
19
|
+
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
20
|
+
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
21
|
+
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
22
|
+
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
23
|
+
# SOFTWARE.
|
|
24
|
+
|
|
25
|
+
from .batch_splat import BatchSplatConfig, BatchSplatRenderer, MjxBatchSplatRenderer, MtxBatchSplatRenderer
|
|
26
|
+
from .core.batch_rasterization import batch_env_render, batch_render, batch_update_gaussians
|
|
27
|
+
from .core.gaussiandata import GaussianBatchData, GaussianData
|
|
28
|
+
from .core.gs_renderer import GSRenderer
|
|
29
|
+
from .core.super_splat_loader import is_super_splat_format, load_super_splat_ply, save_super_splat_ply
|
|
30
|
+
from .core.util_gau import load_ply, save_ply, transform_shs
|
|
31
|
+
from .gs_renderer_motrixsim import GSRendererMotrixSim
|
|
32
|
+
from .gs_renderer_mujoco import GSRendererMuJoCo
|
|
33
|
+
|
|
34
|
+
__version__ = "0.2.0"
|
|
35
|
+
__author__ = "Yufei Jia"
|
|
@@ -0,0 +1,221 @@
|
|
|
1
|
+
# SPDX-License-Identifier: MIT
|
|
2
|
+
#
|
|
3
|
+
# MIT License
|
|
4
|
+
#
|
|
5
|
+
# Copyright (c) 2025 Yufei Jia
|
|
6
|
+
#
|
|
7
|
+
# Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
8
|
+
# of this software and associated documentation files (the "Software"), to deal
|
|
9
|
+
# in the Software without restriction, including without limitation the rights
|
|
10
|
+
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
11
|
+
# copies of the Software, and to permit persons to whom the Software is
|
|
12
|
+
# furnished to do so, subject to the following conditions:
|
|
13
|
+
#
|
|
14
|
+
# The above copyright notice and this permission notice shall be included in all
|
|
15
|
+
# copies or substantial portions of the Software.
|
|
16
|
+
#
|
|
17
|
+
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
18
|
+
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
19
|
+
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
20
|
+
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
21
|
+
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
22
|
+
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
23
|
+
# SOFTWARE.
|
|
24
|
+
|
|
25
|
+
"""
|
|
26
|
+
Minimal batch Gaussian splat renderer that works with mjx state tensors.
|
|
27
|
+
|
|
28
|
+
Usage pattern (3 steps):
|
|
29
|
+
1) init: cfg = BatchSplatConfig(...); renderer = BatchSplatRenderer(cfg)
|
|
30
|
+
2) loop: gsb = renderer.batch_update_gaussians(body_pos, body_quat)
|
|
31
|
+
3) loop: rgb, depth = renderer.batch_env_render(gsb, cam_pos, cam_xmat, height, width, fovy)
|
|
32
|
+
|
|
33
|
+
Where body_pos/body_quat/cam_pos/cam_xmat can come from mjx (MuJoCo JAX) state via
|
|
34
|
+
`torch.utils.dlpack.from_dlpack(state.data.xxx)`.
|
|
35
|
+
"""
|
|
36
|
+
|
|
37
|
+
from dataclasses import dataclass
|
|
38
|
+
from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Tuple
|
|
39
|
+
|
|
40
|
+
import numpy as np
|
|
41
|
+
import torch
|
|
42
|
+
from torch import Tensor
|
|
43
|
+
|
|
44
|
+
from .core.batch_rasterization import (
|
|
45
|
+
batch_env_render as _batch_env_render,
|
|
46
|
+
)
|
|
47
|
+
from .core.batch_rasterization import (
|
|
48
|
+
batch_update_gaussians as _batch_update_gaussians,
|
|
49
|
+
)
|
|
50
|
+
from .core.gaussiandata import GaussianBatchData, GaussianData
|
|
51
|
+
from .core.util_gau import load_ply
|
|
52
|
+
|
|
53
|
+
|
|
54
|
+
@dataclass
|
|
55
|
+
class BatchSplatConfig:
|
|
56
|
+
# Mapping from MuJoCo body name -> PLY path (local coordinate frame).
|
|
57
|
+
body_gaussians: Dict[str, str]
|
|
58
|
+
# Optional static/background PLY that is not attached to any body.
|
|
59
|
+
background_ply: Optional[str] = None
|
|
60
|
+
# Device to place tensors on; if None, will pick CUDA if available else CPU.
|
|
61
|
+
device: Optional[torch.device] = None
|
|
62
|
+
# Minibatch size for rendering
|
|
63
|
+
minibatch: Optional[int] = None
|
|
64
|
+
|
|
65
|
+
|
|
66
|
+
class BatchSplatRenderer:
|
|
67
|
+
def __init__(self, cfg: BatchSplatConfig, body_name_to_id: Optional[Dict[str, int]]) -> None:
|
|
68
|
+
assert torch.cuda.is_available()
|
|
69
|
+
device = cfg.device or torch.device("cuda")
|
|
70
|
+
|
|
71
|
+
self.device = device
|
|
72
|
+
self.minibatch = cfg.minibatch
|
|
73
|
+
|
|
74
|
+
xyz_list: List[Tensor] = []
|
|
75
|
+
rot_list: List[Tensor] = []
|
|
76
|
+
scale_list: List[Tensor] = []
|
|
77
|
+
opacity_list: List[Tensor] = []
|
|
78
|
+
sh_list: List[Tensor] = []
|
|
79
|
+
max_sh_dim = 0
|
|
80
|
+
|
|
81
|
+
self.gs_idx_start: List[int] = []
|
|
82
|
+
self.gs_idx_end: List[int] = []
|
|
83
|
+
self.gs_body_ids: List[int] = []
|
|
84
|
+
|
|
85
|
+
# Load per-body gaussians
|
|
86
|
+
for body_name, ply_path in cfg.body_gaussians.items():
|
|
87
|
+
g = load_ply(ply_path)
|
|
88
|
+
max_sh_dim = max(max_sh_dim, g.sh.shape[1])
|
|
89
|
+
start = len(torch.cat(xyz_list)) if xyz_list else 0
|
|
90
|
+
end = start + len(g.xyz)
|
|
91
|
+
self.gs_idx_start.append(start)
|
|
92
|
+
self.gs_idx_end.append(end)
|
|
93
|
+
if body_name not in body_name_to_id:
|
|
94
|
+
raise ValueError(f"Body '{body_name}' not found in mj_model; available: {list(body_name_to_id.keys())}")
|
|
95
|
+
self.gs_body_ids.append(body_name_to_id[body_name])
|
|
96
|
+
xyz_list.append(torch.tensor(g.xyz, device=device, dtype=torch.float32))
|
|
97
|
+
rot_list.append(torch.tensor(g.rot, device=device, dtype=torch.float32))
|
|
98
|
+
scale_list.append(torch.tensor(g.scale, device=device, dtype=torch.float32))
|
|
99
|
+
opacity_list.append(torch.tensor(g.opacity, device=device, dtype=torch.float32))
|
|
100
|
+
sh_list.append(torch.tensor(g.sh, device=device, dtype=torch.float32))
|
|
101
|
+
|
|
102
|
+
# Optional background/static gaussian
|
|
103
|
+
if cfg.background_ply:
|
|
104
|
+
g = load_ply(cfg.background_ply)
|
|
105
|
+
max_sh_dim = max(max_sh_dim, g.sh.shape[1])
|
|
106
|
+
xyz_list.append(torch.tensor(g.xyz, device=device, dtype=torch.float32))
|
|
107
|
+
rot_list.append(torch.tensor(g.rot, device=device, dtype=torch.float32))
|
|
108
|
+
scale_list.append(torch.tensor(g.scale, device=device, dtype=torch.float32))
|
|
109
|
+
opacity_list.append(torch.tensor(g.opacity, device=device, dtype=torch.float32))
|
|
110
|
+
sh_list.append(torch.tensor(g.sh, device=device, dtype=torch.float32))
|
|
111
|
+
|
|
112
|
+
if sh_list:
|
|
113
|
+
for i, sh in enumerate(sh_list):
|
|
114
|
+
cur_dim = sh.shape[1]
|
|
115
|
+
if cur_dim < max_sh_dim:
|
|
116
|
+
pad = torch.zeros(sh.shape[0], max_sh_dim - cur_dim, device=sh.device, dtype=sh.dtype)
|
|
117
|
+
sh_list[i] = torch.cat([sh, pad], dim=1)
|
|
118
|
+
|
|
119
|
+
# Concatenate template
|
|
120
|
+
xyz_all = torch.cat(xyz_list, dim=0)
|
|
121
|
+
rot_all = torch.cat(rot_list, dim=0)
|
|
122
|
+
scale_all = torch.cat(scale_list, dim=0)
|
|
123
|
+
opacity_all = torch.cat(opacity_list, dim=0)
|
|
124
|
+
sh_all = torch.cat(sh_list, dim=0).reshape(xyz_all.shape[0], -1, 3).contiguous()
|
|
125
|
+
|
|
126
|
+
self.template = GaussianData(
|
|
127
|
+
xyz=xyz_all,
|
|
128
|
+
rot=rot_all,
|
|
129
|
+
scale=scale_all,
|
|
130
|
+
opacity=opacity_all,
|
|
131
|
+
sh=sh_all,
|
|
132
|
+
)
|
|
133
|
+
|
|
134
|
+
self.gs_idx_start = torch.tensor(self.gs_idx_start, dtype=torch.long, device=device)
|
|
135
|
+
self.gs_idx_end = torch.tensor(self.gs_idx_end, dtype=torch.long, device=device)
|
|
136
|
+
self.gs_body_ids = torch.tensor(self.gs_body_ids, dtype=torch.long, device=device)
|
|
137
|
+
|
|
138
|
+
# Precompute point-to-body mapping for vectorized update
|
|
139
|
+
num_points = len(self.template)
|
|
140
|
+
self.dynamic_mask = torch.zeros(num_points, dtype=torch.bool, device=device)
|
|
141
|
+
self.point_to_body_idx = torch.zeros(num_points, dtype=torch.long, device=device)
|
|
142
|
+
|
|
143
|
+
for i in range(len(self.gs_idx_start)):
|
|
144
|
+
start = self.gs_idx_start[i]
|
|
145
|
+
end = self.gs_idx_end[i]
|
|
146
|
+
body_id = self.gs_body_ids[i]
|
|
147
|
+
self.dynamic_mask[start:end] = True
|
|
148
|
+
self.point_to_body_idx[start:end] = body_id
|
|
149
|
+
|
|
150
|
+
def batch_update_gaussians(self, body_pos: Tensor, body_quat: Tensor, scalar_first=True) -> GaussianBatchData:
|
|
151
|
+
"""Update gaussians using body poses.
|
|
152
|
+
|
|
153
|
+
Args:
|
|
154
|
+
body_pos: (Nenv, Nbody, 3) torch tensor
|
|
155
|
+
body_quat: (Nenv, Nbody, 4) torch tensor (wxyz)
|
|
156
|
+
Returns:
|
|
157
|
+
GaussianBatchData with per-env gaussians.
|
|
158
|
+
"""
|
|
159
|
+
# Ensure device
|
|
160
|
+
if not isinstance(body_pos, torch.Tensor):
|
|
161
|
+
body_pos = torch.tensor(body_pos, device=self.device, dtype=torch.float32)
|
|
162
|
+
if not isinstance(body_quat, torch.Tensor):
|
|
163
|
+
body_quat = torch.tensor(body_quat, device=self.device, dtype=torch.float32)
|
|
164
|
+
return _batch_update_gaussians(
|
|
165
|
+
self.template,
|
|
166
|
+
body_pos,
|
|
167
|
+
body_quat,
|
|
168
|
+
point_to_body_idx=self.point_to_body_idx,
|
|
169
|
+
dynamic_mask=self.dynamic_mask,
|
|
170
|
+
scalar_first=scalar_first,
|
|
171
|
+
)
|
|
172
|
+
|
|
173
|
+
def batch_env_render(
|
|
174
|
+
self,
|
|
175
|
+
gsb: GaussianBatchData,
|
|
176
|
+
cam_pos: Tensor,
|
|
177
|
+
cam_xmat: Tensor,
|
|
178
|
+
height: int,
|
|
179
|
+
width: int,
|
|
180
|
+
fovy: np.ndarray,
|
|
181
|
+
bg_imgs: Optional[Tensor] = None,
|
|
182
|
+
):
|
|
183
|
+
"""Render RGBD for batch envs and cameras."""
|
|
184
|
+
if not isinstance(cam_pos, torch.Tensor):
|
|
185
|
+
cam_pos = torch.tensor(cam_pos, device=self.device, dtype=torch.float32)
|
|
186
|
+
if not isinstance(cam_xmat, torch.Tensor):
|
|
187
|
+
cam_xmat = torch.tensor(cam_xmat, device=self.device, dtype=torch.float32)
|
|
188
|
+
return _batch_env_render(gsb, cam_pos, cam_xmat, height, width, fovy, bg_imgs=bg_imgs, minibatch=self.minibatch)
|
|
189
|
+
|
|
190
|
+
|
|
191
|
+
if TYPE_CHECKING:
|
|
192
|
+
import mujoco
|
|
193
|
+
|
|
194
|
+
|
|
195
|
+
class MjxBatchSplatRenderer(BatchSplatRenderer):
|
|
196
|
+
def __init__(self, cfg: BatchSplatConfig, mj_model: "mujoco.MjModel") -> None:
|
|
197
|
+
body_name_to_id = {}
|
|
198
|
+
for i in range(mj_model.nbody):
|
|
199
|
+
body_name = mj_model.body(i).name
|
|
200
|
+
if body_name:
|
|
201
|
+
body_name_to_id[body_name] = i
|
|
202
|
+
super().__init__(cfg, body_name_to_id)
|
|
203
|
+
|
|
204
|
+
|
|
205
|
+
if TYPE_CHECKING:
|
|
206
|
+
import motrixsim
|
|
207
|
+
|
|
208
|
+
|
|
209
|
+
class MtxBatchSplatRenderer(BatchSplatRenderer):
|
|
210
|
+
def __init__(self, cfg, mx_model: "motrixsim.MotrixSimModel") -> None:
|
|
211
|
+
body_name_to_id = {}
|
|
212
|
+
for i, link_name in enumerate(mx_model.link_names):
|
|
213
|
+
if link_name:
|
|
214
|
+
body_name_to_id[link_name] = i
|
|
215
|
+
super().__init__(cfg, body_name_to_id)
|
|
216
|
+
|
|
217
|
+
def batch_update_gaussians(self, body_pos: Tensor, body_quat: Tensor) -> GaussianBatchData:
|
|
218
|
+
return super().batch_update_gaussians(body_pos, body_quat, scalar_first=False)
|
|
219
|
+
|
|
220
|
+
|
|
221
|
+
__all__ = ["BatchSplatConfig", "BatchSplatRenderer", "MjxBatchSplatRenderer", "MtxBatchSplatRenderer"]
|
|
@@ -0,0 +1,28 @@
|
|
|
1
|
+
# SPDX-License-Identifier: MIT
|
|
2
|
+
#
|
|
3
|
+
# MIT License
|
|
4
|
+
#
|
|
5
|
+
# Copyright (c) 2025 Yufei Jia
|
|
6
|
+
#
|
|
7
|
+
# Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
8
|
+
# of this software and associated documentation files (the "Software"), to deal
|
|
9
|
+
# in the Software without restriction, including without limitation the rights
|
|
10
|
+
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
11
|
+
# copies of the Software, and to permit persons to whom the Software is
|
|
12
|
+
# furnished to do so, subject to the following conditions:
|
|
13
|
+
#
|
|
14
|
+
# The above copyright notice and this permission notice shall be included in all
|
|
15
|
+
# copies or substantial portions of the Software.
|
|
16
|
+
#
|
|
17
|
+
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
18
|
+
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
19
|
+
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
20
|
+
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
21
|
+
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
22
|
+
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
23
|
+
# SOFTWARE.
|
|
24
|
+
|
|
25
|
+
from .batch_rasterization import batch_env_render, batch_render, batch_update_gaussians
|
|
26
|
+
from .gaussiandata import GaussianBatchData, GaussianData
|
|
27
|
+
from .gs_renderer import GSRenderer
|
|
28
|
+
from .util_gau import load_ply, save_ply, transform_shs
|