CaiYuanhao/SAX-NeRF

Structure-Aware Sparse-View X-ray 3D Reconstruction

CVPR 2024

15 scenes · 4 application domains · Sparse-view CT reconstruction · X-ray novel view synthesis

Paper · Code · Model · Download Data · Video · 中文解读

3D Reconstruction Gallery

Backpack	Box	Bonsai
Foot	Teapot	Engine

Reconstruction visualizations from the official SAX-NeRF repository; the downloadable dataset consists of processed scene arrays.

Model Description

SAX-NeRF (Structure-Aware X-ray Neural Radiodensity Fields) reconstructs 3D internal structures from sparse-view X-ray projections. Introduced in Structure-Aware Sparse-View X-ray 3D Reconstruction (CVPR 2024), it supports both X-ray novel view synthesis and CT reconstruction.

The method combines a Line Segment-based Transformer (Lineformer), which models dependencies along X-ray segments, with Masked Local-Global (MLG) ray sampling, which captures local context and global geometry in projections. The network maps 3D positions to radiodensity, enabling projection rendering and volume reconstruction. See the paper.

This repository provides 15 scene-specific checkpoints for the X3D dataset. Each checkpoint represents its corresponding scene. To reconstruct a new scene, train a model using that scene's projections and scanning geometry.

Available Checkpoints

Match the checkpoint, processed data, and configuration for the same scene:

Scene	Checkpoint	Dataset file	Configuration
Abdomen	abdomen.tar	abdomen_50.pickle	abdomen_50.yaml
Aneurism	aneurism.tar	aneurism_50.pickle	aneurism_50.yaml
Backpack	backpack.tar	backpack_50.pickle	backpack_50.yaml
Bonsai	bonsai.tar	bonsai_50.pickle	bonsai_50.yaml
Box	box.tar	box_50.pickle	box_50.yaml
Carp	carp.tar	carp_50.pickle	carp_50.yaml
Chest	chest.tar	chest_50.pickle	chest_50.yaml
Engine	engine.tar	engine_50.pickle	engine_50.yaml
Foot	foot.tar	foot_50.pickle	foot_50.yaml
Head	head.tar	head_50.pickle	head_50.yaml
Jaw	jaw.tar	jaw_50.pickle	jaw_50.yaml
Leg	leg.tar	leg_50.pickle	leg_50.yaml
Pancreas	pancreas.tar	pancreas_50.pickle	pancreas_50.yaml
Pelvis	pelvis.tar	pelvis_50.pickle	pelvis_50.yaml
Teapot	teapot.tar	teapot_50.pickle	teapot_50.yaml

The .tar files are PyTorch checkpoints loaded directly by the project's test.py; no manual extraction is required. The evaluation code reads the network state dictionary and the saved epoch from the checkpoint.

Quick Start

1. Set Up the Project

git clone https://github.com/caiyuanhao1998/SAX-NeRF.git
cd SAX-NeRF

Follow the official environment setup to install PyTorch, the CUDA-dependent hash encoder requirements, and the project dependencies. The released evaluation script uses an NVIDIA GPU with CUDA.

Install the Hugging Face download client in that environment:

pip install huggingface_hub

2. Download a Checkpoint and Its Scene Data

Run this Python snippet from the project root:

from huggingface_hub import hf_hub_download

scene = "chest"

hf_hub_download(
    repo_id="CaiYuanhao/SAX-NeRF",
    filename=f"{scene}.tar",
    local_dir="pretrained",
)

hf_hub_download(
    repo_id="CaiYuanhao/SAX-NeRF",
    repo_type="dataset",
    filename=f"{scene}_50.pickle",
    local_dir="data",
)

The resulting files should be:

SAX-NeRF/
├── config/Lineformer/chest_50.yaml
├── data/chest_50.pickle
├── pretrained/chest.tar
└── test.py

To download all checkpoints instead:

from huggingface_hub import snapshot_download

snapshot_download(
    repo_id="CaiYuanhao/SAX-NeRF",
    allow_patterns=["*.tar"],
    local_dir="pretrained",
)

3. Render Projections and Reconstruct CT

python test.py \
    --gpu_id 0 \
    --method Lineformer \
    --category chest \
    --config config/Lineformer/chest_50.yaml \
    --weights pretrained/chest.tar \
    --output_path output

For another scene, replace chest in the category, configuration, checkpoint, and downloaded dataset filename. Keep --method Lineformer for SAX-NeRF.

The evaluation script renders the held-out projections, reconstructs the volume on the scene's voxel grid, and reports projection and 3D reconstruction PSNR/SSIM. It saves predicted and reference projection images and CT slices under:

output/Lineformer/chest/
├── proj_pred/
├── proj_gt/
└── CT/
    ├── H/
    ├── W/
    └── L/

Training

To train SAX-NeRF on a released scene:

python train_mlg.py --config config/Lineformer/chest_50.yaml

The configuration specifies the data path, Lineformer architecture, hash-grid encoder, ray sampling, rendering, and optimization settings. Use the corresponding scene configuration when reproducing results, and record any changes. For new scenes, see the project's data generation instructions.

Evaluation on X3D

The paper evaluates 15 scenes, with 50 training projections and 50 held-out projections per scene. The reference CT volume is used for CT reconstruction evaluation. The table below gives paper-reported averages across scenes, rather than a new evaluation of the hosted weights.

Task	PSNR ↑	SSIM ↑
X-ray novel view synthesis	51.37 dB	0.9994
CT reconstruction	37.25 dB	0.9753

Source: SAX-NeRF paper, Tables 1 and 2. For checkpoint evaluation, use the released configuration and dataset alongside the official test script.

Visual Results

X-ray Novel View Synthesis

More novel view synthesis results

CT Reconstruction

More CT reconstruction results

Comparison figures reproduced from the official project repository.

Citation

If you use these models, please cite the SAX-NeRF paper. The entry below follows the citation supplied by the official code repository:

@inproceedings{sax_nerf,
  title={Structure-Aware Sparse-View X-ray 3D Reconstruction},
  author={Yuanhao Cai and Jiahao Wang and Alan Yuille and Zongwei Zhou and Angtian Wang},
  booktitle={CVPR},
  year={2024}
}