Modification to ClassifyDataLoader

README.md

@@ -1,9 +1,9 @@
 # RectAngle
 Segmentation and classification tool for trans-rectal B-mode ultrasound images.
 
-*Submitted as coursework for MPHY0041: Machine Learning in Medical Imaging.*
+*Submitted as part of the ASMUS2021 Conference for the paper 'Development and evaluation of intraoperative ultrasound segmentation with negative image frames and multiple observer labels'.*
 
-This package contains PyTorch-based implementations of a U-Net based segmentation model, and a DenseNet-based classification model, for the simultaneous detection and segmentation of prostate in rectal b-mode ultrasound images.
+This package contains a PyTorch-based implementation of a U-Net based segmentation model, and a DenseNet-based classification model, for the detection and segmentation of prostate in rectal b-mode ultrasound images.
 
 ## Installation
 
@@ -25,6 +25,8 @@ Once this is activated, the package may be installed using the *setup.py* file:
 
 Following this, training/inference may be performed using objects in the *train* module.
 
-To familiarise with the code used, an interactive notebook used for experiments in the associated report is available below. Please note that data used is proprietary and so has been withheld from the published repository.
+To familiarise yourself with the code used, an interactive notebook used for experiments in the associated report is available below. Please note that data used is proprietary and so has been withheld from the published repository.
 
 <a href="https://colab.research.google.com/github/liamchalcroft/RectAngle/blob/main/demo.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>
+
+(The code relevant to different label sampling methods is in sub-branch: label_method.)

prescreening_strategy2.py

@@ -0,0 +1,156 @@
+from numpy.lib.arraysetops import unique
+import rectangle as rect
+import h5py
+import torch
+import random
+import numpy as np
+import os 
+from rectangle.model.networks import DenseNet as DenseNet 
+
+from torch.utils.data import DataLoader 
+import tensorflow as tf 
+
+def standardise(image):
+
+    batch_ = image.shape[0]
+    for batch_iter_ in range(batch_):
+        image[batch_iter_,...] = (image[batch_iter_,...] - \
+                                torch.mean(image[batch_iter_,...]) / \
+                                torch.std(image[batch_iter_,...]))
+
+    return image
+
+def dice_score2(y_pred, y_true, eps=1e-8):
+    '''
+    y_pred, y_true -> [N, C=1, D, H, W]
+    '''
+    #y_pred[y_pred < 0.5] = 0.
+    #y_pred[y_pred > 0] = 1.
+
+    #Calculate the number of incorrectly labelled pixels 
+
+    numerator = torch.sum(y_true*y_pred, dim=(2,3)) * 2
+    denominator = torch.sum(y_true, dim=(2,3)) + torch.sum(y_pred, dim=(2,3)) + eps
+    return torch.mean(numerator / denominator)
+
+def dice_fp(y_pred, y_true, pos_frames, neg_frames): 
+    """ A function that computes dice score on positive frames, 
+    and FP pixels on negative frames, based off Yipeng's metrics
+    """
+    dice_ = dice_score2(y_pred[pos_frames, :, :], y_true[pos_frames, :, :])
+    fp = torch.sum(y_pred[neg_frames, :, :], dim = [1,2,3])
+
+    return dice_, fp 
+
+
+use_cuda = torch.cuda.is_available()
+
+### Loading ensemble segmentation network ### 
+num_ensemble = 5 
+path_str = '/Users/iani/Documents/Segmentation_project/ensemble/'
+latest_model = ['13.pth', '4.pth', '30.pth', '28.pth', '28.pth'] #Checked manually 
+model_paths = [os.path.join(path_str, 'model_'+ str(idx), latest_model[idx]) 
+for idx in range(num_ensemble)]
+
+depth = 5
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+seg_models = [rect.model.networks.UNet(n_layers=depth, device=device,
+                                    gate=None) for e in range(int(num_ensemble))]
+
+for n, m in enumerate(seg_models):
+    m.load_state_dict(torch.load(model_paths[n], map_location= device))
+
+### Loading classifier network ### 
+class_model = torch.load("/Users/iani/Documents/Segmentation_project/classification_model", map_location = device)
+
+### Inference ###
+
+test_file = h5py.File('/Users/iani/Documents/Reg2Seg/dataset/test.h5', 'r')
+test_DS = rect.utils.io.H5DataLoader(test_file)
+test_DL = DataLoader(test_DS, batch_size = 8, shuffle = False)
+
+segmentation_threshold = 0.5 
+classification_threshold = 0.5 
+
+
+all_dice_screen = []
+all_dice_noscreen = []
+
+all_fp_screen = []
+all_fp_noscreen = [] 
+
+with torch.no_grad(): 
+
+    for jj, (images_test, labels_test) in enumerate(test_DL):
+
+        if use_cuda: 
+            images_test, labels_test = images_test.cuda(), labels_test.cuda()
+
+        #Obtain positive and negative frames 
+        positive_frames = [(1 in label) for label in labels_test]
+        negative_frames = [not(1 in label) for label in labels_test]
+
+        #False positives negative frames
+
+        #Dice score : positive frames 
+
+        #Obtain prediction for classifier 
+        class_preds = class_model(images_test)
+
+        #Normalise images for segmentation network 
+        norm_images_test = standardise(images_test)
+
+        #Obtain predictions for each ensemble model and combine them
+        combined_predictions = torch.zeros_like(labels_test, dtype = float)
+        majority = len(seg_models) - 1 
+
+        for model_ in seg_models:
+            #Obtain predictions
+            model_.eval()
+            seg_predictions = torch.tensor(model_(norm_images_test) > 0.5, dtype = float)
+            combined_predictions += seg_predictions
+
+        #All segmentation results - only on positive frames 
+        combined_predictions = (combined_predictions >= majority) #Majority vote 
+        dice_noscreen, fp_noscreen = dice_fp(combined_predictions, labels_test, positive_frames, negative_frames)
+        all_dice_noscreen.append(dice_noscreen)
+        all_fp_noscreen.append(fp_noscreen)
+
+
+        #dice_noscreen = dice_score(combined_predictions, labels_test)
+        #all_dice_noscreen.append(dice_noscreen)
+
+        #Pre-screened results only 
+        prostate_idx = np.where(class_preds == 1)[0]
+        #dice_screened = dice_score(combined_predictions[prostate_idx, :,:], labels_test[prostate_idx, :,:])
+
+        positive_frames_screened = [positive_frames[i] for i in prostate_idx]
+        negative_frames_screened = [negative_frames[i] for i in prostate_idx]
+
+        dice_screen, fp_screen = dice_fp(combined_predictions[prostate_idx, :,:], labels_test[prostate_idx, :,:], positive_frames_screened, negative_frames_screened)
+        all_dice_screen.append(dice_screen)
+        all_fp_screen.append(fp_screen)
+
+        print(f"Dice scores: Not-screened : {dice_noscreen} | Screened : {dice_screen}")
+        print(f"FP scores: Not-screened : {fp_noscreen} | Screened : {fp_screen}")
+
+
+#Obtaining plots of the histogram 
+
+#Obtain all unique FP scores for screen, no screen method 
+unique_fp_screen = [np.unique(fp_vals) for fp_vals in all_fp_screen if len(fp_vals) > 0]
+unique_fp_screen = np.concatenate(unique_fp_screen, axis = 0)
+
+#Obtain all unique FP scores for screen, no screen method 
+unique_fp_noscreen = [np.unique(fp_vals) for fp_vals in all_fp_noscreen if len(fp_vals) > 0]
+unique_fp_noscreen = np.concatenate(unique_fp_noscreen, axis = 0)
+
+from matplotlib import pyplot as plt
+plt.hist(unique_fp_noscreen, label = "noscreen")
+plt.hist(unique_fp_screen, label = "screen")
+plt.xlabel("Number of FP pixels per negative segmented frame")
+plt.legend() 
+plt.show()
+
+print('Chicken')
+

