flashdrr 0.5__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.
- flashdrr-0.5/LICENSE +21 -0
- flashdrr-0.5/MANIFEST.in +10 -0
- flashdrr-0.5/PKG-INFO +223 -0
- flashdrr-0.5/README.md +168 -0
- flashdrr-0.5/flashdrr/__init__.py +1 -0
- flashdrr-0.5/flashdrr/rendering/__init__.py +36 -0
- flashdrr-0.5/flashdrr/rendering/raycast.py +902 -0
- flashdrr-0.5/flashdrr/rendering/triton_raycast.py +678 -0
- flashdrr-0.5/flashdrr/utils/__init__.py +2 -0
- flashdrr-0.5/flashdrr/utils/volume_utils.py +46 -0
- flashdrr-0.5/flashdrr.egg-info/PKG-INFO +223 -0
- flashdrr-0.5/flashdrr.egg-info/SOURCES.txt +21 -0
- flashdrr-0.5/flashdrr.egg-info/dependency_links.txt +1 -0
- flashdrr-0.5/flashdrr.egg-info/not-zip-safe +1 -0
- flashdrr-0.5/flashdrr.egg-info/requires.txt +39 -0
- flashdrr-0.5/flashdrr.egg-info/top_level.txt +1 -0
- flashdrr-0.5/pyproject.toml +207 -0
- flashdrr-0.5/setup.cfg +4 -0
- flashdrr-0.5/tests/test_camera_helpers.py +78 -0
- flashdrr-0.5/tests/test_piecewise_linear_channelwise.py +82 -0
- flashdrr-0.5/tests/test_regression_ctchest.py +218 -0
- flashdrr-0.5/tests/test_triton_vs_python.py +342 -0
- flashdrr-0.5/tests/test_volume_raycaster.py +311 -0
flashdrr-0.5/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Patrick Carnahan
|
|
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.
|
flashdrr-0.5/MANIFEST.in
ADDED
|
@@ -0,0 +1,10 @@
|
|
|
1
|
+
# Files to include in the source distribution (sdist) beyond what
|
|
2
|
+
# setuptools auto-detects from pyproject.toml.
|
|
3
|
+
|
|
4
|
+
# Standard project metadata files
|
|
5
|
+
include LICENSE
|
|
6
|
+
include README.md
|
|
7
|
+
include pyproject.toml
|
|
8
|
+
|
|
9
|
+
# Ship the small bundled test data
|
|
10
|
+
recursive-include flashdrr/data *.pt
|
flashdrr-0.5/PKG-INFO
ADDED
|
@@ -0,0 +1,223 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: flashdrr
|
|
3
|
+
Version: 0.5
|
|
4
|
+
Summary: Fast differentiable DRR raycasting and volume rendering for PyTorch
|
|
5
|
+
Author: Patrick Carnahan, Arteen Mirzai
|
|
6
|
+
License: MIT
|
|
7
|
+
Project-URL: Homepage, https://github.com/pcarnah/flashdrr
|
|
8
|
+
Project-URL: Repository, https://github.com/pcarnah/flashdrr.git
|
|
9
|
+
Project-URL: Issues, https://github.com/pcarnah/flashdrr/issues
|
|
10
|
+
Keywords: pytorch,volume,raycasting,rendering,medical-imaging,deep-learning,radiograph,drr
|
|
11
|
+
Classifier: Development Status :: 4 - Beta
|
|
12
|
+
Classifier: Intended Audience :: Science/Research
|
|
13
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
14
|
+
Classifier: Operating System :: OS Independent
|
|
15
|
+
Classifier: Programming Language :: Python :: 3
|
|
16
|
+
Classifier: Programming Language :: Python :: 3 :: Only
|
|
17
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
18
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
19
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
20
|
+
Classifier: Topic :: Scientific/Engineering
|
|
21
|
+
Classifier: Topic :: Scientific/Engineering :: Medical Science Apps.
|
|
22
|
+
Classifier: Topic :: Scientific/Engineering :: Visualization
|
|
23
|
+
Requires-Python: >=3.10
|
|
24
|
+
Description-Content-Type: text/markdown
|
|
25
|
+
License-File: LICENSE
|
|
26
|
+
Requires-Dist: torch>=2.5
|
|
27
|
+
Requires-Dist: torchvision>=0.20
|
|
28
|
+
Requires-Dist: numpy
|
|
29
|
+
Requires-Dist: tqdm
|
|
30
|
+
Requires-Dist: monai
|
|
31
|
+
Requires-Dist: triton>=2.1; sys_platform != "win32"
|
|
32
|
+
Requires-Dist: triton-windows>=3.0; sys_platform == "win32"
|
|
33
|
+
Provides-Extra: cu128
|
|
34
|
+
Requires-Dist: torch>=2.7.0; extra == "cu128"
|
|
35
|
+
Requires-Dist: torchvision>=0.22.0; extra == "cu128"
|
|
36
|
+
Provides-Extra: cu130
|
|
37
|
+
Requires-Dist: torch>=2.9.0; extra == "cu130"
|
|
38
|
+
Requires-Dist: torchvision>=0.24.0; extra == "cu130"
|
|
39
|
+
Provides-Extra: cpu
|
|
40
|
+
Requires-Dist: torch>=2.5; extra == "cpu"
|
|
41
|
+
Requires-Dist: torchvision>=0.20; extra == "cpu"
|
|
42
|
+
Provides-Extra: data
|
|
43
|
+
Requires-Dist: nibabel; extra == "data"
|
|
44
|
+
Provides-Extra: docs
|
|
45
|
+
Requires-Dist: sphinx; extra == "docs"
|
|
46
|
+
Requires-Dist: myst-parser; extra == "docs"
|
|
47
|
+
Requires-Dist: sphinx-rtd-theme; extra == "docs"
|
|
48
|
+
Provides-Extra: vis
|
|
49
|
+
Requires-Dist: matplotlib; extra == "vis"
|
|
50
|
+
Requires-Dist: seaborn; extra == "vis"
|
|
51
|
+
Requires-Dist: pandas; extra == "vis"
|
|
52
|
+
Provides-Extra: all
|
|
53
|
+
Requires-Dist: flashdrr[data,docs,vis]; extra == "all"
|
|
54
|
+
Dynamic: license-file
|
|
55
|
+
|
|
56
|
+
# FlashDRR
|
|
57
|
+
Fast differentiable DRR raycasting and volume rendering for PyTorch.
|
|
58
|
+
|
|
59
|
+
`FlashDRR` provides fast, differentiable raycasters that turn CT volumes into
|
|
60
|
+
synthesized X‑ray projections (Digitally Reconstructed Radiographs, DRR) while
|
|
61
|
+
keeping gradients flowing through both the volume and the camera parameters.
|
|
62
|
+
It is designed for deep-learning workloads such as 2D/3D registration,
|
|
63
|
+
pose estimation, and synthetic‑image generation from medical volumes.
|
|
64
|
+
|
|
65
|
+
## [Documentation](https://torchvtk.github.io)
|
|
66
|
+
|
|
67
|
+
## Features
|
|
68
|
+
* **Differentiable DRR raycasting** — `flashdrr.rendering.VolumeRaycaster`
|
|
69
|
+
renders physically‑based DRRs using the Beer‑Lambert law with fully
|
|
70
|
+
differentiable ray sampling.
|
|
71
|
+
* **Fast Triton‑fused renderer** — `flashdrr.rendering.FusedVolumeRenderer`
|
|
72
|
+
runs a CUDA kernel for near‑instant, backprop‑free (or autograd-aware)
|
|
73
|
+
projection at high resolution.
|
|
74
|
+
* **C‑arm / camera helpers** — sample realistic C‑arm views
|
|
75
|
+
(`carm_to_camera_params`, `get_random_carm_views`) and build VTK‑compatible
|
|
76
|
+
camera matrices in RAS space (`get_vtk_view_mat`).
|
|
77
|
+
* **Transfer functions** — map CT Hounsfield units to density/attenuation via
|
|
78
|
+
`flashdrr.rendering.piecewise_linear_channelwise`.
|
|
79
|
+
* **Lightweight volume utilities** — `flashdrr.utils` provides dimension
|
|
80
|
+
helpers (`make_nd`), HU normalization, and voxel‑scale handling.
|
|
81
|
+
|
|
82
|
+
## Benchmark
|
|
83
|
+
|
|
84
|
+
The following trends summarize head-to-head measurements of
|
|
85
|
+
`VolumeRaycaster` (PyTorch `grid_sample` based) against the Triton-fused
|
|
86
|
+
`FusedVolumeRenderer` across resolutions, ray-sample counts, and volume
|
|
87
|
+
sizes. All tests performed using an RTX 5070Ti 16GB.
|
|
88
|
+
|
|
89
|
+
<a href="benchmarks/plots/benchmark_analysis.svg">
|
|
90
|
+
<img src="benchmarks/plots/benchmark_analysis.svg"
|
|
91
|
+
alt="Benchmark comparison of the PyTorch-loop VolumeRaycaster against the Triton-fused FusedVolumeRenderer. Left: peak GPU memory (MiB, log scale) versus ray-sample count across resolutions from 128x128 to 2048x2048, showing linear growth for the PyTorch backend (peaking at ~24.8 GB at 2048x2048, 256 samples) and a near-flat profile below ~430 MiB for Triton. Middle: per-iteration latency (ms, log scale) versus resolution, showing super-linear degradation for the PyTorch backend beyond 1024x1024 (3,690 ms at 2048x2048) and roughly linear scaling for Triton (11.8 ms at 2048x2048, a ~314x speedup). Right: train vs eval memory and latency at 1024x1024 with 256 samples, where Triton uses 92.9 MiB eval / 115.2 MiB train versus 6,254.9 / 6,247.2 MiB for PyTorch, and 4.21 ms eval / 23.50 ms train versus 55.9 / 75.37 ms for PyTorch."
|
|
92
|
+
width="900">
|
|
93
|
+
</a>
|
|
94
|
+
|
|
95
|
+
### Key Performance Trends
|
|
96
|
+
|
|
97
|
+
1. **Memory Scaling: Linear Allocation vs. Near-Constant Footprint**
|
|
98
|
+
* **Without Triton:** Peak VRAM scales directly with total samples
|
|
99
|
+
($\text{Rays} \times \text{Samples/Ray}$). At $2048 \times 2048$
|
|
100
|
+
resolution with 256 ray samples, peak memory hits 24.8 GB in
|
|
101
|
+
evaluation mode, leading to out-of-memory (OOM) failures at
|
|
102
|
+
$\geq 384$ samples.
|
|
103
|
+
* **With Triton:** Fused rendering kernels avoid materializing
|
|
104
|
+
intermediate ray-sample tensors in HBM. Memory footprint remains
|
|
105
|
+
below 430 MiB across all tested configurations — a 99.5% memory
|
|
106
|
+
reduction at high workloads.
|
|
107
|
+
* **Volume Size Impact:** Scaling volume grid resolution from
|
|
108
|
+
$128^3$ to $256^3$ adds a fixed baseline memory overhead
|
|
109
|
+
(~100–300 MiB) for tensor storage, but Triton's intermediate
|
|
110
|
+
sample memory remains invariant to volume size.
|
|
111
|
+
|
|
112
|
+
2. **Time Scaling: Sub-Linear Latency vs. Memory-Wall Thrashing**
|
|
113
|
+
* **Extreme Acceleration:** At $2048 \times 2048$ resolution
|
|
114
|
+
(256 samples, eval), Triton reduces iteration time from
|
|
115
|
+
1,804.32 ms to 9.16 ms — a $197\times$ speedup.
|
|
116
|
+
* **Non-Linear Degradation without Triton:** Beyond
|
|
117
|
+
$1024 \times 1024$, standard PyTorch execution time degrades
|
|
118
|
+
super-linearly due to memory bandwidth limits and cache
|
|
119
|
+
thrashing. Triton maintains predictable linear scaling
|
|
120
|
+
proportional to total rays rendered.
|
|
121
|
+
|
|
122
|
+
3. **Train vs. Eval Dynamics**
|
|
123
|
+
* **Memory Overhead:** Without Triton, evaluation and training both
|
|
124
|
+
consume massive memory ($\approx 12.4$ GB at $1024\times 1024$ with
|
|
125
|
+
512 samples). With Triton, training memory increases only slightly
|
|
126
|
+
over evaluation (121.6 MiB train vs 110.7 MiB eval) due to lightweight
|
|
127
|
+
gradient buffer retention.
|
|
128
|
+
* **Backward Pass Latency:** Training mode adds a larger relative
|
|
129
|
+
compute penalty in Triton than in standard PyTorch (~$3.05\times$ overhead vs ~$1.31\times$
|
|
130
|
+
at $1024\times 1024$, 512 samples). This stems
|
|
131
|
+
from executing backward gradient kernels for custom operators,
|
|
132
|
+
though Triton train step time (15.4 ms) remains dramatically
|
|
133
|
+
faster than non-Triton eval time (82.0 ms).
|
|
134
|
+
|
|
135
|
+
### Configuration snapshot (1024×1024, 512 samples)
|
|
136
|
+
|
|
137
|
+
| Configuration (1024×1024, 512 samples) | Triton OFF | Triton ON | Improvement |
|
|
138
|
+
|----------------------------------------|-----------:| ---: |---------------------|
|
|
139
|
+
| Eval Memory (MiB) | 12,386.7 | 110.7 | $111.9\times$ less |
|
|
140
|
+
| Train Memory (MiB) | 12,409.6 | 115.2 | $102.1\times$ less |
|
|
141
|
+
| Eval Latency (ms/iter) | 82.00 | 5.05 | $16.2\times$ faster |
|
|
142
|
+
| Train Latency (ms/iter) | 107.34 | 15.40 | $7.0\times$ faster |
|
|
143
|
+
|
|
144
|
+
## Installation Instructions
|
|
145
|
+
The project is not yet released to PyPI.
|
|
146
|
+
|
|
147
|
+
[//]: # (```)
|
|
148
|
+
|
|
149
|
+
[//]: # (pip install flashdrr)
|
|
150
|
+
|
|
151
|
+
[//]: # (```)
|
|
152
|
+
|
|
153
|
+
To get the latest master:
|
|
154
|
+
```
|
|
155
|
+
pip install git+https://github.com/pcarnah/flashdrr.git@master#egg=flashdrr
|
|
156
|
+
```
|
|
157
|
+
|
|
158
|
+
### CUDA builds
|
|
159
|
+
PyTorch ships separate wheels per CUDA version. By default `pip install
|
|
160
|
+
flashdrr` pulls the CPU build of `torch`. To select a CUDA build, install the
|
|
161
|
+
corresponding extra and point pip at the matching PyTorch index, for example:
|
|
162
|
+
|
|
163
|
+
```bash
|
|
164
|
+
pip install flashdrr[cu128] --extra-index-url https://download.pytorch.org/whl/cu128
|
|
165
|
+
pip install flashdrr[cu130] --extra-index-url https://download.pytorch.org/whl/cu130
|
|
166
|
+
```
|
|
167
|
+
|
|
168
|
+
The Triton-fused renderer requires a working Triton installation (the
|
|
169
|
+
`triton-windows` fork is used automatically on Windows).
|
|
170
|
+
|
|
171
|
+
### Using uv
|
|
172
|
+
[uv](https://docs.astral.sh/uv/) is supported. CUDA build
|
|
173
|
+
selection uses the same `cu128` / `cu130` / `cpu` extras as with pip — pick
|
|
174
|
+
one and the matching PyTorch index is
|
|
175
|
+
selected automatically via `[tool.uv.sources]`:
|
|
176
|
+
|
|
177
|
+
```bash
|
|
178
|
+
uv sync --extra cu128
|
|
179
|
+
uv sync --extra cu130
|
|
180
|
+
uv sync --extra cpu
|
|
181
|
+
```
|
|
182
|
+
|
|
183
|
+
Dev-only tooling (e.g. `pytest`) lives in PEP 735 dependency groups, which
|
|
184
|
+
must not share names with the CUDA extras — add them with `--group`:
|
|
185
|
+
|
|
186
|
+
```bash
|
|
187
|
+
uv sync --extra cu130 --group test
|
|
188
|
+
```
|
|
189
|
+
|
|
190
|
+
The optional `data` extra (NIfTI file support) is just another extra
|
|
191
|
+
and combines with a CUDA extra:
|
|
192
|
+
|
|
193
|
+
```bash
|
|
194
|
+
uv sync --extra cu128 --extra data
|
|
195
|
+
```
|
|
196
|
+
|
|
197
|
+
## Quick example
|
|
198
|
+
```python
|
|
199
|
+
import torch
|
|
200
|
+
import flashdrr.rendering as R
|
|
201
|
+
from flashdrr.rendering import carm_to_camera_params, get_vtk_view_mat
|
|
202
|
+
|
|
203
|
+
# vol: (B, C, D, H, W) attenuation volume in RAS-aligned voxels (e.g. from MONAI)
|
|
204
|
+
# ras2ijk: (4, 4) RAS -> IJK affine of the volume
|
|
205
|
+
vol = torch.rand(1, 1, 128, 128, 128)
|
|
206
|
+
|
|
207
|
+
raycaster = R.VolumeRaycaster(ray_samples=256, resolution=(512, 512)).cuda()
|
|
208
|
+
|
|
209
|
+
center = torch.tensor([64., 64., 64.], device='cuda')
|
|
210
|
+
pos, focal, up = carm_to_camera_params(sid=1000.0, ap_angle=0.0, lat_angle=-30.0,
|
|
211
|
+
center_ras=center, table_si=0.0)
|
|
212
|
+
view_mat = get_vtk_view_mat(pos, focal, up, device='cuda').unsqueeze(0)
|
|
213
|
+
|
|
214
|
+
# (1, 1, H, W) DRR projection
|
|
215
|
+
drr = raycaster(vol.cuda(), view_mat=view_mat, ras2ijk=ras2ijk.cuda())
|
|
216
|
+
```
|
|
217
|
+
|
|
218
|
+
## Optional dependencies
|
|
219
|
+
* `nibabel` — required only for loading Nifti format data for tests and example scripts
|
|
220
|
+
(`flashdrr[data]`).
|
|
221
|
+
|
|
222
|
+
## Attribution
|
|
223
|
+
If you use this software in your research, please cite it using the "Cite this repository" button or the provided CITATION.cff.
|
flashdrr-0.5/README.md
ADDED
|
@@ -0,0 +1,168 @@
|
|
|
1
|
+
# FlashDRR
|
|
2
|
+
Fast differentiable DRR raycasting and volume rendering for PyTorch.
|
|
3
|
+
|
|
4
|
+
`FlashDRR` provides fast, differentiable raycasters that turn CT volumes into
|
|
5
|
+
synthesized X‑ray projections (Digitally Reconstructed Radiographs, DRR) while
|
|
6
|
+
keeping gradients flowing through both the volume and the camera parameters.
|
|
7
|
+
It is designed for deep-learning workloads such as 2D/3D registration,
|
|
8
|
+
pose estimation, and synthetic‑image generation from medical volumes.
|
|
9
|
+
|
|
10
|
+
## [Documentation](https://torchvtk.github.io)
|
|
11
|
+
|
|
12
|
+
## Features
|
|
13
|
+
* **Differentiable DRR raycasting** — `flashdrr.rendering.VolumeRaycaster`
|
|
14
|
+
renders physically‑based DRRs using the Beer‑Lambert law with fully
|
|
15
|
+
differentiable ray sampling.
|
|
16
|
+
* **Fast Triton‑fused renderer** — `flashdrr.rendering.FusedVolumeRenderer`
|
|
17
|
+
runs a CUDA kernel for near‑instant, backprop‑free (or autograd-aware)
|
|
18
|
+
projection at high resolution.
|
|
19
|
+
* **C‑arm / camera helpers** — sample realistic C‑arm views
|
|
20
|
+
(`carm_to_camera_params`, `get_random_carm_views`) and build VTK‑compatible
|
|
21
|
+
camera matrices in RAS space (`get_vtk_view_mat`).
|
|
22
|
+
* **Transfer functions** — map CT Hounsfield units to density/attenuation via
|
|
23
|
+
`flashdrr.rendering.piecewise_linear_channelwise`.
|
|
24
|
+
* **Lightweight volume utilities** — `flashdrr.utils` provides dimension
|
|
25
|
+
helpers (`make_nd`), HU normalization, and voxel‑scale handling.
|
|
26
|
+
|
|
27
|
+
## Benchmark
|
|
28
|
+
|
|
29
|
+
The following trends summarize head-to-head measurements of
|
|
30
|
+
`VolumeRaycaster` (PyTorch `grid_sample` based) against the Triton-fused
|
|
31
|
+
`FusedVolumeRenderer` across resolutions, ray-sample counts, and volume
|
|
32
|
+
sizes. All tests performed using an RTX 5070Ti 16GB.
|
|
33
|
+
|
|
34
|
+
<a href="benchmarks/plots/benchmark_analysis.svg">
|
|
35
|
+
<img src="benchmarks/plots/benchmark_analysis.svg"
|
|
36
|
+
alt="Benchmark comparison of the PyTorch-loop VolumeRaycaster against the Triton-fused FusedVolumeRenderer. Left: peak GPU memory (MiB, log scale) versus ray-sample count across resolutions from 128x128 to 2048x2048, showing linear growth for the PyTorch backend (peaking at ~24.8 GB at 2048x2048, 256 samples) and a near-flat profile below ~430 MiB for Triton. Middle: per-iteration latency (ms, log scale) versus resolution, showing super-linear degradation for the PyTorch backend beyond 1024x1024 (3,690 ms at 2048x2048) and roughly linear scaling for Triton (11.8 ms at 2048x2048, a ~314x speedup). Right: train vs eval memory and latency at 1024x1024 with 256 samples, where Triton uses 92.9 MiB eval / 115.2 MiB train versus 6,254.9 / 6,247.2 MiB for PyTorch, and 4.21 ms eval / 23.50 ms train versus 55.9 / 75.37 ms for PyTorch."
|
|
37
|
+
width="900">
|
|
38
|
+
</a>
|
|
39
|
+
|
|
40
|
+
### Key Performance Trends
|
|
41
|
+
|
|
42
|
+
1. **Memory Scaling: Linear Allocation vs. Near-Constant Footprint**
|
|
43
|
+
* **Without Triton:** Peak VRAM scales directly with total samples
|
|
44
|
+
($\text{Rays} \times \text{Samples/Ray}$). At $2048 \times 2048$
|
|
45
|
+
resolution with 256 ray samples, peak memory hits 24.8 GB in
|
|
46
|
+
evaluation mode, leading to out-of-memory (OOM) failures at
|
|
47
|
+
$\geq 384$ samples.
|
|
48
|
+
* **With Triton:** Fused rendering kernels avoid materializing
|
|
49
|
+
intermediate ray-sample tensors in HBM. Memory footprint remains
|
|
50
|
+
below 430 MiB across all tested configurations — a 99.5% memory
|
|
51
|
+
reduction at high workloads.
|
|
52
|
+
* **Volume Size Impact:** Scaling volume grid resolution from
|
|
53
|
+
$128^3$ to $256^3$ adds a fixed baseline memory overhead
|
|
54
|
+
(~100–300 MiB) for tensor storage, but Triton's intermediate
|
|
55
|
+
sample memory remains invariant to volume size.
|
|
56
|
+
|
|
57
|
+
2. **Time Scaling: Sub-Linear Latency vs. Memory-Wall Thrashing**
|
|
58
|
+
* **Extreme Acceleration:** At $2048 \times 2048$ resolution
|
|
59
|
+
(256 samples, eval), Triton reduces iteration time from
|
|
60
|
+
1,804.32 ms to 9.16 ms — a $197\times$ speedup.
|
|
61
|
+
* **Non-Linear Degradation without Triton:** Beyond
|
|
62
|
+
$1024 \times 1024$, standard PyTorch execution time degrades
|
|
63
|
+
super-linearly due to memory bandwidth limits and cache
|
|
64
|
+
thrashing. Triton maintains predictable linear scaling
|
|
65
|
+
proportional to total rays rendered.
|
|
66
|
+
|
|
67
|
+
3. **Train vs. Eval Dynamics**
|
|
68
|
+
* **Memory Overhead:** Without Triton, evaluation and training both
|
|
69
|
+
consume massive memory ($\approx 12.4$ GB at $1024\times 1024$ with
|
|
70
|
+
512 samples). With Triton, training memory increases only slightly
|
|
71
|
+
over evaluation (121.6 MiB train vs 110.7 MiB eval) due to lightweight
|
|
72
|
+
gradient buffer retention.
|
|
73
|
+
* **Backward Pass Latency:** Training mode adds a larger relative
|
|
74
|
+
compute penalty in Triton than in standard PyTorch (~$3.05\times$ overhead vs ~$1.31\times$
|
|
75
|
+
at $1024\times 1024$, 512 samples). This stems
|
|
76
|
+
from executing backward gradient kernels for custom operators,
|
|
77
|
+
though Triton train step time (15.4 ms) remains dramatically
|
|
78
|
+
faster than non-Triton eval time (82.0 ms).
|
|
79
|
+
|
|
80
|
+
### Configuration snapshot (1024×1024, 512 samples)
|
|
81
|
+
|
|
82
|
+
| Configuration (1024×1024, 512 samples) | Triton OFF | Triton ON | Improvement |
|
|
83
|
+
|----------------------------------------|-----------:| ---: |---------------------|
|
|
84
|
+
| Eval Memory (MiB) | 12,386.7 | 110.7 | $111.9\times$ less |
|
|
85
|
+
| Train Memory (MiB) | 12,409.6 | 115.2 | $102.1\times$ less |
|
|
86
|
+
| Eval Latency (ms/iter) | 82.00 | 5.05 | $16.2\times$ faster |
|
|
87
|
+
| Train Latency (ms/iter) | 107.34 | 15.40 | $7.0\times$ faster |
|
|
88
|
+
|
|
89
|
+
## Installation Instructions
|
|
90
|
+
The project is not yet released to PyPI.
|
|
91
|
+
|
|
92
|
+
[//]: # (```)
|
|
93
|
+
|
|
94
|
+
[//]: # (pip install flashdrr)
|
|
95
|
+
|
|
96
|
+
[//]: # (```)
|
|
97
|
+
|
|
98
|
+
To get the latest master:
|
|
99
|
+
```
|
|
100
|
+
pip install git+https://github.com/pcarnah/flashdrr.git@master#egg=flashdrr
|
|
101
|
+
```
|
|
102
|
+
|
|
103
|
+
### CUDA builds
|
|
104
|
+
PyTorch ships separate wheels per CUDA version. By default `pip install
|
|
105
|
+
flashdrr` pulls the CPU build of `torch`. To select a CUDA build, install the
|
|
106
|
+
corresponding extra and point pip at the matching PyTorch index, for example:
|
|
107
|
+
|
|
108
|
+
```bash
|
|
109
|
+
pip install flashdrr[cu128] --extra-index-url https://download.pytorch.org/whl/cu128
|
|
110
|
+
pip install flashdrr[cu130] --extra-index-url https://download.pytorch.org/whl/cu130
|
|
111
|
+
```
|
|
112
|
+
|
|
113
|
+
The Triton-fused renderer requires a working Triton installation (the
|
|
114
|
+
`triton-windows` fork is used automatically on Windows).
|
|
115
|
+
|
|
116
|
+
### Using uv
|
|
117
|
+
[uv](https://docs.astral.sh/uv/) is supported. CUDA build
|
|
118
|
+
selection uses the same `cu128` / `cu130` / `cpu` extras as with pip — pick
|
|
119
|
+
one and the matching PyTorch index is
|
|
120
|
+
selected automatically via `[tool.uv.sources]`:
|
|
121
|
+
|
|
122
|
+
```bash
|
|
123
|
+
uv sync --extra cu128
|
|
124
|
+
uv sync --extra cu130
|
|
125
|
+
uv sync --extra cpu
|
|
126
|
+
```
|
|
127
|
+
|
|
128
|
+
Dev-only tooling (e.g. `pytest`) lives in PEP 735 dependency groups, which
|
|
129
|
+
must not share names with the CUDA extras — add them with `--group`:
|
|
130
|
+
|
|
131
|
+
```bash
|
|
132
|
+
uv sync --extra cu130 --group test
|
|
133
|
+
```
|
|
134
|
+
|
|
135
|
+
The optional `data` extra (NIfTI file support) is just another extra
|
|
136
|
+
and combines with a CUDA extra:
|
|
137
|
+
|
|
138
|
+
```bash
|
|
139
|
+
uv sync --extra cu128 --extra data
|
|
140
|
+
```
|
|
141
|
+
|
|
142
|
+
## Quick example
|
|
143
|
+
```python
|
|
144
|
+
import torch
|
|
145
|
+
import flashdrr.rendering as R
|
|
146
|
+
from flashdrr.rendering import carm_to_camera_params, get_vtk_view_mat
|
|
147
|
+
|
|
148
|
+
# vol: (B, C, D, H, W) attenuation volume in RAS-aligned voxels (e.g. from MONAI)
|
|
149
|
+
# ras2ijk: (4, 4) RAS -> IJK affine of the volume
|
|
150
|
+
vol = torch.rand(1, 1, 128, 128, 128)
|
|
151
|
+
|
|
152
|
+
raycaster = R.VolumeRaycaster(ray_samples=256, resolution=(512, 512)).cuda()
|
|
153
|
+
|
|
154
|
+
center = torch.tensor([64., 64., 64.], device='cuda')
|
|
155
|
+
pos, focal, up = carm_to_camera_params(sid=1000.0, ap_angle=0.0, lat_angle=-30.0,
|
|
156
|
+
center_ras=center, table_si=0.0)
|
|
157
|
+
view_mat = get_vtk_view_mat(pos, focal, up, device='cuda').unsqueeze(0)
|
|
158
|
+
|
|
159
|
+
# (1, 1, H, W) DRR projection
|
|
160
|
+
drr = raycaster(vol.cuda(), view_mat=view_mat, ras2ijk=ras2ijk.cuda())
|
|
161
|
+
```
|
|
162
|
+
|
|
163
|
+
## Optional dependencies
|
|
164
|
+
* `nibabel` — required only for loading Nifti format data for tests and example scripts
|
|
165
|
+
(`flashdrr[data]`).
|
|
166
|
+
|
|
167
|
+
## Attribution
|
|
168
|
+
If you use this software in your research, please cite it using the "Cite this repository" button or the provided CITATION.cff.
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
__version__ = '0.5'
|
|
@@ -0,0 +1,36 @@
|
|
|
1
|
+
"""Differentiable volume rendering for PyTorch.
|
|
2
|
+
|
|
3
|
+
The package exposes two renderers and a handful of camera / geometry helpers:
|
|
4
|
+
|
|
5
|
+
* :class:`VolumeRaycaster` - the main Python-loop differentiable raycaster
|
|
6
|
+
(optionally with the :class:`DepthAwareScatter` and :class:`ASPP` building
|
|
7
|
+
blocks for scatter-based DRR rendering).
|
|
8
|
+
* :class:`FusedVolumeRenderer` - a Triton-fused forward kernel for fast
|
|
9
|
+
backprop-free rendering on CUDA.
|
|
10
|
+
|
|
11
|
+
The camera helpers (:func:`get_view_mat`, :func:`get_vtk_view_mat`,
|
|
12
|
+
:func:`get_proj_mat`, :func:`get_random_carm_views`, ...) are the same ones
|
|
13
|
+
the renderers use internally and are re-exported for downstream code.
|
|
14
|
+
"""
|
|
15
|
+
|
|
16
|
+
from .raycast import (
|
|
17
|
+
ASPP,
|
|
18
|
+
DepthAwareScatter,
|
|
19
|
+
VolumeRaycaster,
|
|
20
|
+
carm_to_camera_params,
|
|
21
|
+
get_random_carm_views,
|
|
22
|
+
get_vtk_view_mat,
|
|
23
|
+
piecewise_linear_channelwise,
|
|
24
|
+
)
|
|
25
|
+
from .triton_raycast import FusedVolumeRenderer
|
|
26
|
+
|
|
27
|
+
__all__ = [
|
|
28
|
+
'ASPP',
|
|
29
|
+
'DepthAwareScatter',
|
|
30
|
+
'FusedVolumeRenderer',
|
|
31
|
+
'VolumeRaycaster',
|
|
32
|
+
'carm_to_camera_params',
|
|
33
|
+
'get_random_carm_views',
|
|
34
|
+
'get_vtk_view_mat',
|
|
35
|
+
'piecewise_linear_channelwise',
|
|
36
|
+
]
|