colearn

A deep learning technique for the spatially varying fusion of spatially-aligned feature maps from multi-modality medical images.

Dependencies

tensorflow-gpu 1.10+
CUDA 9.0
CuDNN 7.0
python3
numpy
Pillow

This version has only been tested on Ubuntu 16.04 on an NVIDIA GTX 1080 Ti.

Usage

An example architecture is included here. A main function to use the architecture on PET-CT data is also included.

Data Preparation

The example main file expects PET-CT lung data in the TFRecord format.

    'height': tf.FixedLenFeature([], tf.int64)  # height of the image
        'width': tf.FixedLenFeature([], tf.int64)   # width of the image
        'depth': tf.FixedLenFeature([], tf.int64)   # deprecated will be removed
        'ct': tf.FixedLenFeature([], tf.string)     # CT image (height * width * 1)
        'pt': tf.FixedLenFeature([], tf.string)     # PET image (height * width * 1)
        'ctlb': tf.FixedLenFeature([], tf.string)   # binary masks of anatomical regions (height * width * 2)
        'ptlb': tf.FixedLenFeature([], tf.string)   # binary mask of tumours (height * w * 1)
        'bglb': tf.FixedLenFeature([], tf.string)   # binary mask of background (height * width * 1)

For the anatomical binary masks, the two channels (depth dimension of 2) are for the lung fields and the mediastinum.

Training

To train using default parameters on a GPU, use:

python3 colearn_cnn_petct_main.py \
    --mode=train \
    --train_data_path=TRAIN_DATA_PATH \
    [--valid_data_path=VALID_DATA_PATH] \
    --log_root=LOG_ROOT \
    --train_dir=LOG_ROOT/TRAIN_DIR \
    [--valid_dir=LOG_ROOT/VALID_DIR]

where:

TRAIN_DATA_PATH is the directory with the *.tfrecord files for the training data
VALID_DATA_PATH is the directory with the *.tfrecord files for the validation data
LOG_ROOT is the directory where the logs for this run will be stored
TRAIN_DIR is the subdirectory of LOG_ROOT for training logs
VALID_DIR is the subdirectory of LOG_ROOT for validation logs

If VALID_DATA_PATH is not specified, there will be no outputs for model validation. If VALID_DATA_PATH is specified, then VALID_DIR must also be specified.

Evaluation

To evaluate your trained model, use:

python3 colearn_cnn_petct_main.py \
    --mode=eval \
    --eval_data_path=EVAL_DATA_PATH \
    --log_root=LOG_ROOT \
    --eval_dir=LOG_ROOT/EVAL_DIR \
    --checkpoint_to_eval=CKPT_TO_EVAL \

where:

EVAL_DATA_PATH is the directory with the *.tfrecord files for the evaluation data
LOG_ROOT is the directory where the logs for this run will be stored
EVAL_DIR is the subdirectory of LOG_ROOT for evaluation logs
CKPT_TO_EVAL is the integer number of the training checkpoint to evaluate

Citation

If you make use of this code in your work, please cite the following paper:

A. Kumar, M. Fulham, D. Feng, and J. Kim, "Co-Learning Feature Fusion Maps from PET-CT Images of Lung Cancer", arXiv:1810.02492, 2018.

Link to Paper.

ashnilkumar / colearn

readme