scripts/baseline_mean.qsub.sh

@@ -0,0 +1,23 @@
+#$ -S /bin/bash
+#$ -l tmem=32G
+#$ -l h_vmem=32G
+#$ -l h_rt=40:00:00
+
+#$ -l gpu=true
+#$ -N baseline
+
+#$ -cwd
+
+#module purge
+#module load default/python/3.8.5
+source /share/apps/source_files/python/python-3.8.5.source
+source rectenv/bin/activate
+
+python ./RectAngle/train.py --train ./miccai_us_data/train.h5 \
+--val ./miccai_us_data/val.h5 \
+--ensemble 5 \
+--lr_schedule exponential \
+--label mean \
+--odir ./baseline_data/mean \
+--epochs 50 \
+--seed 0

scripts/baseline_random.qsub.sh

@@ -0,0 +1,23 @@
+#$ -S /bin/bash
+#$ -l tmem=32G
+#$ -l h_vmem=32G
+#$ -l h_rt=40:00:00
+
+#$ -l gpu=true
+#$ -N baseline
+
+#$ -cwd
+
+#module purge
+#module load default/python/3.8.5
+source /share/apps/source_files/python/python-3.8.5.source
+source rectenv/bin/activate
+
+python ./RectAngle/train.py --train ./miccai_us_data/train.h5 \
+--val ./miccai_us_data/val.h5 \
+--ensemble 5 \
+--lr_schedule exponential \
+--label random \
+--odir ./baseline_data/random \
+--epochs 50 \
+--seed 0

scripts/baseline_vote.qsub.sh

@@ -0,0 +1,23 @@
+#$ -S /bin/bash
+#$ -l tmem=32G
+#$ -l h_vmem=32G
+#$ -l h_rt=40:00:00
+
+#$ -l gpu=true
+#$ -N baseline
+
+#$ -cwd
+
+#module purge
+#module load default/python/3.8.5
+source /share/apps/source_files/python/python-3.8.5.source
+source rectenv/bin/activate
+
+python ./RectAngle/train.py --train ./miccai_us_data/train.h5 \
+--val ./miccai_us_data/val.h5 \
+--ensemble 5 \
+--label vote \
+--lr_schedule exponential \
+--odir ./baseline_data/vote \
+--epochs 50 \
+--seed 0

