CTTok

A discrete autoregressive approach that extends the vocabulary of a pre-trained language model with anatomical tokens representing 3D CT patches.

Fig.1 in manuscript

All training scripts, data loaders, and configuration files are located in the scripts/ directory.

To start training or reproduce our results, please navigate to the scripts folder:

cd scripts

This directory contains the core scripts for training the image denoising model. The pipeline is designed to learn the mapping from blurry, degraded inputs to clear, high-fidelity outputs, utilizing masks as structural conditions where applicable.

• train.py: The main entry point for the training loop. It initializes the model, optimizer, and handles the forward/backward passes.
• train_config.json: The configuration file containing all hyperparameters (e.g., learning rate, batch size, epoch count) and data paths.
• dataloader.py: Handles the critical data pipeline. It reads the image data, processes the blurry-to-clear image pairs, applies data augmentations, and yields the condition masks.
• muon_fixed.py: Contains the implementation of the custom optimizer or specific mathematical operations used during training.
• utils.py: A collection of helper functions (e.g., logging, saving checkpoints, calculating metrics like FID).
• train_intro.md: This documentation file.

• Configuration Paths: Before running the code, please open train_config.json and ensure you manually update all the required paths (e.g., dataset directories, output logging paths) to match your local environment.
• Pre-trained Weights: The complete set of pre-trained weights for this project will be publicly released on Hugging Face upon the acceptance of the paper.

We use Hugging Face accelerate for distributed training. Make sure your environment is properly set up before launching.

To start the training process, run the following command from the terminal. Replace <num_gpus> with the actual number of GPUs you want to use (e.g., 4 or 8):

accelerate launch \
  --multi_gpu \
  --num_processes <num_gpus> \
  --mixed_precision bf16 \
  train.py \
  --config train_config.json \
  --auto_resume

Additional visualization results and comparison figures can be found in the assets directory:

cd assets