diff --git a/README.md b/README.md
index b749bf7..5ab8fc2 100644
--- a/README.md
+++ b/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.)
diff --git a/prescreening_strategy2.py b/prescreening_strategy2.py
new file mode 100644
index 0000000..99719bc
--- /dev/null
+++ b/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')
+
diff --git a/scripts/baseline_mean.qsub.sh b/scripts/baseline_mean.qsub.sh
new file mode 100644
index 0000000..318a674
--- /dev/null
+++ b/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
diff --git a/scripts/baseline_random.qsub.sh b/scripts/baseline_random.qsub.sh
new file mode 100644
index 0000000..9e8179f
--- /dev/null
+++ b/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
diff --git a/scripts/baseline_vote.qsub.sh b/scripts/baseline_vote.qsub.sh
new file mode 100644
index 0000000..6b11ca8
--- /dev/null
+++ b/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
diff --git a/src/rectangle/utils/io.py b/src/rectangle/utils/io.py
index 8f8dfe3..962dcb6 100644
--- a/src/rectangle/utils/io.py
+++ b/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.
diff --git a/src/rectangle/utils/metrics.py b/src/rectangle/utils/metrics.py
index 6734015..74e75ee 100644
--- a/src/rectangle/utils/metrics.py
+++ b/src/rectangle/utils/metrics.py
@@ -3,6 +3,26 @@
 
 # Loss function
 
