training script from a pretrained model

.gitignore

@@ -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

README.md

@@ -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
 

data.py

+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)

eval.py

@@ -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)

inference.py

@@ -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.')

train_tf.py

+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!")