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

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:

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

The root config.json lists this repository's checkpoints, supporting files, and implementation. Include it when downloading selected weights so those downloads can be counted by the Hugging Face Hub. Use the linked implementation and its matching runtime configuration to load the checkpoints.

Run this Python snippet from the project root:

import json
from pathlib import Path
from huggingface_hub import hf_hub_download

scene = "chest"

config_path = hf_hub_download(
    repo_id="CaiYuanhao/SAX-NeRF",
    filename="config.json",
    local_dir="pretrained",
)
checkpoint_files = json.loads(Path(config_path).read_text())["checkpoint_files"]
checkpoint = next(name for name in checkpoint_files if name == f"{scene}.tar")
hf_hub_download(
    repo_id="CaiYuanhao/SAX-NeRF",
    filename=checkpoint,
    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=["config.json", "*.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.

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}
}