This repository contains the code I used for my final project in the class Machine Learning for Health Care at NYU. You will find three subfolders:
The data
folder contains:
- a script written to convert the nerve data to the COCO format (
nerve_coco_data.py
) - an
sbatch
file to run the script on NYU's HPC cluster (coco.sbatch
) - a script written to split the nerve data into train and test (
split_train_test.py
) - an
sbatch
file to run the script on NYU's HPC cluster (split.sbatch
)
Note that the sbatch
files will need to be modified to contain the path to your nerve dataset and environments.
The unet
folder contains:
train.py
: the python file containing the code to train the modeltrain.sbatch
: executestrain.py
on NYU's HPC clusterbatch_norm_train.py
: the python file containing the code to train the model with batchnormb-train.sbatch
: executesbatch_norm_train.py
on NYU's HPC clusterdata.py
: loads the data files into.npy
format for faster loadingdata.sbatch
: executesdata.py
on NYU's HPC clusterget_masks.py
: loads in a trained model, computes the predicted masks on the test set, and saves the resultsinspect_model.ipynb
allows you to interactively look at the results of your trained modeljupyter.sbatch
allows you to run jupyter notebooks on NYU's HPC cluster (requires extra config)
If you'd like to run training, you must download the dataset from the Kaggle competition, run the split_train_test.py
script, and either symlink or move train
and val
to /path/to/repo/unet/raw
.
This folder contains the implementation of the Mask R-CNN. This implementation was adapted from Facebook's repository written in Pytorch 1.0. I had to modify various internal files in order to get my code to work, so I copied their repository in here. Follow INSTALL.md
to understand what needs to be downloaded in order to run training/evaluation.
This folder contains three relevant batch files:
train.sbatch
: trains a Mask R-CNN model on the nerve datasetval.sbatch
: runs evaluation on a trained modeljupyter.sbatch
: allows you to run jupyter notebooks on NYU's HPC cluster (requires extra config)
The relevant jupyter notebook is found under demos
: Mask_R-CNN_demo.ipynb
If you'd like to run training, you must download the dataset from the Kaggle competition, run the nerve_coco_data.py
script, and symlink annotations
, train
, and val
to /path/to/repo/mrcnn/datasets/nerve
.