Skip to content

Aayushjshah/MRI_reconstruction

Repository files navigation

FastMRI Reconstruction with Pretrained Models

A comprehensive implementation for MRI reconstruction using pretrained deep learning models. This project supports XPDNet (TensorFlow/Keras) and VarNet (PyTorch/ATOMMIC) for brain multicoil data with 4× acceleration. No training required - uses inference-only with pretrained weights.

Features

  • XPDNet (TensorFlow/Keras) for brain multicoil 4× acceleration
  • VarNet path with ATOMMIC + Hugging Face checkpoints (PyTorch)
  • ✅ FastMRI dataset handling and preprocessing
  • ✅ Automatic sensitivity map estimation
  • ✅ Easy demo functionality with visualization
  • ✅ Checkpoint swapping capabilities
  • ✅ Comprehensive output formats (NIfTI, NumPy, PNG)

Quick Start

1. Installation

# Clone or download the project files
git clone <your-repo> # or download reconstructFastMRI.py + requirements.txt

# Install dependencies
pip install -r requirements.txt

# Optional: For VarNet path (PyTorch + ATOMMIC)
pip install torch torchvision transformers
pip install git+https://github.com/wdika/atommic.git

2. Run Demo

# Run demo with XPDNet (downloads sample data automatically)
python reconstructFastMRI.py --demo

# Run demo with VarNet
python reconstructFastMRI.py --demo --method varnet

3. Use Your Own Data

# Reconstruct with XPDNet
python reconstructFastMRI.py --h5 path/to/your/brain_multicoil.h5

# Reconstruct with VarNet
python reconstructFastMRI.py --h5 path/to/your/brain_multicoil.h5 --method varnet

# Custom output directory and acceleration
python reconstructFastMRI.py --h5 data.h5 --output my_results --acceleration 8

FastMRI Dataset

This project works with the public fastMRI dataset. To get access:

  1. Visit fastMRI website
  2. Request access (free for research)
  3. Download brain multicoil data
  4. Use with this reconstruction tool

Expected data format: .h5 files with kspace dataset containing multicoil k-space data [slices, coils, ky, kx].

Usage Examples

Basic Reconstruction

from reconstructFastMRI import FastMRIReconstructor

# Initialize reconstructor
reconstructor = FastMRIReconstructor(method='xpdnet')

# Load pretrained weights
reconstructor.load_pretrained_weights()

# Run reconstruction
results = reconstructor.reconstruct_volume(
    h5_path='brain_multicoil_sample.h5',
    output_dir='reconstruction_output',
    acceleration=4
)

print(f"Results saved to: {results['nifti_path']}")

Advanced Usage

# Custom checkpoint
reconstructor = FastMRIReconstructor(method='xpdnet')
reconstructor.load_pretrained_weights('/path/to/custom/checkpoint.h5')

# Different acceleration factors
results = reconstructor.reconstruct_volume(
    h5_path='data.h5',
    output_dir='results_8x',
    acceleration=8,
    center_fraction=0.04  # Smaller center for higher acceleration
)

VarNet with PyTorch

# Using VarNet (requires PyTorch + ATOMMIC)
reconstructor = FastMRIReconstructor(method='varnet')
reconstructor.load_pretrained_weights()

results = reconstructor.reconstruct_volume(
    h5_path='data.h5',
    output_dir='varnet_results'
)

Command Line Interface

Basic Commands

# Minimal usage
python reconstructFastMRI.py --h5 data.h5

# Specify output directory
python reconstructFastMRI.py --h5 data.h5 --output my_results

# Choose reconstruction method
python reconstructFastMRI.py --h5 data.h5 --method xpdnet  # or varnet

# Custom acceleration
python reconstructFastMRI.py --h5 data.h5 --acceleration 8 --center-fraction 0.04

All Options

python reconstructFastMRI.py \
    --h5 path/to/data.h5 \                    # Input fastMRI h5 file
    --output ./results \                       # Output directory
    --method xpdnet \                         # xpdnet or varnet
    --acceleration 4 \                        # Acceleration factor
    --center-fraction 0.08 \                  # Center fraction for mask
    --checkpoint path/to/custom/weights.h5 \  # Custom checkpoint
    --demo                                    # Run demo mode

Output Files

After reconstruction, you'll find:

output_directory/
├── reconstruction_xpdnet.nii.gz     # NIfTI volume
├── reconstruction_xpdnet.npy        # NumPy array
├── comparison_xpdnet.png            # Comparison plot
└── slices/                          # Individual slice images
    ├── slice_000_xpdnet.png
    ├── slice_010_xpdnet.png
    └── ...