src/rectangle/utils/io.py

@@ -3,7 +3,45 @@
 import numpy as np
 import random
 
+def plot_example(frame, labels, gt_method=None, savefig=None):
+    '''
+    frame = image array
+    labels = single label or list of labels
+    savefig= filepath to save, if None (default) use plt.show()
+    gt_method = String to describe ground truth method used 
+    '''    
+
+    colors=['lime', 'red', 'blue', 'orange'] # cycles through colors in this order
+    if gt_method is None:
+      gt_method = 'Vote'
+
+    legend_names = ['Label 1', 'Label 2', 'Label 3', gt_method]
+    # 0 = label 1 = lime
+    # 1 = label 2 = red
+    # 3 = label 3 = blue
+    # 4 = ground truth = orange (if included in list)
+
+    plt.figure(figsize=(12, 12))
+    plt.imshow(frame, cmap='gray')
+    if type(labels) == list:
+      for i, label in enumerate(labels):
+        plt.contour(label, colors=colors[i], linewidths=2)
+    else:
+      plt.contour(label, colors=colors[0], linewidths=2)
+    plt.tight_layout()
+    plt.axis('off')      
+
+    # make legend
+    patches = [ mpatches.Patch(color=colors[i], label=legend_names[i]) for i in range(len(labels) ) ]
+    # put those patched as legend-handles into the legend
+    plt.legend(handles=patches, bbox_to_anchor=(0.97, 0.97), loc=1, borderaxespad=0., fontsize=30)
 
+    if savefig is None:
+      plt.show()
+    else:
+      plt.savefig(savefig)
+
+
 def train_val_test(file, ratio=(0.6, 0.2, 0.2)):
   """ Generate list of keys for file based on index values
   Input arguments:
@@ -94,6 +132,7 @@ def __getitem__(self, index):
     image = torch.unsqueeze(torch.tensor(
         self.file['frame_%05d' % (subj_ix, 
                                        )][()].astype('float32')), dim=0)
+
     if self.label == 'random':                                
       label = torch.unsqueeze(torch.tensor(
           self.file['label_%05d_%02d' % (subj_ix, 
@@ -164,6 +203,55 @@ def __getitem__(self, index):
       label = torch.tensor([0.0])
     return(image, label)
 
+class ClassifyDataLoader_v2(torch.utils.data.Dataset):
+
+  def __init__(self, file, keys=None):
+    """ Dataloader for hdf5 files, with labels converted to classifier labels
+    Input arguments:
+      file : h5py File object
+             Loaded using h5py.File(path : string)
+      keys : list, default = None
+             Keys from h5py file to use. Useful for train-val-test split.
+             If None, keys generated from entire file.
+    """
+
+    super().__init__()
+
+    self.file = file
+    if not keys:
+      keys = list(file.keys())
+
+    self.split_keys = [key.split('_') for key in keys]
+    start_subj = int(self.split_keys[0][1])
+    last_subj = int(self.split_keys[-1][1])
+    self.num_subjects = (last_subj - start_subj)+ 1  #Add 1 to account for 0 idx python
+    self.subjects = [key[1] for key in self.split_keys if key[0] == 'frame']
+    #self.subjects = np.linspace(start_subj, last_subj, 
+    #                            self.num_subjects+1, dtype=int)
+
+  def __len__(self):
+        return self.num_subjects
+
+  def __getitem__(self, index):
+
+    subj_ix = self.subjects[index]
+    image_key = 'frame_' + subj_ix
+    image = torch.unsqueeze(torch.tensor(self.file[image_key][()].astype('float32')), dim=0)
+
+    label_batch = torch.cat([torch.unsqueeze(torch.tensor(
+        self.file[f'label_{subj_ix}_0{label_ix}' ]
+        [()].astype('float32')), dim=0) for label_ix in range(3)])
+
+    label_vote = torch.sum(label_batch, dim=(1,2))
+    sum_vote = torch.sum(label_vote != 0)
+
+    #print(sum_vote)
+    if sum_vote >= 2:
+      label = torch.tensor([1.0])
+    else:
+      label = torch.tensor([0.0])
+
+    return(image, label)
 
 class TestPlotLoader(torch.utils.data.Dataset):
   def __init__(self, file, keys=None, label='vote'):
@@ -231,7 +319,6 @@ def __getitem__(self, index):
       label = torch.unsqueeze(torch.mean(label_batch, dim=0), dim=0)
     return(image, label)
 
-
 class PreScreenLoader(torch.utils.data.Dataset):
   def __init__(self, model, file, keys=None, label='random', threshold=0.5):
     """ Dataloader for hdf5 files.