TRACE

KO | EN

gitlite — search

#python #java #python3 #arduino #golang #machine-learning #rust #html #flask #javascript #seismology #nodejs

★ 12 Open GitHub ↗

Offical Repo of "TRACE: Temporally Reliable Anatomically-Conditioned 3D CT Generation with Enhanced Efficiency" (MICCAI 2025)

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

# *MICCAI 2025*
# **TRACE: Temporally Reliable Anatomically-Conditioned 3D CT Generation with Enhanced Efficiency**

[![paper](https://img.shields.io/badge/arXiv-Paper-red.svg)](https://arxiv.org/pdf/2507.00802)
[![Overview](https://img.shields.io/badge/Overview-Read-orange.svg)](https://raw.githubusercontent.com/VinyehShaw/TRACE/main/)
[![1. Download](https://img.shields.io/badge/Datasets-Download-yellow.svg)](https://raw.githubusercontent.com/VinyehShaw/TRACE/main/)
[![Models](https://img.shields.io/badge/Training-Inference-purple.svg)](https://raw.githubusercontent.com/VinyehShaw/TRACE/main/)
[![BibTeX](https://img.shields.io/badge/BibTeX-Cite-blueviolet.svg)](https://raw.githubusercontent.com/VinyehShaw/TRACE/main/)

Official implementation of the  Medical Image Computing and Computer Assisted Intervention – MICCAI 2025 paper [TRACE: Temporally Reliable Anatomically-Conditioned 3D CT Generation with Enhanced Efficiency](https://link.springer.com/chapter/10.1007/978-3-032-04965-0_59), provided by [Minye Shao](https://www.linkedin.com/in/minyeshao/).

## 🩻 TL;DR

### 🚨 Problems in Prior 3D Medical Image Generation
1. **Limited Anatomical Fidelity** — Existing 3D diffusion models often produce anatomically inconsistent CT volumes, showing **organ displacement, missing or duplicated structures, and discontinuity along the Z-axis**.  
2. **Fixed Cubic Assumption** — Conventional **isotropic 3D diffusion** constrains axial length and ignores real-world **anisotropic CT shapes**, leading to poor spatial representation.  
3. **High Computational Demand** — Full 3D diffusion incurs **heavy GPU memory usage**, slow convergence, and limited scalability in clinical or low-resource settings.  
4. **Inadequate Evaluation** — Most prior works rely solely on **generic FID/FVD metrics**, **lacking medical-level anatomical assessments** to measure true fidelity and diagnostic reliability.

---

### 💡 Our Solution: **TRACE**
**TRACE** redefines 3D CT synthesis as a **video generation task** using a **2D multimodal-conditioned diffusion model**.  
It models 3D volumes as sequential 2D slices (“frames”), guided by **segmentation masks**, **radiology reports**, **optical flow**, and **temporal position embeddings**, then reconstructs coherent 3D volumes via **Overlapping Frame Guidance**.
By bridging **video diffusion paradigms** with **medical priors**, it enables **flexible-length CT synthesis** suitable for data augmentation, privacy-preserving research, and personalized modeling in clinical workflows.

## Overview
<div align=center>  
<img src='figs/ov.png' width="100%">
</div>
 During training, it denoises frame pairs with varying skip intervals, guided by four modalities: anatomical masks (VISTA3D), optical flow between frames (RAFT), report embeddings (CLIP), and relative position embeddings. The optical flow and text embeddings pass through trainable adapters before entering the diffusion model. Inference employs an overlapping-frame guidance strategy to synthesize semantically aligned frame pairs, generating anatomically consistent CT sequences, which are then reconstructed back to 3D volumes.

Some experimental results for illustration. For more details, please refer to our paper.

---

## 🛠️ Install Dependencies
> ✅ Tested on Ubuntu 20.04 + Pytorch 2.1.2 + Python 3.11

Please follow the pip list file `requirements.txt`.

Extract optical flow between frames. For installation, please follow [RAFT's](https://github.com/princeton-vl/RAFT) offical repo.

For 3D volume segmentation, please follow the official [VISTA-3D](https://github.com/Project-MONAI/VISTA/tree/main/vista3d) repo for installation.

For VBench evaluation, please follow the installation instructions in their official [VBench repo](https://github.com/Vchitect/VBench).

## 📦 Datasets Preparation
### 1. Download
 Datasets Preparation
Please download [CT-RATE](https://huggingface.co/datasets/ibrahimhamamci/CT-RATE/tree/main/dataset) dataset from the Hugging face.

Please select the desired patient CT from either the [training](https://huggingface.co/datasets/ibrahimhamamci/CT-RATE/tree/main/dataset/train) or [validation](https://huggingface.co/datasets/ibrahimhamamci/CT-RATE/tree/main/dataset/valid) dataset.

> Please note that some data contain inconsistent Window Width/Window Level settings or are non-thoracic CT volumes. Please avoid these cases.

> 🆕 The CT-RATE team has now provided a [list of non-thoracic CT](https://huggingface.co/datasets/ibrahimhamamci/CT-RATE/blob/main/dataset/metadata/no_chest_train.txt) volumes.
The [metadata file](https://huggingface.co/datasets/ibrahimhamamci/CT-RATE/blob/main/dataset/metadata/train_metadata.csv) contains information such as z-spacing, patient age, and sex. The corresponding clinical reports for each volume can be found [here](https://huggingface.co/datasets/ibrahimhamamci/CT-RATE/blob/main/dataset/radiology_text_reports/train_reports.csv), download them.

### 2. Preparation
After downloading the volumes, follow VISTA-3D’s official documentation to run segmentation. To retain specific anatomy (organs, bones, tissues), choose the corresponding classes using the provided [label dictionary](https://github.com/Project-MONAI/VISTA/blob/main/vista3d/data/jsons/label_dict.json).

Then split each patient’s CT and its segmentation volume into 2D slices and save them, following the format in the provided `example/frame_figs` and `/frame_segs` folder. 
> When slicing the original CT volumes, do not apply any additional processing such as normalization or windowing, as this may cause slice-to-slice intensity inconsistency and alter the original HU values.

Next, extract inter-frame optical flow at temporal strides 1, 2, and 4 and store it for later training. We use RAFT with Sintel weights by default, but you may try other pretrained flow models; save results following the `example/frame_flows` format.

For each CT volume, follow [GenerateCT’s](https://github.com/ibrahimethemhamamci/GenerateCT) processing description; combine metadata (age, sex) with the report to produce “{age} years old {sex}: {impression}”, for reference please see ours `example/example_summary.csv`.

Finally, specify the paths for splitting your dataset into training and inference subsets in a JSON file, following the format shown in `example/config/split.json`.

## Training
Run command line like this
```bash 
python3  main.py --mode train --json_dir your_path_to/frame_yours.json --img_dir your_path_to/frame_figs --flow_dir_512 your_path_to/frame_flows --seg_dir your_path_to/frame_segs --model_type DDIM --num_epochs 1200 
```

`--json_dir your_path_to/frame_yours.json`: path to trian/inference split JSON file

`--img_dir your_path_to/frame_figs`: directory of input image frames data

` --flow_dir_512 your_path_to/frame_flows`: directory of prepared optical flows

`--seg_dir your_path_to/frame_segs`: directory of segmentation maps

`--mixed_precision`: set to `fp32` by default; you can try half precision (`fp16`), though in our experiments it did not yield satisfactory results.

## ⚙️ Inference
After training, your model will be automatically saved in:

> output_dir = '{}-{}-{}'.format(model_type.lower(), dataset, img_size)

To run Inference, execute the following command:

```bash
python3 main.py --mode eval
```
Set the checkpoint by loading from your output directory:
`model = model.from_pretrained(os.path.join('your_path_to.safetensors','unet'), use_safetensors=True)
checkpoint = torch.load(os.path.join('your_model_path', 'optimizer_scheduler.pt'))`

Then, uncomment the corresponding evaluation block in the `training.py`:

```python
# model.eval()
#     ...
```
During inference, if the axial (z) dimension is too long, you may need to adjust it manually, as adjacent slices might become visibly misaligned.

## Citation

We are among the first to **explore 3D CT generation through 2D multimodal-conditioned diffusion**, aiming to achieve anatomically faithful and temporally coherent results with improved computational efficiency.
There is still **room for improvement**, future work may explore **flow matching**, **ControlNet-based fine-tuning** of pretrained diffusion models (e.g., Stable Diffusion), and better **overlapping-frame consistency control**. We also **welcome extensions** of TRACE to other volumetric modalities such as **brain or cardiac MRI** to further assess its generalization.

> **If you find our work helpful in your research, clinical tool development, or inspiring for your own generative modeling efforts, please consider citing us:**

```
@inproceedings{shao2025trace,
  title={TRACE: Temporally Reliable Anatomically-Conditioned 3D CT Generation with Enhanced Efficiency},
  author={Shao, Minye and Miao, Xingyu and Duan, Haoran and Wang, Zeyu and Chen, Jingkun and Huang, Yawen and Wu, Xian and Deng, Jingjing and Long, Yang and Zheng, Yefeng},
  booktitle={International Conference on Medical Image Computing and Computer-Assisted Intervention},
  pages={627--637},
  year={2025},
  organization={Springer}
}
```

## Acknowledgments
This work was inspired by [Konz, N. et al](https://papers.miccai.org/miccai-2024/075-Paper0584.html).  We also sincerely thank the developers of [**RAFT**](https://arxiv.org/pdf/2003.12039), [**VISTA3D**](https://arxiv.org/pdf/2406.05285), [**VBench**](https://arxiv.org/abs/2311.17982), and [**GenerateCT**](https://arxiv.org/abs/2305.16037) for their outstanding contributions that made this research possible.

TRACE

Explore Similar Repositories

Was this content helpful?

Recent Feedback

Download README