Commit 767802bb authored by lindawangg's avatar lindawangg

training script from a pretrained model

parent 74dfa4c8
......@@ -12,10 +12,10 @@ train_3000.txt
test_split_v2.txt
gensynth/
train.py
data.py
export_to_meta.py
model.py
train_tf.py
archive/
requirements.txt
test_dups.py
test_COVIDx2.txt
train_COVIDx2.txt
......@@ -18,7 +18,7 @@ For a detailed description of the methodology behind COVID-Net and a full descri
Currently, the COVID-Net team is working on COVID-RiskNet, a deep neural network tailored for COVID-19 risk stratification. Stay tuned as we make it available soon.
If you would like to **contribute COVID-19 x-ray images**, please contact us at linda.wang513@gmail.com and a28wong@uwaterloo.ca or alex@darwinai.ca. Lets all work together to stop the spread of COVID-19!
If you would like to **contribute COVID-19 x-ray images**, please submit to https://figure1.typeform.com/to/lLrHwv. Lets all work together to stop the spread of COVID-19!
If you are a researcher or healthcare worker and you would like access to the **GSInquire tool to use to interpret COVID-Net results** on your data or existing data, please reach out to a28wong@uwaterloo.ca or alex@darwinai.ca
......@@ -74,7 +74,7 @@ Additional requirements to generate dataset:
The current COVIDx dataset is constructed by the following open source chest radiography datasets:
* https://github.com/ieee8023/covid-chestxray-dataset
* https://www.kaggle.com/c/rsna-pneumonia-detection-challenge
* https://www.kaggle.com/c/rsna-pneumonia-detection-challenge (which came from: https://nihcc.app.box.com/v/ChestXray-NIHCC)
We especially thank the Radiological Society of North America and others involved in the RSNA Pneumonia Detection Challenge, and Dr. Joseph Paul Cohen and the team at MILA involved in the COVID-19 image data collection project, for making data available to the global community.
......@@ -108,18 +108,18 @@ The network takes as input an image of shape (N, 224, 224, 3) and outputs the so
If using the TF checkpoints, here are some useful tensors:
* input tensor: `input_1:0`
* logit tensor: `dense_3/MatMul:0`
* output tensor: `dense_3/Softmax:0`
* label tensor: `dense_3_target:0`
* class weights tensor: `dense_3_sample_weights:0`
* loss tensor: `loss/mul:0`
### Steps for training
Releasing TF training script from pretrained model soon.
<!--1. To train from scratch, `python train.py`
2. To train from an existing hdf5 file, `python train.py --checkpoint output/example/cp-0.hdf5`
3. For more options and information, `python train.py --help`
4. If you have a GenSynth account, to convert hdf5 file to TF checkpoints,
`python export_to_meta.py --weightspath output/example --weightspath cp-0.hdf5`-->
TF training script from a pretrained model:
1. We provide you with the tensorflow evaluation script, [train_tf.py](train_tf.py)
2. Locate the tensorflow checkpoint files (location of pretrained model)
3. To train from a pretrained model, `python train_tf.py --weightspath models/COVIDNetv2 --metaname model.meta_train --ckptname model-2069`
4. For more options and information, `python train_tf.py --help`
### Steps for evaluation
......
import tensorflow as tf
from tensorflow import keras
import numpy as np
import os
import cv2
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from sklearn.utils import shuffle
class BalanceDataGenerator(keras.utils.Sequence):
'Generates data for Keras'
def __init__(self,
dataset,
is_training=True,
batch_size=8,
input_shape=(224,224),
n_classes=3,
num_channels=3,
mapping={'normal': 0, 'pneumonia': 1, 'COVID-19': 2},
shuffle=True,
augmentation=True,
datadir='data',
class_weights=[1., 1., 25.]
):
'Initialization'
self.datadir = datadir
self.dataset = dataset
self.is_training = is_training
self.batch_size = batch_size
self.N = len(self.dataset)
self.input_shape = input_shape
self.n_classes = n_classes
self.num_channels = num_channels
self.mapping = mapping
self.shuffle = True
self.n = 0
self.class_weights = class_weights
if augmentation:
self.augmentation = ImageDataGenerator(
featurewise_center=False,
featurewise_std_normalization=False,
rotation_range=10,
width_shift_range=0.1,
height_shift_range=0.1,
horizontal_flip=True,
brightness_range=(0.9, 1.1),
fill_mode='constant',
cval=0.,
)
datasets = {'normal': [], 'pneumonia': [], 'COVID-19': []}
for l in dataset:
datasets[l.split()[-1]].append(l)
self.datasets = [
datasets['normal'] + datasets['pneumonia'],
datasets['COVID-19'],
]
print(len(self.datasets[0]), len(self.datasets[1]))
self.on_epoch_end()
def __next__(self):
# Get one batch of data
batch_x, batch_y, weights = self.__getitem__(self.n)
# Batch index
self.n += 1
# If we have processed the entire dataset then
if self.n >= self.__len__():
self.on_epoch_end
self.n = 0
return batch_x, batch_y, weights
def __len__(self):
return int(np.ceil(len(self.datasets[0]) / float(self.batch_size)))
def on_epoch_end(self):
'Updates indexes after each epoch'
if self.shuffle == True:
for v in self.datasets:
np.random.shuffle(v)
def __getitem__(self, idx):
batch_x, batch_y = np.zeros((self.batch_size, *self.input_shape, self.num_channels)), np.zeros(self.batch_size)
batch_files = self.datasets[0][idx*self.batch_size : (idx+1)*self.batch_size]
batch_files[np.random.randint(self.batch_size)] = np.random.choice(self.datasets[1])
for i in range(self.batch_size):
sample = batch_files[i].split()
if self.is_training:
folder = 'train'
else:
folder = 'test'
x = cv2.imread(os.path.join(self.datadir, folder, sample[1]))
x = cv2.resize(x, self.input_shape)
if self.is_training and hasattr(self, 'augmentation'):
x = self.augmentation.random_transform(x)
x = x.astype('float32') / 255.0
y = self.mapping[sample[2]]
batch_x[i] = x
batch_y[i] = y
weights = np.take(self.class_weights, batch_y.astype('int64'))
return batch_x, keras.utils.to_categorical(batch_y, num_classes=self.n_classes), weights
class DataGenerator(keras.utils.Sequence):
'Generates data for Keras'
def __init__(self,
dataset,
is_training=True,
batch_size=8,
input_shape=(224,224),
n_classes=3,
num_channels=3,
mapping={'normal': 0, 'pneumonia': 1, 'COVID-19': 2},
shuffle=True):
'Initialization'
self.dataset = dataset
self.is_training = is_training
self.batch_size = batch_size
self.N = len(self.dataset)
self.input_shape = input_shape
self.n_classes = n_classes
self.num_channels = num_channels
self.mapping = mapping
self.shuffle = True
self.on_epoch_end()
def __len__(self):
return int(np.ceil(self.N / float(self.batch_size)))
def on_epoch_end(self):
self.dataset = shuffle(self.dataset, random_state=0)
def __getitem__(self, idx):
batch_x, batch_y = np.zeros((self.batch_size, *self.input_shape, self.num_channels)), np.zeros(self.batch_size)
for i in range(self.batch_size):
index = min((idx * self.batch_size) + i, self.N-1)
sample = self.dataset[index].split()
if self.is_training:
folder = 'train'
else:
folder = 'test'
x = cv2.imread(os.path.join('data', folder, sample[1]))
x = cv2.resize(x, self.input_shape)
x = x.astype('float32') / 255.0
#y = int(sample[1])
y = self.mapping[sample[2]]
batch_x[i] = x
batch_y[i] = y
return batch_x, keras.utils.to_categorical(batch_y, num_classes=self.n_classes)
......@@ -4,51 +4,56 @@ import tensorflow as tf
import os, argparse
import cv2
parser = argparse.ArgumentParser(description='COVID-Net Evaluation')
parser.add_argument('--weightspath', default='output', type=str, help='Path to output folder')
parser.add_argument('--metaname', default='model.meta', type=str, help='Name of ckpt meta file')
parser.add_argument('--ckptname', default='model', type=str, help='Name of model ckpts')
parser.add_argument('--testfile', default='test_COVIDx.txt', type=str, help='Name of testfile')
parser.add_argument('--testfolder', default='test', type=str, help='Folder where test data is located')
args = parser.parse_args()
mapping = {'normal': 0, 'pneumonia': 1, 'COVID-19': 2}
sess = tf.Session()
tf.get_default_graph()
saver = tf.train.import_meta_graph(os.path.join(args.weightspath, args.metaname))
saver.restore(sess, os.path.join(args.weightspath, args.ckptname))
graph = tf.get_default_graph()
image_tensor = graph.get_tensor_by_name("input_1:0")
pred_tensor = graph.get_tensor_by_name("dense_3/Softmax:0")
file = open(args.testfile, 'r')
testfile = file.readlines()
y_test = []
pred = []
for i in range(len(testfile)):
line = testfile[i].split()
x = cv2.imread(os.path.join('data', args.testfolder, line[1]))
x = cv2.resize(x, (224, 224))
x = x.astype('float32') / 255.0
y_test.append(mapping[line[2]])
pred.append(np.array(sess.run(pred_tensor, feed_dict={image_tensor: np.expand_dims(x, axis=0)})).argmax(axis=1))
y_test = np.array(y_test)
pred = np.array(pred)
matrix = confusion_matrix(y_test, pred)
matrix = matrix.astype('float')
#cm_norm = matrix / matrix.sum(axis=1)[:, np.newaxis]
print(matrix)
#class_acc = np.array(cm_norm.diagonal())
class_acc = [matrix[i,i]/np.sum(matrix[i,:]) if np.sum(matrix[i,:]) else 0 for i in range(len(matrix))]
print('Sens Normal: {0:.3f}, Pneumonia: {1:.3f}, COVID-19: {2:.3f}'.format(class_acc[0],
class_acc[1],
class_acc[2]))
ppvs = [matrix[i,i]/np.sum(matrix[:,i]) if np.sum(matrix[:,i]) else 0 for i in range(len(matrix))]
print('PPV Normal: {0:.3f}, Pneumonia {1:.3f}, COVID-19: {2:.3f}'.format(ppvs[0],
ppvs[1],
ppvs[2]))
def eval(sess, graph, testfile, testfolder):
image_tensor = graph.get_tensor_by_name("input_1:0")
pred_tensor = graph.get_tensor_by_name("dense_3/Softmax:0")
y_test = []
pred = []
for i in range(len(testfile)):
line = testfile[i].split()
x = cv2.imread(os.path.join('data', testfolder, line[1]))
x = cv2.resize(x, (224, 224))
x = x.astype('float32') / 255.0
y_test.append(mapping[line[2]])
pred.append(np.array(sess.run(pred_tensor, feed_dict={image_tensor: np.expand_dims(x, axis=0)})).argmax(axis=1))
y_test = np.array(y_test)
pred = np.array(pred)
matrix = confusion_matrix(y_test, pred)
matrix = matrix.astype('float')
#cm_norm = matrix / matrix.sum(axis=1)[:, np.newaxis]
print(matrix)
#class_acc = np.array(cm_norm.diagonal())
class_acc = [matrix[i,i]/np.sum(matrix[i,:]) if np.sum(matrix[i,:]) else 0 for i in range(len(matrix))]
print('Sens Normal: {0:.3f}, Pneumonia: {1:.3f}, COVID-19: {2:.3f}'.format(class_acc[0],
class_acc[1],
class_acc[2]))
ppvs = [matrix[i,i]/np.sum(matrix[:,i]) if np.sum(matrix[:,i]) else 0 for i in range(len(matrix))]
print('PPV Normal: {0:.3f}, Pneumonia {1:.3f}, COVID-19: {2:.3f}'.format(ppvs[0],
ppvs[1],
ppvs[2]))
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='COVID-Net Evaluation')
parser.add_argument('--weightspath', default='output', type=str, help='Path to output folder')
parser.add_argument('--metaname', default='model.meta', type=str, help='Name of ckpt meta file')
parser.add_argument('--ckptname', default='model', type=str, help='Name of model ckpts')
parser.add_argument('--testfile', default='test_COVIDx.txt', type=str, help='Name of testfile')
parser.add_argument('--testfolder', default='test', type=str, help='Folder where test data is located')
args = parser.parse_args()
sess = tf.Session()
tf.get_default_graph()
saver = tf.train.import_meta_graph(os.path.join(args.weightspath, args.metaname))
saver.restore(sess, os.path.join(args.weightspath, args.ckptname))
graph = tf.get_default_graph()
file = open(args.testfile, 'r')
testfile = file.readlines()
eval(sess, graph, testfile, args.testfolder)
......@@ -30,5 +30,7 @@ x = x.astype('float32') / 255.0
pred = sess.run(pred_tensor, feed_dict={image_tensor: np.expand_dims(x, axis=0)})
print('Prediction: {}'.format(inv_mapping[pred.argmax(axis=1)[0]]))
print('Confidence')
print('Normal: {:.3f}, Pneumonia: {:.3f}, Normal: {:.3f}'.format(pred[0][0], pred[0][1], pred[0][2]))
print('**DISCLAIMER**')
print('Do not use this prediction for self-diagnosis. You should check with your local authorities for the latest advice on seeking medical assistance.')
from __future__ import print_function
import tensorflow as tf
import os, argparse, pathlib
from eval import eval
from data import BalanceDataGenerator
parser = argparse.ArgumentParser(description='COVID-Net Training Script')
parser.add_argument('--epochs', default=10, type=int, help='Number of epochs')
parser.add_argument('--lr', default=0.00002, type=float, help='Learning rate')
parser.add_argument('--bs', default=8, type=int, help='Batch size')
parser.add_argument('--weightspath', default='models/COVIDNetv2', type=str, help='Path to output folder')
parser.add_argument('--metaname', default='model.meta_train', type=str, help='Name of ckpt meta file')
parser.add_argument('--ckptname', default='model-2069', type=str, help='Name of model ckpts')
parser.add_argument('--trainfile', default='train_COVIDx.txt', type=str, help='Name of train file')
parser.add_argument('--testfile', default='test_COVIDx.txt', type=str, help='Name of test file')
parser.add_argument('--name', default='COVIDNet', type=str, help='Name of folder to store training checkpoints')
parser.add_argument('--datadir', default='data', type=str, help='Path to data folder')
args = parser.parse_args()
# Parameters
learning_rate = args.lr
batch_size = args.bs
display_step = 1
# output path
outputPath = './output/'
runID = args.name + '-lr' + str(learning_rate)
runPath = outputPath + runID
pathlib.Path(runPath).mkdir(parents=True, exist_ok=True)
print('Output: ' + runPath)
with open(args.trainfile) as f:
trainfiles = f.readlines()
with open(args.testfile) as f:
testfiles = f.readlines()
generator = BalanceDataGenerator(trainfiles, datadir=args.datadir, class_weights=[1., 1., 25.])
# Create a dataset tensor from the images and the labels
'''dataset = tf.data.Dataset.from_generator(lambda: generator,
output_types=(tf.float32, tf.float32, tf.float32),
output_shapes=([batch_size, 224, 224, 3],
[batch_size, 3],
[batch_size]))'''
# Create an iterator over the dataset
#iterator = dataset.make_initializable_iterator()
# Neural Net Input (images, labels, weights)
#batch_x, batch_y, weights = iterator.get_next()
with tf.Session() as sess:
tf.get_default_graph()
saver = tf.train.import_meta_graph(os.path.join(args.weightspath, args.metaname))
graph = tf.get_default_graph()
image_tensor = graph.get_tensor_by_name("input_1:0")
labels_tensor = graph.get_tensor_by_name("dense_3_target:0")
sample_weights = graph.get_tensor_by_name("dense_3_sample_weights:0")
pred_tensor = graph.get_tensor_by_name("dense_3/MatMul:0")
# loss expects unscaled logits since it performs a softmax on logits internally for efficiency
# Define loss and optimizer
loss_op = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits_v2(
logits=pred_tensor, labels=labels_tensor)*sample_weights)
optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)
train_op = optimizer.minimize(loss_op)
# Initialize the variables
init = tf.global_variables_initializer()
# Run the initializer
sess.run(init)
# load weights
saver.restore(sess, os.path.join(args.weightspath, args.ckptname))
#saver.restore(sess, tf.train.latest_checkpoint(args.weightspath))
# save base model
saver.save(sess, os.path.join(runPath, 'model'))
print('Saved baseline checkpoint')
print('Baseline eval:')
eval(sess, graph, testfiles, 'test')
# Training cycle
print('Training started')
total_batch = len(generator)
progbar = tf.keras.utils.Progbar(total_batch)
for epoch in range(args.epochs):
for i in range(total_batch):
# Run optimization
batch_x, batch_y, weights = next(generator)
sess.run(train_op, feed_dict={image_tensor: batch_x,
labels_tensor: batch_y,
sample_weights: weights})
progbar.update(i+1)
if epoch % display_step == 0:
pred = sess.run(pred_tensor, feed_dict={image_tensor:batch_x})
loss = sess.run(loss_op, feed_dict={pred_tensor: pred,
labels_tensor: batch_y,
sample_weights: weights})
print("Epoch:", '%04d' % (epoch + 1), "Minibatch loss=", "{:.9f}".format(loss))
eval(sess, graph, testfiles, 'test')
saver.save(sess, os.path.join(runPath, 'model'), global_step=epoch+1, write_meta_graph=False)
print('Saving checkpoint at epoch {}'.format(epoch + 1))
print("Optimization Finished!")
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment