Medical-Image-Classifier

Neural network classifier for medical image using Tensorflow

Content

1. Introduction
2. Features
3. Networks
4. Usage
   1. Required Libraries
   2. Software Configuration
   3. Example usage
      1. Folder Hierarchy
      2. Data preprocessing
      3. Training
      4. Evaluate
5. References
6. Authors

Introduction

This is a Tensorflow implementation for medical image classification tasks. The repository integrates a pre-processing pipeline optimized for 3D imaging data that suitable for multiple images input with multi-class output. The original application of this framework is to detect intracranial hemorrhage types with 3D CT images, however with slightly modification the framework is general applicable for all medical image classification tasks.

Features

• 3D data processing ready
• Augumented patching technique, requires less image input for training
• Multichannel input and multiclass output
• Generic image reader with SimpleITK support (Currently only support .nii/.nii.gz format for convenience, easy to expand to DICOM, tiff and jpg format)
• Medical image pre-post processing with SimpleITK filters
• Easy network replacement structure

Networks

Currently the classifier provides following classification networks:

• LeNet
• AlexNet
• GoogLeNetv1
• Vgg
• ResNet

Usage

Required Libraries

Known good dependencies:

• Anaconda 3.6
• Tensorflow 1.15
• SimpleITK
• tqdm
• pylatex (for PDF report output)

Software Configuration

All necessary network configurations are stored in config.json except for GPU specification. Modify the config.json file to fit your application.

Example usage

Folder Hierarchy

All training, testing and evaluation data should put in the directory that specified in config.json. Note that you should also provide the class labels in CSV format file.

In the default config.json, the data folder structure is as following:

.                                       # Repository root
├── ...
├── data                      
│   └── example                         # Data directory
│       ├── training                    # Put all training data here
│       │   ├── case1                   # foldername for the cases is arbitrary
│       │   |   └── image_brain_mni.nii # The image name is specified in config.json
│       │   ├── case2
│       │   |   └── image_brain_mni.nii
│       │   ├──	...
│       ├── testing                     # Put all testing data here
│       │   ├── case1                   # follow the same folder structure as the training one
│       │   |   └── image_brain_mni.nii
│       │   ├── case2
│       │   |   └── image_brain_mni.nii
│       │   ├── ...
│       ├── evaluation                  # Put all evaluation data here
│       │   ├── case1                   # follow the same folder structure as the training one
│       │   |   └── image_brain_mni.nii
│       │   ├── case2
│       │   |   └── image_brain_mni.nii
│       │   ├──	...
│       └── labels.csv                  # CSV file stores all labels for training and testing
├── tmp
│   ├── cktp                            # Tensorflow checkpoints
|   └── log                             # Tensorboard logging folder
├── ...

In training phase, the the program will automatically scan all the data in training and testing folder. The case/subject name is identified by subfolder's name and all of them should also be listed in labels.csv. Filenames could be altered in config.csv.

An example dataset folder is provided in ./data/example.

Data preprocessing

The framework provides 3D image preprocessing pipeline for end-to-end training. You may modify self.train_transforms with SimpleITK backed classes in NiftiDataset.py.

The example cases are used to classify intracranial hemorrhage types. To provide a normalized spacing and orientation for all data from different CT machines, the data are registered to MNI-spacing as one of the standard neurological coordinate system. Here we provide MNI templates in ./data/atlas for gantry tile correction and coordinate system normalization. Cerebral tissues are pre-extracted for a faster training time. The pre-processing procedures are completed by FSL.

GPU selection

By default the software use the 0th GPU in the system. Use argument option --gpu 1 to perform on GPU 1.

Training

1. Put all the data in accordance to Folder Hierarchy
2. Run the command:

   python main.py -p TRAIN --gpu 0

3. Open Tensorboard:

   tensorboard --logdir="./tmp/log"

Note:

• If you want to resume training from previous checkpoint, set "Restore": true in config.json.
• By default all the input images are in MNI space with size (91x109x91), voxel size = (2x2x2)mm^3.
• Change batch size to optimize the training according to your GPU memory. You may check GPU memory usage with nvidia-smi

Evaluate

Edit PredictionSetting in ./configs/config.json, then run with following command:

$ python main.py --phase PREDICT --config_json ./configs/config.json

Grid Search Optimal Hyperparameters

It is advisible to use Guild AI for machine learning model tracking. For detail check this file for further information.

References

• About the MNI space(s)
• FMRIB Software Library v6.0

Authors

Jacky Ko [email protected]