VERBOSE_TF = False
import os
if not VERBOSE_TF:
import warnings
warnings.filterwarnings('ignore', category=FutureWarning)
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
import tensorflow as tf
tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
else:
import tensorflow as tf
from random import shuffle
import numpy as np
import argparse
from time import time, strftime
from tqdm import tqdm
from os.path import join, dirname, realpath
from collections import defaultdict
from sklearn.metrics import roc_auc_score
from datetime import date
from pathlib import Path
from lungs.preprocess import preprocess, walk_dicom_dirs, walk_np_files
from lungs import utils
from lungs.i3d import InceptionI3d
class I3dForCTVolumes:
def __init__(self, args):
self.args = args
# This is the shape of both dimensions of each slice of the volume.
# The final volume shape fed to the model is [self.args['num_slices, 224, 224]
self.slice_size = 224
# pylint: disable=not-context-manager
with tf.Graph().as_default():
global_step = tf.get_variable(
'global_step',
[],
initializer=tf.constant_initializer(0),
trainable=False
)
# Placeholders
self.volumes_placeholder, self.labels_placeholder, self.is_training_placeholder = utils.placeholder_inputs(
num_slices=self.args['num_slices'],
crop_size=self.slice_size,
rgb_channels=3
)
# Learning rate and optimizer
lr = tf.train.exponential_decay(self.args['lr'], global_step, decay_steps=5000, decay_rate=0.1, staircase=True)
optimizer = tf.train.AdamOptimizer(lr)
# Init I3D model
with tf.device('/device:' + self.args['device'] + ':0'):
with tf.compat.v1.variable_scope('RGB'):
_, end_points = InceptionI3d(num_classes=2, final_endpoint='Predictions')\
(self.volumes_placeholder, self.is_training_placeholder, dropout_keep_prob=args['keep_prob'])
self.logits = end_points['Logits']
self.preds = end_points['Predictions']
# Loss function
# self.loss = utils.focal_loss(self.logits[:, 1], self.labels_placeholder)
self.loss = utils.cross_entropy_loss(self.logits, self.labels_placeholder)
# Evaluation metrics
self.get_preds = utils.get_preds(self.preds)
self.get_logits = utils.get_logits(self.logits)
self.accuracy = utils.accuracy(self.logits, self.labels_placeholder)
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
grads = optimizer.compute_gradients(self.loss)
apply_gradient = optimizer.apply_gradients(grads, global_step=global_step)
self.train_op = tf.group(apply_gradient)
# Create a saver for loading pretrained checkpoints.
pretrained_variable_map = {}
for variable in tf.global_variables():
if variable.name.split('/')[0] == 'RGB' and 'Adam' not in variable.name.split('/')[-1] \
and variable.name.split('/')[2] != 'Logits':
pretrained_variable_map[variable.name.replace(':0', '')] = variable
self.pretrained_saver = tf.train.Saver(var_list=pretrained_variable_map, reshape=True)
# Create a saver for writing training checkpoints.
self.saver = tf.train.Saver()
# Init local and global vars
init = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer())
# Create a session for running Ops on the Graph.
run_config = tf.ConfigProto(allow_soft_placement=True)
self.sess = tf.Session(config=run_config)
self.sess.run(init)
def train_loop(self, train_list, metrics_dir):
train_batches = utils.batcher(train_list, self.args['batch_size'])
for coupled_batch in tqdm(train_batches):
feed_dict, _ = self.process_data_into_to_dict(coupled_batch, is_training=True)
self.sess.run(self.train_op, feed_dict=feed_dict)
metrics = self.evaluate(train_list, ds='Train')
utils.write_number_list(metrics[-1], join(metrics_dir, 'tr_true'), verbose=self.args['verbose'])
return metrics
def evaluate(self, coupled_list, ds='Val.'):
coupled_batches = utils.batcher(coupled_list, self.args['batch_size'])
loss_list, acc_list, preds_list, labels_list = [], [], [], []
print('\nINFO: ++++++++++++++++++++ {} Evaluation ++++++++++++++++++++'.format(ds))
for coupled_batch in tqdm(coupled_batches):
feed_dict, labels = self.process_data_into_to_dict(coupled_batch)
acc, loss, preds = self.sess.run([self.accuracy, self.loss, self.get_preds], feed_dict=feed_dict)
loss_list.append(loss)
acc_list.append(acc)
preds_list.extend(preds)
labels_list.extend(labels)
if self.args['verbose']:
print('\nDEBUG: {}. Preds/Labels: {}'.format(ds, list(zip(preds_list, labels_list))))
print('\nDEBUG: {} Batch accuracy/loss: {}'.format(ds, list(zip(acc_list, loss_list))))
mean_acc = np.mean(acc_list)
mean_loss = np.mean(loss_list)
auc_score = roc_auc_score(labels_list, preds_list)
print('\n' + '=' * 34)
print("|| INFO: {} Accuracy: {:.4f} ||".format(ds, mean_acc))
print("|| INFO: {} Loss: {:.4f} ||".format(ds, mean_loss))
print("|| INFO: {} AUC: {:.4f} ||".format(ds, auc_score))
print('=' * 34)
return mean_loss, mean_acc, auc_score, preds_list, labels_list
def predict(self, inference_data):
errors_map = defaultdict(int)
volume_iterator = walk_np_files(inference_data) if self.args['preprocessed'] else walk_dicom_dirs(inference_data)
for i, volume_path in enumerate(volume_iterator):
try:
if not self.args['preprocessed']:
print('\nINFO: Preprocessing volume...')
preprocessed, _ = preprocess(volume_path, errors_map, self.args['num_slices'], self.slice_size, \
sample_volume=False, verbose=self.args['verbose'])
else:
preprocessed = self.load_np_volume(volume_path)
# preprocessed = np.expand_dims(preprocessed, axis=0)
except ValueError as e:
raise e
print('\nINFO: Predicting cancer for volume no. {}...'.format(i + 1))
singleton_batch = [[preprocessed, None]]
feed_dict, _ = self.process_data_into_to_dict(singleton_batch, from_paths=False)
preds = self.sess.run([self.get_preds], feed_dict=feed_dict)
print('\nINFO: Probability of cancer within 1 year: {:.5f}\n\n'.format(preds[0][0]))
def process_data_into_to_dict(self, coupled_batch, from_paths=True, is_training=False):
volumes = []
labels = []
for volume, label in coupled_batch:
try:
if from_paths:
volume = self.load_np_volume(volume)
# Crop volume to shape (self.args['num_slices'], 224, 224)
crop_start = volume.shape[0] // 2 - self.args['num_slices'] // 2
volume = volume[crop_start: crop_start + self.args['num_slices']]
volumes.append(volume)
if label is not None:
labels.append(label)
except:
print('\nERROR! Could not load:', volume)
# Perform windowing online volume, to save storage space of preprocessed volumes
volumes = np.array(volumes)
volume_batch = utils.apply_window(volumes)
if labels:
labels_np = np.array(labels).astype(np.int64)
else:
labels_np = np.zeros(volume_batch.shape[0], dtype=np.int64)
feed_dict = {self.volumes_placeholder: volume_batch, self.labels_placeholder: labels_np, self.is_training_placeholder: is_training}
return feed_dict, labels
def load_np_volume(self, volume_file):
if volume_file.endswith('.npz'):
scan_arr = np.load(join(self.args['data_dir'], volume_file))['data']
else:
scan_arr = np.load(join(self.args['data_dir'], volume_file)).astype(np.float32)
return scan_arr
def create_output_dirs(args):
# Create model dir and log dir if they doesn't exist
timestamp = date.today().strftime("%A_") + strftime("%H:%M:%S")
out_dir_time = Path(str(args['out_dir']) + '_' + timestamp)
save_dir = out_dir_time / 'models'
metrics_dir = out_dir_time / 'metrics'
val_preds_dir = metrics_dir / 'val_preds'
tr_preds_dir = metrics_dir / 'tr_preds'
plots_dir = out_dir_time / 'plots'
for new_dir in out_dir_time, save_dir, val_preds_dir, tr_preds_dir, plots_dir:
os.makedirs(new_dir, exist_ok=True)
return save_dir, metrics_dir, plots_dir
def main(args):
print('\nINFO: Initializing...')
# Set GPU
if args['device'] == 'GPU':
os.environ["CUDA_VISIBLE_DEVICES"] = str(args['gpu_id'])
# Init model wrapper
model = I3dForCTVolumes(args)
# Load pre-trained weights
pre_trained_ckpt = utils.load_pretrained_ckpt(args['ckpt'], args['data_dir'])
model.pretrained_saver.restore(model.sess, pre_trained_ckpt)
if args['input']:
print('\nINFO: Begin Inference \n')
model.predict(args['input'])
else:
print('\nINFO: Begin Training')
print('\nINFO: Hyperparams:')
print('\n'.join([str(item) for item in args.items()]))
save_dir, metrics_dir, plots_dir = create_output_dirs(args)
train_list = utils.load_data_list(args['train'])
val_list = utils.load_data_list(args['val'])
val_labels = utils.get_list_labels(val_list)
utils.write_number_list(val_labels, join(metrics_dir, 'val_true'), verbose=args['verbose'])
metrics = defaultdict(list)
for epoch in range(1, args['epochs'] + 1):
print('\nINFO: +++++++++++++++++++++ EPOCH {} +++++++++++++++++++++'.format(epoch))
start_time = time()
shuffle(train_list)
# Run training for 1 epoch and save weights to file
tr_epoch_metrics = model.train_loop(train_list, metrics_dir)
print("\nINFO: Saving Weights...")
model.saver.save(model.sess, "{}/epoch_{}/model.ckpt".format(save_dir, epoch))
train_end_time = time()
print('\nINFO: Train epoch duration: {:.2f} secs'.format(train_end_time - start_time))
# Run validation at end of each epoch
print("\nINFO: Begin Validation")
val_metrics = model.evaluate(val_list)
print('\nINFO: Val duration: {:.2f} secs'.format(time() - train_end_time))
print('\nINFO: Writing metrics plotting them...')
utils.write_metrics(metrics, tr_epoch_metrics, val_metrics, metrics_dir, epoch, verbose=args['verbose'])
utils.plot_metrics(epoch, metrics_dir, plots_dir)
def train(**kwargs):
'''
Run prediction.
For arguments description, see General and Training sections in params() function below.
'''
final_kwargs = params()
# Override default parameters with given arguments
for key, value in kwargs.items():
final_kwargs[key] = value
main(final_kwargs)
def predict(**kwargs):
'''
Run prediction.
For arguments description, see General and Inference sections in params() function below.
'''
final_kwargs = params()
# Override default parameters with given arguments
for key, value in kwargs.items():
final_kwargs[key] = value
main(final_kwargs)
def params():
parser = argparse.ArgumentParser()
default_out_dir = Path.home() / 'Lung-Cancer-Risk-Prediction' / 'out'
default_data_dir = Path.home() / 'Lung-Cancer-Risk-Prediction' / 'data'
lists_dir = default_data_dir / 'lists'
######################################## General parameters #########################################
parser.add_argument('--ckpt', default='cancer_fine_tuned', type=str, help="pre-trained weights to load. \
Either 'i3d_imagenet', 'cancer_fine_tuned' or a path to a directory containing model.ckpt file")
parser.add_argument('--num_slices', default=220, type=int, \
help='number of slices (z dimension) from the volume to be used by the model')
parser.add_argument('--verbose', default=False, type=bool, help='whether to print detailed logs')
######################################## Training parameters ########################################
parser.add_argument('--epochs', default=40, type=int, help='the number of epochs')
parser.add_argument('--lr', default=0.0001, type=int, help='initial learning rate')
parser.add_argument('--keep_prob', default=0.8, type=int, help='dropout keep prob')
parser.add_argument('--batch_size', default=2, type=int, help='the batch size for training/validation')
parser.add_argument('--gpu_id', default=1, type=int, help='gpu id')
parser.add_argument('--device', default='GPU', type=str, help='the device to execute on')
parser.add_argument('--data_dir', default=default_data_dir, \
help='path to data directory (for raw/processed volumes, train/val lists, checkpoints etc.)')
parser.add_argument('--train', default=lists_dir / 'train.list', help='path to train data .list file')
parser.add_argument('--val', default=lists_dir / 'val.list', help='path to validation data .list file')
parser.add_argument('--out_dir', default=default_out_dir, help='path to output dir for models, metrics and plots')
######################################## Inference parameters ########################################
parser.add_argument('--input', default=None, type=str, help="path to directory of volumes for cancer prediction.")
parser.add_argument('--preprocessed', default=False, type=bool, help='whether data for inference is \
preprocessed (.npz files) or raw volumes (dirs of .dcm files)')
parser.set_defaults()
args, _ = parser.parse_known_args()
kwargs = vars(args)
return kwargs
if __name__ == "__main__":
kwargs = params()
main(kwargs)
Datasets
Models