Model Architecture

XPDNet (Default)

  • Framework: TensorFlow/Keras
  • Architecture: Primal-dual network with iterative reconstruction
  • Domain: Alternates between image and k-space domains
  • Iterations: 10 (configurable)
  • Pretrained: Brain multicoil, 4× acceleration

VarNet (Optional)

  • Framework: PyTorch via ATOMMIC
  • Architecture: Variational network with cascaded refinement
  • Domain: Primarily k-space with learned sensitivity encoding
  • Cascades: 12 (typical)
  • Pretrained: Available via Hugging Face

Checkpoint Management

Using Pretrained Weights

The system automatically downloads appropriate pretrained weights:

  • XPDNet: Downloads from facebook/xpdnet-brain-multicoil or fallback
  • VarNet: Downloads from facebook/varnet-brain-multicoil

Custom Checkpoints

# Load custom checkpoint
reconstructor.load_pretrained_weights('/path/to/custom/model.h5')

# Via command line
python reconstructFastMRI.py --h5 data.h5 --checkpoint custom_model.h5

Swapping Checkpoints

# Runtime checkpoint swapping
reconstructor = FastMRIReconstructor(method='xpdnet')

# Load first checkpoint
reconstructor.load_pretrained_weights('checkpoint_1.h5')
results_1 = reconstructor.reconstruct_volume('data.h5', 'output_1')

# Swap to second checkpoint
reconstructor.load_pretrained_weights('checkpoint_2.h5')
results_2 = reconstructor.reconstruct_volume('data.h5', 'output_2')

Technical Details

Data Processing Pipeline

  1. Load: Read multicoil k-space from h5 file
  2. Undersample: Apply acceleration mask
  3. Sensitivity: Estimate coil sensitivity maps
  4. Reconstruct: Run through neural network
  5. Save: Export as NIfTI, NumPy, and visualizations

Sensitivity Map Estimation

Uses ESPIRiT-like method:

  • Extract calibration region (24×24 default)
  • Convert to image domain
  • Normalize by root-sum-of-squares

Undersampling Patterns

  • Equispaced: Regular sampling with center calibration
  • Center fraction: 8% default (0.08)
  • Acceleration: 4× default (configurable)

System Requirements

Minimum Requirements

  • Python 3.7+
  • 8GB RAM
  • 2GB disk space

Recommended

  • Python 3.8+
  • 16GB RAM
  • GPU with CUDA support (for VarNet)
  • 5GB disk space

Dependencies

  • numpy>=1.21.0
  • tensorflow>=2.8.0
  • matplotlib>=3.5.0
  • h5py>=3.7.0
  • nibabel>=3.2.0
  • huggingface_hub>=0.16.0

Optional for VarNet:

  • torch>=1.12.0
  • atommic (install from GitHub)

Troubleshooting

Common Issues

1. Missing Dependencies

pip install -r requirements.txt

2. CUDA Out of Memory (VarNet)

# Use CPU instead
reconstructor = FastMRIReconstructor(method='xpdnet')  # Falls back to CPU

3. Download Errors

  • Check internet connection
  • Ensure Hugging Face Hub access
  • Try with VPN if behind firewall

4. Data Format Issues

  • Ensure h5 file has kspace dataset
  • Check data dimensions: [slices, coils, ky, kx]
  • Verify complex data format

Performance Tips

For Large Datasets:

# Process slice by slice to save memory
# (automatically handled by the implementation)

For Faster Processing:

# Use GPU if available (VarNet)
export CUDA_VISIBLE_DEVICES=0
python reconstructFastMRI.py --h5 data.h5 --method varnet

Contributing

  1. Fork the repository
  2. Create feature branch
  3. Add tests for new functionality
  4. Submit pull request

License

This project is provided for research and educational purposes. Please cite appropriate papers when using in publications:

  • XPDNet: [Original Paper]
  • VarNet: [Original Paper]
  • FastMRI: [FastMRI Dataset Paper]

Citation

If you use this implementation, please cite:

@software{fastmri_reconstruction,
  title={FastMRI Reconstruction with Pretrained Models},
  author={Your Name},
  year={2024},
  url={https://github.com/your-repo}
}

References

  1. FastMRI Dataset: https://fastmri.med.nyu.edu/
  2. XPDNet: [Link to paper]
  3. VarNet: [Link to paper]
  4. ATOMMIC: https://github.com/wdika/atommic

Need help? Open an issue or check the troubleshooting section above.

About

No description, website, or topics provided.

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

 
 
 

Contributors

Languages