+class WeightedBCE(nn.Module):
+  def __init__(self, weights=None):
+      super().__init__()
+      self.weights = weights
+      
+  def forward(self, inputs, targets):
+    inputs = inputs.view(-1).float()
+    targets = targets.view(-1).float()
+
+    if self.weights is not None:
+        assert len(self.weights) == 2
+
+        loss = weights[1] * (targets * torch.log(inputs)) + \
+        weights[0] * ((1 - targets) * torch.log(1 - inputs))
+    else:
+        loss = targets * torch.log(inputs) + (1 - targets) * torch.log(1 - inputs)
+
+    return loss
+
+
 class DiceLoss(nn.Module):
   """ Loss function based on Dice-Sorensen Coefficient (L = 1 - Dice)
   Input arguments:
@@ -29,7 +49,7 @@ def forward(self, inputs, targets):
     # Seems to perform very well without binary - soft dice?
 
     if not self.soft:
-      inputs = BinaryDice(inputs, self.threshold)
+      inputs = self.BinaryDice(inputs, self.threshold)
 
     inputs = inputs.view(-1).float()
     targets = targets.view(-1).float()
diff --git a/src/rectangle/utils/train.py b/src/rectangle/utils/train.py
index a6802b8..8293f06 100644
--- a/src/rectangle/utils/train.py
+++ b/src/rectangle/utils/train.py
@@ -1,6 +1,7 @@
 import torch
 from torch import nn
 from rectangle.utils.metrics import DiceLoss, Precision, Recall, Accuracy
+from rectangle.utils.transforms import Binary
 from torch.optim import Adam
 from copy import deepcopy
 from os import path, makedirs
@@ -9,13 +10,15 @@
 from torch.utils.data import DataLoader, random_split, ConcatDataset
 import numpy as np
 from scipy.ndimage import laplace
+from torch.utils.tensorboard import SummaryWriter
+from torchvision.utils import make_grid
 
 
 class Trainer(nn.Module):
     def __init__(self, model, nb_epochs=200, outdir='./logs',
      loss=DiceLoss(), metric=DiceLoss(), opt='adam',
-     print_interval=1, val_interval=5, device='cuda',
-     early_stop=5, ensemble=None):
+     print_interval=1, val_interval=1, device='cuda',
+     early_stop=10, lr_schedule=None, ensemble=None): 
 
         super().__init__()
 
@@ -29,6 +32,7 @@ def __init__(self, model, nb_epochs=200, outdir='./logs',
         self.ensemble = ensemble
         self.outdir = outdir
         self.device = device
+        self.bin = Binary()
 
         if self.ensemble == 0:
             self.ensemble = None
@@ -44,15 +48,26 @@ def __init__(self, model, nb_epochs=200, outdir='./logs',
                 for i in range(self.ensemble):
                     self.model_ensemble.append(deepcopy(model))
 
+            self.writer = [SummaryWriter(log_dir=path.join(outdir,'runs/model_{}'.format(i))) for i in range(self.ensemble)]
+        else:
+            self.writer = SummaryWriter(log_dir=path.join(outdir,'runs'))
+
         if opt == 'adam':
             if self.ensemble: 
-                opt = [Adam(model.parameters()) for model in self.model_ensemble]
+                opt = [Adam(model.parameters(), lr=0.0001) for model in self.model_ensemble]
             else:
-                opt = Adam(model.parameters())
+                opt = Adam(model.parameters(), lr=0.0001)
             # opt = Adam(model.parameters())
 
         self.opt = opt
 
+        if lr_schedule:
+            if lr_schedule not in ['lambda', 'exponential', 'reduce_on_plateau']:
+                raise ValueError('Available learning rate schedules are LambdaLR, ExponentialLR or ReduceLROnPlateau.')
+            elif self.ensemble:
+                lr_schedule = [lr_schedule for model in self.model_ensemble]
+
+        self.lr_schedule = lr_schedule
 
     def train(self, train_data, val_data=None, oname=None, 
         train_pre=None, train_post=None, train_batch=128, train_shuffle=True,
@@ -93,6 +108,18 @@ def train(self, train_data, val_data=None, oname=None,
                 train = train_list[i]
                 val = val_list[i]
                 opt_ = self.opt[i]
+                writer_ = self.writer[i]
+                if self.lr_schedule:
+                    lr_schedule_ = self.lr_schedule[i]
+                    if lr_schedule_ == 'lambda':
+                        lr_schedule_ = torch.optim.lr_scheduler.LambdaLR(opt_, lambda epoch: 0.95 ** epoch)
+                    elif lr_schedule_ == 'exponential':
+                        lr_schedule_ = torch.optim.lr_scheduler.ExponentialLR(opt_, 0.95)
+                    else:
+                        lr_schedule_ = torch.optim.lr_scheduler.ReduceLROnPlateau(opt_)
+                else:
+                    lr_schedule_ = None
+
                 print('Beginning training of model #{}'.format(i))
                 for epoch in range(self.nb_epochs):
                     if self.early_stop:
@@ -105,7 +132,13 @@ def train(self, train_data, val_data=None, oname=None,
                         opt_.zero_grad()
                         if train_pre:
                             for aug in train_pre:
-                                input = aug(input)
+                                if aug.__class__.__name__ == 'Flip' or 'Affine':
+                                    input = torch.cat([input, label])
+                                    input = aug(input)
+                                    input, label = torch.chunk(input, 2)
+                                    label = self.bin(label)
+                                else:
+                                    input = aug(input)
                         pred = model(input)
                         if train_post:
                             for aug in train_post:
@@ -114,8 +147,12 @@ def train(self, train_data, val_data=None, oname=None,
                         loss_.backward()
                         opt_.step()
                         loss_epoch.append(loss_.item())
+                    if lr_schedule_ and lr_schedule_ != 'reduce_on_plateau':
+                        lr_schedule_.step()
                     loss_log_ensemble[i,epoch] = np.nanmean(loss_epoch)
                     if epoch % self.print_interval == 0:
+                        writer_.add_scalar('train/dice_loss_ensemble', np.nanmean(dice_epoch), epoch)
+                        writer_.add_scalar('train/dice_coefficient_ensemble', 1-np.nanmean(dice_epoch), epoch)
                         print('Epoch #{}: Mean Dice Loss: {}'.format(epoch, loss_log_ensemble[i,epoch]))
                     if epoch % self.val_interval == 0:
                         dice_epoch = []
@@ -125,7 +162,13 @@ def train(self, train_data, val_data=None, oname=None,
                                 input, label = input.to(self.device), label.to(self.device)
                                 if val_pre:
                                     for aug in val_pre:
-                                        input = aug(input)
+                                        if aug.__class__.__name__ == 'Flip' or 'Affine':
+                                            input = torch.cat([input, label])
+                                            input = aug(input)
+                                            input, label = torch.chunk(input, 2)
+                                            label = self.bin(label)
+                                        else:
+                                            input = aug(input)
                                 pred = model(input)
                                 if val_post:
                                     for aug in val_post:
@@ -133,6 +176,27 @@ def train(self, train_data, val_data=None, oname=None,
                                 dice_metric = self.metric(pred, label)
                                 dice_epoch.append(1 - dice_metric.item())
                             dice_log_ensemble[i,int(epoch//self.val_interval)] = np.nanmean(dice_epoch)
+                            if lr_schedule_ == 'reduce_on_plateau':
+                                lr_schedule_.step(1-np.nanmean(dice_epoch))
+
+                            writer_.add_scalar('val/dice_loss_ensemble', np.nanmean(dice_epoch), epoch)
+                            writer_.add_scalar('val/dice_coefficient_ensemble', 1-np.nanmean(dice_epoch), epoch)
+
+                            ## show some (e.g.,10) example images in tensorboard
+                            ex_num = 10
+                            ex_label = label[:ex_num,0]
+                            ex_pred = pred[:ex_num,0]
+                            ex_image = torch.cat([ex_label,ex_pred], dim=2)
+
+                            ex_images = ex_image.reshape(-1,ex_image.shape[2])
+                            image_grid = (make_grid(ex_images, nrow=ex_num)[0]+0.5)/ex_num                     
+                            writer_.add_images(
+                                "val/example_images_ensemble",
+                                image_grid,
+                                epoch,
+                                dataformats="HW",
+                            )
+
                         if epoch >= self.val_interval:
                             if dice_log_ensemble[i,int(epoch//self.val_interval)] > dice_max:
                                 early_ = 0
@@ -155,6 +219,14 @@ def train(self, train_data, val_data=None, oname=None,
             early_ = 0
             dice_max = 0
             model = self.model
+            if self.lr_schedule:
+                if self.lr_schedule == 'lambda':
+                    self.lr_schedule = torch.optim.lr_scheduler.LambdaLR(self.opt, lambda epoch: 0.95 ** epoch)
+                elif self.lr_schedule == 'exponential':
+                    self.lr_schedule = torch.optim.lr_scheduler.ExponentialLR(self.opt, 0.95)
+                else:
+                    self.lr_schedule = torch.optim.lr_scheduler.ReduceLROnPlateau(self.opt)
+
             for epoch in range(self.nb_epochs):
                 if self.early_stop:
                     if early_ == self.early_stop:
@@ -166,7 +238,13 @@ def train(self, train_data, val_data=None, oname=None,
                     self.opt.zero_grad()
                     if train_pre:
                         for aug in train_pre:
-                            input = aug(input)
+                            if aug.__class__.__name__ == 'Flip' or 'Affine':
+                                input = torch.cat([input, label])
+                                input = aug(input)
+                                input, label = torch.chunk(input, 2)
+                                label = self.bin(label)
+                            else:
+                                input = aug(input)
                     pred = model(input)
                     if train_post:
                         for aug in train_post:
@@ -175,8 +253,12 @@ def train(self, train_data, val_data=None, oname=None,
                     loss_.backward()
                     self.opt.step()
                     loss_epoch.append(loss_.item())
+                if self.lr_schedule and self.lr_schedule != 'reduce_on_plateau':
+                    self.lr_schedule.step()
                 loss_log[epoch] = np.nanmean(loss_epoch)
                 if epoch % self.print_interval == 0:
+                    self.writer.add_scalar('train/dice_loss', np.nanmean(loss_epoch), epoch)
+                    self.writer.add_scalar('train/dice_coefficient', 1-np.nanmean(loss_epoch), epoch)
                     print('Epoch #{}: Mean Dice Loss: {}'.format(epoch, loss_log[epoch]))
                 if epoch % self.val_interval == 0:
                     dice_epoch = []
@@ -186,14 +268,43 @@ def train(self, train_data, val_data=None, oname=None,
                             input, label = input.to(self.device), label.to(self.device)
                             if val_pre:
                                 for aug in val_pre:
-                                    input = aug(input)
+                                    if aug.__class__.__name__ == 'Flip' or 'Affine':
+                                        input = torch.cat([input, label])
+                                        input = aug(input)
+                                        input, label = torch.chunk(input, 2)
+                                        label = self.bin(label)
+                                    else:
+                                        input = aug(input)
                             pred = model(input)
                             if val_post:
                                 for aug in val_post:
                                     pred = aug(pred)
                             dice_metric = self.metric(pred, label)
+                            
                             dice_epoch.append(1 - dice_metric.item())
                         dice_log[int(epoch//self.val_interval)] = np.nanmean(dice_epoch)
+
+                        if self.lr_schedule == 'reduce_on_plateau':
+                            self.lr_schedule.step(1-np.nanmean(dice_epoch)) # monitors validation loss 
+                        self.writer.add_scalar('val/dice_loss', np.nanmean(dice_epoch), epoch)
+                        self.writer.add_scalar('val/dice_coefficient', 1-np.nanmean(dice_epoch), epoch)
+
+
+                        ## show some (e.g.,10) example images in tensorboard
+                        ex_num = 10
+                        ex_label = label[:ex_num,0]
+                        ex_pred = pred[:ex_num,0]
+                        ex_image = torch.cat([ex_label,ex_pred], dim=2)
+
+                        ex_images = ex_image.reshape(-1,ex_image.shape[2])
+                        image_grid = (make_grid(ex_images, nrow=ex_num)[0]+0.5)/ex_num                       
+                        self.writer.add_images(
+                            "val/example_images",
+                            image_grid,
+                            epoch,
+                            dataformats="HW",
+                        )
+
                     if epoch % self.print_interval == 0:
                         print('Mean Validation Dice: {}'.format(dice_log[int(epoch//self.val_interval)]))
                     if epoch >= self.val_interval:
@@ -266,19 +377,35 @@ def test(self, test_data, oname=None,
             oname = date.today()
             oname = oname.strftime("%b-%d-%Y")
 
+        path_ = path.join(self.outdir,\
+                            'testing/plots')
+        if not path.exists(path_):
+            makedirs(path_)
+
         test = DataLoader(test_data, 1, shuffle=False)
         dice_log = []
         prec_log = []
         rec_log = []
+        neg_log = []
+        pos_log = []
         precision = Precision()
         recall = Recall()
-        self.model.eval()
+        if self.ensemble:
+            self.model = [model.eval() for model in self.model]
+        else:
+            self.model.eval()
         with torch.no_grad():
             for i, (input, label) in enumerate(test):
                 input, label = input.to(self.device), label.to(self.device)
                 if test_pre:
                     for aug in test_pre:
-                        input = aug(input)
+                        if aug.__class__.__name__ == 'Flip' or 'Affine':
+                            input = torch.cat([input, label])
+                            input = aug(input)
+                            input, label = torch.chunk(input, 2)
+                            label = self.bin(label)
+                        else:
+                            input = aug(input)
                 if self.ensemble:
                     pred = [model(input) for model in self.model_ensemble]
                     pred = torch.cat(pred, dim=0)
@@ -288,82 +415,83 @@ def test(self, test_data, oname=None,
                 if test_post:
                     for aug in test_post:
                         pred = aug(pred)
-                dice_metric = self.metric(pred, label)
-                dice_log.append(1-dice_metric.item())
-                prec_log.append(precision(pred, label))
-                rec_log.append(recall(pred, label))
-
-                input_img = input.detach().cpu().numpy()
-                pred_img = pred.detach().cpu().numpy()
-                label_img = label.detach().cpu().numpy()
-
-                input_img = np.squeeze(input_img)
-                pred_img = np.squeeze(pred_img)
-                label_img = np.squeeze(label_img)
-
-                if overlap=='contour':
-                    input_img -= input_img.min()
-                    input_img *= 1.0/input_img.max()
-                    label_img = laplace(label_img)
-                    pred_img = laplace(pred_img)
-                    label_img = (label_img != 0)
-                    pred_img = (pred_img != 0)
-                    label_img = np.ma.masked_where(label_img == 0, label_img)
-                    pred_img = np.ma.masked_where(pred_img == 0, pred_img)
-                    plt.figure()
-                    plt.imshow(input_img, cmap='gray', vmin=0, vmax=1)
-                    plt.axis('off')
-                    plt.imshow(label_img, cmap='Greens', vmin=0, vmax=1)
-                    plt.axis('off')
-                    plt.imshow(pred_img, cmap='Reds', vmin=0, vmax=1)
-                    plt.axis('off')
-                elif overlap=='mask':
-                    input_img -= input_img.min()
-                    input_img *= 1.0/input_img.max()
-                    label_img = np.ma.masked_where(label_img == 0, label_img)
-                    pred_img = np.ma.masked_where(pred_img == 0, pred_img)
-                    plt.figure()
-                    plt.imshow(input_img, cmap='gray', vmin=0, vmax=1)
-                    plt.axis('off')
-                    plt.imshow(label_img, cmap='Greens', vmin=0, vmax=1)
-                    plt.axis('off')
-                    plt.imshow(pred_img, cmap='Reds', vmin=0, vmax=1)
-                    plt.axis('off')
+
+                if pred.sum() == 0:
+                    if label.sum() == 0:
+                        neg_log.append(1.)
+                    else:
+                        neg_log.append(0.)
                 else:
-                    plt.figure()
-                    plt.subplot(131)
-                    plt.imshow(input_img, cmap='gray')
-                    plt.axis('off')
-                    plt.title('Image')
-                    plt.subplot(132)
-                    plt.imshow(pred_img, cmap='gray', vmin=0, vmax=1)
-                    plt.axis('off')
-                    plt.title('Prediction (DSC={:.2f})'.format(dice_log[i]))
-                    plt.subplot(133)
-                    plt.imshow(label_img, cmap='gray', vmin=0, vmax=1)
-                    plt.axis('off')
-                    plt.title('Ground Truth')
-
-                path_ = path.join(self.outdir,\
-                                'testing/plots')
-                if not path.exists(path_):
-                    makedirs(path_)
-                plt.savefig(path.join(path_, 'pred{}_{}.png'.format(i, oname)))
+                    if label.sum() == 0:
+                        pos_log.append(0.)
+                    else:
+                        pos_log.append(1.)
+                    dice_metric = self.metric(pred, label)
+                    dice_log.append(1-dice_metric.detach().cpu().numpy())
+                    prec_log.append(precision(pred, label).detach().cpu().numpy())
+                    rec_log.append(recall(pred, label).detach().cpu().numpy())
+
+                    input_img = input.detach().cpu().numpy()
+                    pred_img = pred.detach().cpu().numpy()
+                    label_img = label.detach().cpu().numpy()
+
+                    input_img = np.squeeze(input_img)
+                    pred_img = np.squeeze(pred_img)
+                    label_img = np.squeeze(label_img)
+
+                    if i % 50==0:
+                        plt.figure()
+                        plt.imshow(pred_img, cmap='gray', vmin=0, vmax=1)
+                        plt.axis('off')
+                        plt.title('Prediction (DSC={:.2f})'.format(dice_log[-1]))
+                        plt.savefig(path.join(path_, 'pred_{}_{}.png'.format(i, oname)))
+
+                        if overlap:
+                            if overlap=='contour':
+                                input_img -= input_img.min()
+                                input_img *= 1.0/input_img.max()
+                                # label_img = np.ma.masked_where(label_img == 0, label_img)
+                                # pred_img = np.ma.masked_where(pred_img == 0, pred_img)
+                                plt.figure()
+                                plt.imshow(input_img, cmap='gray', vmin=0, vmax=1)
+                                plt.axis('off')
+                                plt.contour(label_img, cmap='Greens', linewidths=1)
+                                plt.axis('off')
+                                plt.contour(pred_img, cmap='Reds', linewidths=1)
+                                plt.title('Prediction (DSC={:.2f})'.format(dice_log[-1]))
+                                plt.axis('off')
+                            elif overlap=='mask':
+                                input_img -= input_img.min()
+                                input_img *= 1.0/input_img.max()
+                                label_img = np.ma.masked_where(label_img == 0, label_img)
+                                pred_img = np.ma.masked_where(pred_img == 0, pred_img)
+                                plt.figure()
+                                plt.imshow(input_img, cmap='gray', vmin=0, vmax=1)
+                                plt.axis('off')
+                                plt.imshow(label_img, cmap='Greens', vmin=0, vmax=1, alpha=0.3)
+                                plt.axis('off')
+                                plt.imshow(pred_img, cmap='Reds', vmin=0, vmax=1, alpha=0.3)
+                                plt.title('Prediction (DSC={:.2f})'.format(dice_log[-1]))
+                                plt.axis('off')
+                                
+                            plt.savefig(path.join(path_, 'pred_overlap_{}_{}.png'.format(i, oname)))
 
         dice_log = np.array(dice_log, dtype=float)
         prec_log = np.array(prec_log, dtype=float)
         rec_log = np.array(rec_log, dtype=float)
+        neg_log = np.array(neg_log, dtype=float)
+        pos_log = np.array(pos_log, dtype=float)
 
         plt.figure()
-        plt.scatter(rec_log, prec_log)
-        plt.plot([0,0.5,1], [0.5,0.5,0.5], '--')
+        plt.scatter(rec_log, prec_log, alpha=0.4)
+        # plt.plot([0,0.5,1], [0.5,0.5,0.5], '--')
         plt.xlabel('Recall')
         plt.ylabel('Precision')
         plt.title('AUC = {:.2f}'.format(np.sum(prec_log * rec_log)/np.size(prec_log)))
         plt.savefig(path.join(path_, 'prec_rec_{}'.format(oname)))
 
 
-        print('Mean Dice score: {:.2f}±{:.3f}, Mean Precision: {:.2f}±{:.3f}, Mean Recall: {:.2f}±{:.3f}'.format(np.mean(dice_log), np.std(dice_log), np.mean(prec_log), np.std(prec_log), np.mean(rec_log), np.std(rec_log)))
+        print('Mean Dice score: {:.2f}±{:.3f}, Mean Precision: {:.2f}±{:.3f}, Mean Recall: {:.2f}±{:.3f} \n TP Rate: {:.2f}±{:.3f}, TN Rate: {:.2f}±{:.3f}'.format(np.mean(dice_log), np.std(dice_log), np.mean(prec_log), np.std(prec_log), np.mean(rec_log), np.std(rec_log), np.mean(pos_log), np.std(pos_log), np.mean(neg_log), np.std(neg_log)))
         path_ = path.join(self.outdir,\
                                 'testing/table')
         if not path.exists(path_):
@@ -374,6 +502,10 @@ def test(self, test_data, oname=None,
                 prec_log, delimiter=',')
         np.savetxt(path.join(path_, 'recall_{}.csv'.format(oname)),\
                 rec_log, delimiter=',')
+        np.savetxt(path.join(path_, 'negative_{}.csv'.format(oname)),\
+                neg_log, delimiter=',')
+        np.savetxt(path.join(path_, 'positive_{}.csv'.format(oname)),\
+                pos_log, delimiter=',')
 
         print('Testing complete')
 
@@ -420,6 +552,8 @@ def __init__(self, model, nb_epochs=200, outdir='./logs',
 
         self.opt = opt
 
+        self.writer = SummaryWriter(log_dir=path.join(outdir,'runs'))
+
 
     def train(self, train_data, val_data=None, oname=None, 
         train_pre=None, train_post=None, train_batch=128, train_shuffle=True,
@@ -472,7 +606,13 @@ def train(self, train_data, val_data=None, oname=None,
                         opt_.zero_grad()
                         if train_pre:
                             for aug in train_pre:
-                                input = aug(input)
+                                if aug.__class__.__name__ == 'Flip' or 'Affine':
+                                    input = torch.cat([input, label])
+                                    input = aug(input)
+                                    input, label = torch.chunk(input, 2)
+                                    label = self.bin(label)
+                                else:
+                                    input = aug(input)
                         pred = model(input)
                         if train_post:
                             for aug in train_post:
@@ -483,6 +623,8 @@ def train(self, train_data, val_data=None, oname=None,
                         loss_epoch.append(loss_.item())
                     loss_log_ensemble[i,epoch] = np.nanmean(loss_epoch)
                     if epoch % self.print_interval == 0:
+                        self.writer.add_scalar('class_train/dice_loss_ensemble', np.nanmean(loss_epoch), epoch) 
+                        self.writer.add_scalar('class_train/dice_coefficient_ensemble', 1-np.nanmean(loss_epoch), epoch)
                         print('Epoch #{}: Mean acc Loss: {}'.format(epoch, loss_log_ensemble[i,epoch]))
                     if epoch % self.val_interval == 0:
                         acc_epoch = []
@@ -492,7 +634,13 @@ def train(self, train_data, val_data=None, oname=None,
                                 input, label = input.to(self.device), label.to(self.device)
                                 if val_pre:
                                     for aug in val_pre:
-                                        input = aug(input)
+                                        if aug.__class__.__name__ == 'Flip' or 'Affine':
+                                            input = torch.cat([input, label])
+                                            input = aug(input)
+                                            input, label = torch.chunk(input, 2)
+                                            label = self.bin(label)
+                                        else:
+                                            input = aug(input)
                                 pred = model(input)
                                 if val_post:
                                     for aug in val_post:
@@ -500,6 +648,8 @@ def train(self, train_data, val_data=None, oname=None,
                                 acc_metric = self.metric(pred, label)
                                 acc_epoch.append(acc_metric)
                             acc_log_ensemble[i,int(epoch//self.val_interval)] = np.nanmean(acc_epoch)
+                            self.writer.add_scalar('class_val/dice_loss_ensemble', np.nanmean(acc_epoch), epoch) 
+                            self.writer.add_scalar('class_val/dice_coefficient_ensemble', 1-np.nanmean(acc_epoch), epoch)
                         if epoch >= self.val_interval:
                             if acc_log_ensemble[i,int(epoch//self.val_interval)] > acc_max:
                                 early_ = 0
@@ -533,7 +683,13 @@ def train(self, train_data, val_data=None, oname=None,
                     self.opt.zero_grad()
                     if train_pre:
                         for aug in train_pre:
-                            input = aug(input)
+                            if aug.__class__.__name__ == 'Flip' or 'Affine':
+                                input = torch.cat([input, label])
+                                input = aug(input)
+                                input, label = torch.chunk(input, 2)
+                                label = self.bin(label)
+                            else:
+                                input = aug(input)
                     pred = model(input)
                     if train_post:
                         for aug in train_post:
@@ -544,6 +700,8 @@ def train(self, train_data, val_data=None, oname=None,
                     loss_epoch.append(loss_.item())
                 loss_log[epoch] = np.nanmean(loss_epoch)
                 if epoch % self.print_interval == 0:
+                    self.writer.add_scalar('class_train/dice_loss', np.nanmean(loss_epoch), epoch) 
+                    self.writer.add_scalar('class_train/dice_coefficient', 1-np.nanmean(loss_epoch), epoch)
                     print('Epoch #{}: Mean acc Loss: {}'.format(epoch, loss_log[epoch]))
                 if epoch % self.val_interval == 0:
                     acc_epoch = []
@@ -553,7 +711,13 @@ def train(self, train_data, val_data=None, oname=None,
                             input, label = input.to(self.device), label.to(self.device)
                             if val_pre:
                                 for aug in val_pre:
-                                    input = aug(input)
+                                    if aug.__class__.__name__ == 'Flip' or 'Affine':
+                                        input = torch.cat([input, label])
+                                        input = aug(input)
+                                        input, label = torch.chunk(input, 2)
+                                        label = self.bin(label)
+                                    else:
+                                        input = aug(input)
                             pred = model(input)
                             if val_post:
                                 for aug in val_post:
@@ -561,6 +725,8 @@ def train(self, train_data, val_data=None, oname=None,
                             acc_metric = self.metric(pred, label)
                             acc_epoch.append(acc_metric)
                         acc_log[int(epoch//self.val_interval)] = np.nanmean(acc_epoch)
+                        self.writer.add_scalar('class_val/dice_loss', np.nanmean(acc_epoch), epoch) 
+                        self.writer.add_scalar('class_val/dice_coefficient', 1-np.nanmean(acc_epoch), epoch)
                     if epoch % self.print_interval == 0:
                         print('Mean Validation acc: {}'.format(acc_log[int(epoch//self.val_interval)]))
                     if epoch >= self.val_interval:
diff --git a/src/rectangle/utils/transforms.py b/src/rectangle/utils/transforms.py
index dbe1ebe..05a84a1 100644
--- a/src/rectangle/utils/transforms.py
+++ b/src/rectangle/utils/transforms.py
@@ -32,8 +32,8 @@ def __call__(self, 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_,...]))
+                                torch.mean(image[batch_iter_,...]))/ \
+                                torch.std(image[batch_iter_,...])
     return image
 
 
@@ -336,6 +336,7 @@ def __init__(self, threshold=0.5):
     super().__init__()
     self.threshold = threshold
 
+  # TODO: check for torch auto binarising
   def __call__(self, image):
     return (image > self.threshold).int()
 
diff --git a/test.py b/test.py
new file mode 100644
index 0000000..ef4175d
--- /dev/null
+++ b/test.py
@@ -0,0 +1,139 @@
+import os
+import argparse
+
+parser = argparse.ArgumentParser(prog='test',
+                                description="Test RectAngle model. See list of available arguments for more info.")
+
+parser.add_argument('--test',
+                    '--te',
+                    metavar='test',
+                    type=str,
+                    action='store',
+                    default=None,
+                    help='Path to test data.')
+
+parser.add_argument('--ensemble',
+                    '--en',
+                    metavar='ensemble',
+                    type=str,
+                    action='store',
+                    default=None,
+                    help='Number of ensembled models.')
+
+parser.add_argument('--weights',
+                    '--w',
+                    metavar='weights',
+                    type=str,
+                    nargs='*',
+                    action='store',
+                    default=None,
+                    help='Path to saved model weights.')
+
+parser.add_argument('--gate',
+                    '--g',
+                    metavar='gate',
+                    type=str,
+                    action='store',
+                    default=None,
+                    help='(Optional) Attention gating.')
+
+parser.add_argument('--odir',
+                    '--o',
+                    metavar='odir',
+                    type=str,
+                    action='store',
+                    default='./',
+                    help='Path to output folder.')
+
+parser.add_argument('--depth',
+                    '--d',
+                    metavar='depth',
+                    type=str,
+                    action='store',
+                    default='5',
+                    help='Depth of U-Net architecture used.')
+
+parser.add_argument('--classifier',
+                    '--c',
+                    metavar='classifier',
+                    type=bool,
+                    action='store',
+                    default=False,
+                    help='Use of classifier for pre-screening. If selected will train without and then perform test without + with.')
+
+parser.add_argument('--classweights',
+                    '--cw',
+                    metavar='classweights',
+                    type=str,
+                    action='store',
+                    default=None,
+                    help='Path to trained weights for classifier.')
+
+parser.add_argument('--threshold',
+                    '--th',
+                    metavar='threshold',
+                    type=str,
+                    action='store',
+                    default='0.5',
+                    help='Activation threshold for classifier.')
+
+parser.add_argument('--seed',
+                    '--s',
+                    metavar='seed',
+                    type=str,
+                    action='store',
+                    default=None,
+                    help='Random seed for training.')
+
+args = parser.parse_args()
+
+## convert arguments to useable form
+if args.ensemble:
+    ensemble = int(args.ensemble)
+else:
+    if args.weights:
+        ensemble=int(len(args.weights))
+    else:
+        ensemble = 1
+
+## run training
+import rectangle as rect
+import h5py
+import torch
+import random
+import numpy as np
+
+# set seeds for repeatable results
+if args.seed:
+    seed = int(args.seed)
+    torch.manual_seed(seed)
+    random.seed(seed)
+    np.random.seed(seed)
+
+if torch.cuda.is_available():
+    device = torch.device('cuda')
+    torch.backends.cudnn.benchmark = True
+else:
+    device = torch.device('cpu')
+
+model = [rect.model.networks.UNet(n_layers=int(args.depth), device=device,
+                                    gate=args.gate) for e in range(ensemble)]
+
+for n, m in enumerate(model):
+    m.load_state_dict(torch.load(args.weights[n], map_location=device))
+
+if args.classifier==True:
+    class_model = rect.model.networks.MakeDenseNet(freeze_weights=False).to(device)
+if args.classweights:
+    class_model.load_state_dict(torch.load(args.classweights))
+
+f_test = h5py.File(args.test, 'r')
+if args.classifier:
+    test_data = rect.utils.io.PreScreenLoader(class_model.eval(), f_test, label=args.label, threshold=float(args.thresh))
+else:
+    test_data = rect.utils.io.H5DataLoader(f_test, label='vote')
+
+trainer = rect.utils.train.Trainer(model, ensemble=ensemble, outdir=args.odir, device=device)
+
+trainer.test(test_data, test_pre=[rect.utils.transforms.z_score()], oname='run', 
+            test_post=[rect.utils.transforms.Binary()], overlap='contour')
diff --git a/train.py b/train.py
index 973a6d6..c6077df 100644
--- a/train.py
+++ b/train.py
@@ -1,8 +1,8 @@
-## CLI for running training
-
 import os
 import argparse
 
+from rectangle.utils.transforms import Affine
+
 parser = argparse.ArgumentParser(prog='train',
                                 description="Train RectAngle model. See list of available arguments for more info.")
 
@@ -22,14 +22,6 @@
                     default=None,
                     help='Path to validation data.')
 
-parser.add_argument('--test',
-                    '--te',
-                    metavar='test',
-                    type=str,
-                    action='store',
-                    default=None,
-                    help='Path to test data.')
-
 parser.add_argument('--label',
                     '--l',
                     metavar='label',
@@ -54,6 +46,14 @@
                     default=None,
                     help='(Optional) Attention gating.')
 
+parser.add_argument('--lr_schedule',
+                    '--lrs',
+                    metavar='lr_schedule',
+                    type=str,
+                    action='store',
+                    default=None,
+                    help="Method for scheduling of learning rate. {None, 'lambda', 'exponential', 'reduce_on_plateau'}")
+
 parser.add_argument('--odir',
                     '--o',
                     metavar='odir',
@@ -86,14 +86,6 @@
                     default='32',
                     help='Batch size. Note images are large (~400x~300).')
 
-parser.add_argument('--classifier',
-                    '--c',
-                    metavar='classifier',
-                    type=bool,
-                    action='store',
-                    default=True,
-                    help='Use of classifier for pre-screening. If selected will train without and then perform test without + with.')
-
 parser.add_argument('--seed',
                     '--s',
                     metavar='seed',
@@ -102,6 +94,14 @@
                     default=None,
                     help='Random seed for training.')
 
+parser.add_argument('--earlystop',
+                    '--e',
+                    metavar='earlystop',
+                    type=str,
+                    action='store',
+                    default='10',
+                    help='Number of val steps with no improvement before stopping training early.')
+
 
 args = parser.parse_args()
 
@@ -118,6 +118,9 @@
 import random
 import numpy as np
 
+print("Code running")
+#os.environ["CUDA_VISIBLE_DEVICES"]="0"
+
 # set seeds for repeatable results
 if args.seed:
     seed = int(args.seed)
@@ -131,52 +134,29 @@
 if torch.cuda.is_available():
     device = torch.device('cuda')
     torch.backends.cudnn.benchmark = True
+    print("Cuda available!")
 else:
     device = torch.device('cpu')
+    print("Using CPU!")
 
 if args.val:
     f_val = h5py.File(args.val, 'r')
     val_data = rect.utils.io.H5DataLoader(f_val, label='vote')
 
-if args.test:
-    f_test = h5py.File(args.test, 'r')
-    test_data = rect.utils.io.H5DataLoader(f_test, label='vote')
-
 model = rect.model.networks.UNet(n_layers=int(args.depth), device=device,
                                     gate=args.gate)
 
-trainer = rect.utils.train.Trainer(model, ensemble=ensemble, outdir=args.odir,
-                                    nb_epochs=int(args.epochs))
+trainer = rect.utils.train.Trainer(model, ensemble=ensemble, outdir=args.odir, device=device,
+                                    nb_epochs=int(args.epochs), lr_schedule=args.lr_schedule,
+                                    early_stop=int(args.earlystop))
+
+#Manually setting Affine Transforms
+AffineTransform = rect.utils.transforms.Affine(prob = 0.3, scale = (1,1), degrees = 5, shear = 0, translate = 0)
 
 if args.val:
-    trainer.train(train_data, val_data, train_pre=[rect.utils.transforms.z_score(), rect.utils.transforms.Flip(), rect.utils.transforms.Affine(), rect.utils.transforms.SpeckleNoise()], 
+    trainer.train(train_data, val_data, train_pre=[rect.utils.transforms.z_score(), rect.utils.transforms.Flip(), AffineTransform],
                     val_pre=[rect.utils.transforms.z_score()], train_batch=int(args.batch))
 else:
-    trainer.train(train_data, train_pre=[rect.utils.transforms.z_score(), rect.utils.transforms.Flip(), rect.utils.transforms.Affine(), rect.utils.transforms.SpeckleNoise()], 
+    trainer.train(train_data, train_pre=[rect.utils.transforms.z_score(), rect.utils.transforms.Flip(), AffineTransform], 
                     val_pre=[rect.utils.transforms.z_score()], train_batch=int(args.batch))
 
-if args.test:
-    trainer.test(test_data, test_pre=[rect.utils.transforms.z_score()], 
-              test_post=[rect.utils.transforms.Binary(), rect.utils.transforms.KeepLargestComponent()])
-
-if args.classifier:
-    class_train_data = rect.utils.io.ClassifyDataLoader(f_train)
-    if args.val:
-        class_val_data = rect.utils.io.ClassifyDataLoader(f_val)
-    else:
-        class_val_data = None
-
-    class_model = rect.model.networks.MakeDenseNet(freeze_weights=False).to(device)
-    class_trainer = rect.utils.train.ClassTrainer(class_model, outdir=os.path.join(args.odir, 'classlogs'),
-                                         ensemble=None, early_stop=1000)
-
-    class_trainer.train(class_train_data, class_val_data, train_batch=int(args.batch))
-
-    threshRange = np.linspace(0, 0.6, 20)
-
-    if args.test:
-        for i, thresh in enumerate(threshRange):
-            test_screen_data = rect.utils.io.PreScreenLoader(class_model.eval(), f_test, label='vote', threshold = thresh)
-            trainer.test(test_screen_data, test_pre=[rect.utils.transforms.z_score()], 
-                        test_post=[rect.utils.transforms.Binary(), rect.utils.transforms.KeepLargestComponent()], oname='class_thresh_{}'.format(i))
-
diff --git a/train_classifier.py b/train_classifier.py
new file mode 100644
index 0000000..bc9d670
--- /dev/null
+++ b/train_classifier.py
@@ -0,0 +1,131 @@
+import os
+import argparse
+
+parser = argparse.ArgumentParser(prog='train',
+                                description="Train RectAngle model. See list of available arguments for more info.")
+
+parser.add_argument('--train',
+                    '--tr',
+                    metavar='train',
+                    type=str,
+                    action='store',
+                    default='./miccai_us_data/train.h5',
+                    help='Path to training data. Note that for ensemble this should include train + val pre-split.')
+
+parser.add_argument('--val',
+                    '--v',
+                    metavar='val',
+                    type=str,
+                    action='store',
+                    default=None,
+                    help='Path to validation data.')
+
+parser.add_argument('--test',
+                    '--te',
+                    metavar='test',
+                    type=str,
+                    action='store',
+                    default=None,
+                    help='Path to test data.')
+
+parser.add_argument('--ensemble',
+                    '--en',
+                    metavar='ensemble',
+                    type=str,
+                    action='store',
+                    default=None,
+                    help='Number of ensembled models.')
+
+parser.add_argument('--freeze',
+                    '--f',
+                    metavar='freeze',
+                    type=bool,
+                    action='store',
+                    default=False,
+                    help='Freeze CNN weights (pre-trained on ImageNet).')
+
+parser.add_argument('--odir',
+                    '--o',
+                    metavar='odir',
+                    type=str,
+                    action='store',
+                    default='./',
+                    help='Path to output folder.')
+
+parser.add_argument('--epochs',
+                    '--ep',
+                    metavar='epochs',
+                    type=str,
+                    action='store',
+                    default='200',
+                    help='Max number of training epochs per model.')
+
+parser.add_argument('--batch',
+                    '--b',
+                    metavar='batch',
+                    type=str,
+                    action='store',
+                    default='32',
+                    help='Batch size. Note images are large (~400x~300).')
+
+parser.add_argument('--seed',
+                    '--s',
+                    metavar='seed',
+                    type=str,
+                    action='store',
+                    default=None,
+                    help='Random seed for training.')
+
+
+args = parser.parse_args()
+
+## convert arguments to useable form
+if args.ensemble:
+    ensemble = int(args.ensemble)
+else:
+    ensemble = None
+
+## run training
+import rectangle as rect
+import h5py
+import torch
+import random
+import numpy as np
+
+# set seeds for repeatable results
+if args.seed:
+    seed = int(args.seed)
+    torch.manual_seed(seed)
+    random.seed(seed)
+    np.random.seed(seed)
+
+if torch.cuda.is_available():
+    device = torch.device('cuda')
+    torch.backends.cudnn.benchmark = True
+else:
+    device = torch.device('cpu')
+
+f_train = h5py.File(args.train, 'r')
+
+if args.val:
+    f_val = h5py.File(args.val, 'r')
+
+if args.test:
+    f_test = h5py.File(args.test, 'r')
+
+class_train_data = rect.utils.io.ClassifyDataLoader_v2(f_train)
+if args.val:
+    class_val_data = rect.utils.io.ClassifyDataLoader_v2(f_val)
+else:
+    class_val_data = None
+if args.test:
+    class_test_data = rect.utils.io.ClassifyDataLoader_v2(f_test)
+else:
+    class_test_data = None
+
+class_model = rect.model.networks.MakeDenseNet(freeze_weights=args.freeze).to(device)
+class_trainer = rect.utils.train.ClassTrainer(class_model, outdir=os.path.join(args.odir),
+                                        ensemble=ensemble, nb_epochs=int(args.epochs), device=device)
+
+class_trainer.train(class_train_data, class_val_data, train_pre=[rect.utils.transforms.z_score(), rect.utils.transforms.Flip(), rect.utils.transforms.Affine()],
+                    val_pre=[rect.utils.transforms.z_score()], train_batch=int(args.batch))