Pulmonary Vascular Segmentation

A research project for automated pulmonary artery and vein segmentation from CT scans using deep learning-based super-resolution reconstruction and hierarchical segmentation.

Project Status

🔄 Work in Progress - This is an ongoing research project

• ✅ I2SR Super-Resolution: Implemented and functional
• ✅ Data preprocessing and HDF5 conversion pipeline
• ✅ Vessel enhancement filtering (Jerman filter)
• ✅ Analysis and visualization tools
• 🔄 HiPaS Segmentation: Basic implementation available
• ❌ STS (Segmentation Training System): Not yet implemented

Overview

This project implements a multi-stage approach for pulmonary vessel segmentation:

1. I2SR: 3D super-resolution reconstruction (5x slice interpolation)
2. HiPaS: Hierarchical vessel segmentation network
3. STS: Segmentation training system (planned)

Current Features

• 3D CT super-resolution with 5x slice enhancement
• NPZ to HDF5 data format conversion
• Jerman vessel enhancement filtering
• Comprehensive data visualization tools
• Model training and evaluation framework

Installation

# Create environment
conda create -n pulmonary-seg python=3.10
conda activate pulmonary-seg

# Install PyTorch with CUDA
conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia

# Install dependencies
pip install numpy scipy h5py matplotlib tqdm Pillow
pip install scikit-learn scikit-image opencv-python

Project Structure

├── data/                    # CT scans and annotations (.npz format)
├── HiPaS/                   # Segmentation models (basic implementation)
├── Reconstruction/          # I2SR super-resolution (functional)
│   ├── train.py            # I2SR training
│   ├── reconstruct.py      # Batch reconstruction to HDF5
│   ├── view_npz.py         # Data visualization
│   └── test_*.py           # Model evaluation
├── analysis/                # Analysis tools
├── filter.py               # Jerman vessel enhancement
└── requirements.txt

Quick Start

1. Data Exploration

cd Reconstruction
python view_npz.py --scan        # Dataset overview
python view_npz.py --sample      # View sample data
python view_npz.py --file ../data/train/001.npz  # Specific file

2. I2SR Training (Functional)

cd Reconstruction
python train.py --batchSize 2 --lr 1e-4 --nEpochs 200
python train.py --resume "../models/reconstruction/latest_checkpoint.pth"

3. Data Reconstruction

# Single file reconstruction
python reconstruct.py --input ../data/train/001.npz --output output.h5

# Batch processing with vessel enhancement
python reconstruct.py --batch --input_dir ../data/train/ --output_dir ../data/i2sr_hdf5/

4. Model Evaluation

python test_model.py          # Basic metrics
python visual_test.py         # Visual comparison

Current Implementation

I2SR Super-Resolution (Completed)

• Architecture: RRDB + FRB blocks, 577K parameters
• Input: 5 consecutive CT slices [1,1,5,512,512]
• Output: 5x interpolated slices with enhanced resolution
• Training: Functional with L1+MSE+Gradient loss
• Performance: ~15.5 dB PSNR on test set

Data Pipeline (Completed)

• Input Format: NPZ files [H,W,D], normalized [0,1]
• Output Format: HDF5 [2,D,H,W] with vessel enhancement
• Tools: Batch processing, visualization, format conversion

HiPaS Segmentation (Basic Implementation)

• Architecture: 3D U-Net with ResBlocks
• Status: Code structure available, needs refinement
• Multi-stage: Progressive vessel refinement planned

STS System (Planned)

• Segmentation training framework
• Advanced loss functions for vessel segmentation
• Integration with I2SR outputs

Research Notes

Known Issues

• STS segmentation system not yet implemented
• HiPaS network needs optimization for current data format
• Multi-GPU training support needs testing

TODO

• Complete STS implementation
• Optimize HiPaS for HDF5 input format
• Add comprehensive evaluation metrics
• Implement data augmentation strategies
• Add model export functionality

Hardware Requirements

• GPU: 8GB+ VRAM (RTX 3080 recommended)
• RAM: 16GB+
• Storage: 50GB+ for datasets and models

Contributing

This is a research project. Contributions welcome for:

• STS system implementation
• Performance optimization
• Documentation improvements
• Bug fixes and testing

License

MIT License - Research and educational use.