Upload 5 files

augment.py

@@ -0,0 +1,89 @@
+from torchvision import  transforms
+import random
+import numpy as np
+
+class RandAug:
+    """Randomly chosen image augmentations."""
+
+    def __init__(self, img_size, choice=None):
+        # Augmentation options
+        self.trans = ['identity', 'rotate', 'color', 'sharpness', 'blur', 'padding' ,'perspective']
+        self.img_size = img_size
+        self.choice = choice
+
+    def __call__(self, img):
+        if self.choice == None:
+            # Weights set 40% probability for the 'identity' augmentation choice
+            self.choice = random.choices(self.trans, weights=(40, 10, 10, 10, 10, 10, 10))[0]
+
+        if self.choice == 'identity':
+            trans = transforms.Compose([
+                            transforms.Resize((self.img_size,self.img_size)),
+                            transforms.ToTensor()
+                        ])
+            img = trans(img)
+
+        elif self.choice == 'rotate':
+            degrees = random.uniform(0, 180)
+            rand_fill = random.choice([0,1])
+            trans = transforms.Compose([
+                            transforms.Resize((self.img_size,self.img_size)),
+                            transforms.ToTensor(),
+                            transforms.RandomRotation(degrees, expand=True, fill=rand_fill),
+                            transforms.Resize((self.img_size,self.img_size))
+                        ])
+            img = trans(img)
+
+        elif self.choice == 'color':
+            rand_brightness = random.uniform(0, 0.3)
+            rand_hue = random.uniform(0, 0.5)
+            rand_contrast = random.uniform(0, 0.5)
+            rand_saturation = random.uniform(0, 0.5)
+            trans = transforms.Compose([
+                            transforms.Resize((self.img_size,self.img_size)),
+                            transforms.ToTensor(),
+                            transforms.ColorJitter(brightness=rand_brightness, contrast=rand_contrast, saturation=rand_saturation, hue=rand_hue)
+                        ])
+            img = trans(img)
+
+        elif self.choice=='sharpness':
+            sharpness = 1+(np.random.exponential()/2)
+            trans = transforms.Compose([
+                            transforms.Resize((self.img_size,self.img_size)),
+                            transforms.ToTensor(),
+                            transforms.RandomAdjustSharpness(sharpness, p=1)
+                        ])
+            img = trans(img)
+
+        elif self.choice=='blur':
+            kernel = random.choice([1,3,5])
+            trans = transforms.Compose([
+                            transforms.Resize((self.img_size,self.img_size)),
+                            transforms.ToTensor(),
+                            transforms.GaussianBlur(kernel, sigma=(0.1, 2.0))
+                        ])
+            img = trans(img)
+
+        elif self.choice=='padding':
+            pad = random.choice([3,10,25])
+            rand_fill = random.choice([0,1])
+            trans = transforms.Compose([
+                            transforms.Resize((self.img_size,self.img_size)),
+                            transforms.ToTensor(),
+                            transforms.Pad(pad, fill=rand_fill, padding_mode='constant'),
+                            transforms.Resize((self.img_size,self.img_size))
+                        ])
+            img = trans(img)
+
+        elif self.choice=='perspective':
+            scale = random.uniform(0.1, 0.5)
+            rand_fill = random.choice([0,1])
+            trans = transforms.Compose([
+                            transforms.Resize((self.img_size,self.img_size)),
+                            transforms.ToTensor(),
+                            transforms.RandomPerspective(distortion_scale=scale, p=1.0, fill=rand_fill),
+                            transforms.Resize((self.img_size,self.img_size))
+                        ])
+            img = trans(img)
+            
+        return img

requirements.txt

@@ -0,0 +1,10 @@
+--extra-index-url https://download.pytorch.org/whl/cu116
+torch==1.12.1+cu116
+torchvision==0.13.1+cu116
+scikit-learn==1.0.2
+numpy==1.21.6
+pillow==9.3.0 
+matplotlib==3.5.3
+onnx==1.13.0 
+onnxruntime==1.13.1
+tqdm==4.64.1

test.py

@@ -0,0 +1,192 @@
+from __future__ import print_function
+from __future__ import division
+import torch
+import onnxruntime
+import numpy as np
+import pandas as pd
+import torchvision
+from torchvision import transforms
+import matplotlib.pyplot as plt
+from sklearn.metrics import precision_recall_fscore_support, accuracy_score, confusion_matrix, ConfusionMatrixDisplay
+import seaborn as sn
+import random
+import time
+import json
+from PIL import Image
+from PIL import ImageFile
+from pathlib import Path
+import argparse
+print("PyTorch Version: ",torch.__version__)
+print("Torchvision Version: ",torchvision.__version__)
+
+parser = argparse.ArgumentParser('arguments for testing the model')
+
+parser.add_argument('--ts_empty_folder', type=str, default="/data/taulukot/solukuvat/empty/test/",
+                    help='path to test data')
+parser.add_argument('--ts_ok_folder', type=str, default="/data/taulukot/solukuvat/ok/test/",
+                    help='path to test data')
+parser.add_argument('--results_folder', type=str, default="./results/aug_28022024/",
+                    help='Folder for saving results')
+parser.add_argument('--model_path', type=str, default="/koodit/table_segmentation/empty_cell_detection/train/models/aug_b32_lr0001_28022024.onnx",
+                    help='path to load model file from')
+parser.add_argument('--batch_size', type=int, default=16,
+                    help='batch_size')
+parser.add_argument('--num_classes', type=int, default=2,
+                    help='number of classes for classification')
+parser.add_argument('--name', type=str, default='empty_cell_augment_28022024',
+                    help='name given to result files')
+
+start = time.time()
+
+# nohup python test.py > logs/aug_test_28022024.txt 2>&1 &
+# echo $! > output/save_pid.txt
+
+torch.manual_seed(67)
+random.seed(67)
+
+args = parser.parse_args()
+
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+Image.MAX_IMAGE_PIXELS = None
+
+# https://pytorch.org/tutorials/beginner/finetuning_torchvision_models_tutorial.html
+
+
+def get_data():
+    empty_path = Path(args.ts_empty_folder)
+    ok_path = Path(args.ts_ok_folder)
+
+    empty_files = list(empty_path.glob('*.jpg'))
+    ok_files = list(ok_path.glob('*.jpg'))
+
+    empty_labels = np.zeros(len(empty_files))
+    ok_labels = np.ones(len(ok_files))
+
+    #ts_data_files = ts_data_files[:20]
+    #ts_data_labels = ts_data_labels[:20]
+    #ts_ok_files = ts_ok_files[:20]
+    #ts_ok_labels = ts_ok_labels[:20]
+
+    ts_files = empty_files + ok_files
+    ts_labels = np.concatenate((empty_labels, ok_labels))
+
+    print('Test data with empty cells: ', len(empty_files))
+    print('Test data without empty cells: ', len(ok_files))
+
+    return ts_files, ts_labels
+
+
+def initialize_model():
+    model = onnxruntime.InferenceSession(args.model_path)
+    input_size = 224
+    return model, input_size
+
+# Function for getting precision, recall and F-score metrics
+def get_precision_recall(y_true, y_pred):
+    precision_recall_fscore = precision_recall_fscore_support(y_true, y_pred, average=None)
+
+    prec_0 = precision_recall_fscore[0][0]
+    rec_0 = precision_recall_fscore[1][0]
+    F_0 = precision_recall_fscore[2][0]
+
+    prec_1 = precision_recall_fscore[0][1]
+    rec_1 = precision_recall_fscore[1][1]
+    F_1 = precision_recall_fscore[2][1]
+
+    print('\nPrecision for ok: %.2f'%prec_1)
+    print('Recall for ok: %.2f'%rec_1)
+    print('F-score for ok: %.2f'%F_1)
+
+    print('Precision for empty: %.2f'%prec_0 )
+    print('Recall for empty: %.2f'%rec_0)
+    print('F-score for empty: %.2f'%F_0)
+
+
+def createConfusionMatrix(y_true, y_pred):
+    classes = np.array(['empty', 'ok'])
+
+    # Build confusion matrix
+    cf_matrix = confusion_matrix(y_true, y_pred)
+    print(cf_matrix)
+    df_cm = pd.DataFrame(cf_matrix, index=classes,
+                         columns=classes)
+    plt.figure(figsize=(12, 7))    
+    return sn.heatmap(df_cm, annot=True).get_figure()
+
+def save_preds(y_true, y_pred, paths):
+    # Identifies images that were not classified correctly
+    incorrect_indices = np.where(y_true != y_pred)
+    incorrectly_predicted_images = paths[incorrect_indices]
+    correct_labels = y_true[incorrect_indices].astype(str)
+    incorrect_preds = dict(zip(incorrectly_predicted_images, correct_labels))
+
+    print(f'{len(incorrect_preds)} incorrect predictions')
+
+    # Save file names and labels of incorrectly classified images
+    with open(args.results_folder + args.name + '_incorrect_preds', "w") as fp:
+        json.dump(incorrect_preds, fp)
+
+# Initialize the model for this run
+model, input_size = initialize_model()
+
+# Print the model we just instantiated
+#print(model_ft)
+
+data_transforms = transforms.Compose([
+        transforms.Resize((input_size, input_size)),
+        transforms.ToTensor()
+    ])
+
+print("Initializing Datasets and Dataloaders...")
+
+ts_files, ts_labels = get_data()
+
+# Function for getting model predictions on test data
+def test_model(model, ts_files, ts_labels):
+    since = time.time()
+    label_preds = []
+    true_labels = []
+    paths = []
+    n = len(ts_files)
+    # Iterate over data
+    for i in range(n):
+        print(f'{i}/{n}')
+        image = Image.open(ts_files[i])
+        label = ts_labels[i]
+        image = data_transforms(image.convert("RGB")).unsqueeze(0)
+        # Transform tensor to numpy array
+        img = image.detach().cpu().numpy()
+        input = {model.get_inputs()[0].name: img}
+        # Run model prediction
+        output = model.run(None, input)
+        # Get predicted class
+        pred = np.argmax(output[0], 1)
+        pred_class = pred.item()
+        label_preds.append(pred_class)
+        true_labels.append(label)
+        paths.append(str(ts_files[i]))
+
+    time_elapsed = time.time() - since
+    print('Testing complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
+
+    return np.array(label_preds), np.array(true_labels), np.array(paths)
+
+ts_labels = np.array(ts_labels)
+
+# Test model
+y_pred, y_true, paths = test_model(model, ts_files, ts_labels)
+# Saves information of incorrect predictions
+save_preds(y_true, y_pred, paths)
+# Calculates and prints precision, recall and F-score metrics
+get_precision_recall(y_true, y_pred)
+
+# Save confusion matrix to Tensorboard
+#cm = createConfusionMatrix(y_true, y_pred)
+#writer.add_figure("Confusion matrix", cm)
+# Create and save confusion matrix of the predictions and true labels
+conf_matrix = ConfusionMatrixDisplay.from_predictions(y_true, y_pred, normalize='true', display_labels=np.array(['empty', 'ok']))
+plt.savefig(args.results_folder + args.name + '_conf_matrix.jpg', bbox_inches='tight')
+
+end = time.time()
+time_in_mins = (end - start) / 60
+print('Time: %.2f minutes' % time_in_mins)

train.py

@@ -0,0 +1,332 @@
+from __future__ import print_function
+from __future__ import division
+import torch
+import torchvision
+import torch.nn as nn
+from torch.utils.data import Dataset, DataLoader
+from torchvision import  models
+from torch.optim.lr_scheduler import ReduceLROnPlateau
+from sklearn.utils import class_weight
+from sklearn.metrics import precision_recall_fscore_support
+import numpy as np
+import time
+import argparse
+from tqdm import tqdm
+from PIL import Image, ImageFile
+from pathlib import Path
+
+from augment import RandAug
+import utils
+
+print("PyTorch Version: ",torch.__version__)
+print("Torchvision Version: ",torchvision.__version__)
+
+# Much of the code is a modified version of the code available at
+# https://pytorch.org/tutorials/beginner/finetuning_torchvision_models_tutorial.html
+
+
+# nohup python train.py > logs/empty_cell_aug_28032024.txt 2>&1 &
+# echo $! > logs/save_pid.txt
+
+parser = argparse.ArgumentParser('arguments for training')
+
+parser.add_argument('--tr_empty_folder', type=str, default="/data/taulukot/solukuvat/empty/train/",
+                    help='path to training data with empty images')
+parser.add_argument('--val_empty_folder', type=str, default="/data/taulukot/solukuvat/empty/val/",
+                    help='path to validation data with empty images')
+parser.add_argument('--tr_ok_folder', type=str, default="/data/taulukot/solukuvat/ok/train/",
+                    help='path to training data with ok images')
+parser.add_argument('--val_ok_folder', type=str, default="/data/taulukot/solukuvat/ok/val/",
+                    help='path to validation data with ok images')
+parser.add_argument('--results_folder', type=str, default="results/28032024_aug/",
+                    help='Folder for saving training results.')
+parser.add_argument('--save_model_path', type=str, default="./models/",
+                    help='Path for saving model file.')
+parser.add_argument('--batch_size', type=int, default=32,
+                    help='Batch size used for model training. ')
+parser.add_argument('--lr', type=float, default=0.0001,
+                    help='Base learning rate.')
+parser.add_argument('--device', type=str, default='cpu',
+                    help='Defines whether the model is trained using cpu or gpu.')
+parser.add_argument('--num_classes', type=int, default=2,
+                    help='Number of classes used in classification.')
+parser.add_argument('--num_epochs', type=int, default=15,
+                    help='Number of training epochs.')
+parser.add_argument('--random_seed', type=int, default=8765,
+                    help='Number used for initializing random number generation.')
+parser.add_argument('--early_stop_threshold', type=int, default=3,
+                    help='Threshold value of epochs after which training stops if validation accuracy does not improve.')
+parser.add_argument('--save_model_format', type=str, default='torch',
+                    help='Defines the format for saving the model.')
+parser.add_argument('--augment_choice', type=str, default=None,
+                    help='Defines which image augmentation(s) are used. Defaults to randomly selected augmentations.')
+parser.add_argument('--model_name', type=str, default='aug_b32_lr0001',
+                    help='Current date.')
+parser.add_argument('--date', type=str, default=time.strftime("%d%m%Y"),
+                    help='Current date.')
+
+args = parser.parse_args()
+
+# PIL settings to avoid errors caused by truncated and large images
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+Image.MAX_IMAGE_PIXELS = None
+
+# List for saving the names of damaged images
+damaged_images = []
+
+def get_datapaths():
+    """Function for loading train and validation data."""
+    tr_empty_files = list(Path(args.tr_empty_folder).glob('*'))
+    tr_ok_files = list(Path(args.tr_ok_folder).glob('*'))
+    val_empty_files = list(Path(args.val_empty_folder).glob('*'))
+    val_ok_files = list(Path(args.val_ok_folder).glob('*'))
+    # Create labels for train and validation data
+    tr_labels = np.concatenate((np.zeros(len(tr_empty_files)), np.ones(len(tr_ok_files))))
+    val_labels = np.concatenate((np.zeros(len(val_empty_files)), np.ones(len(val_ok_files))))
+    # Combine faulty and non-faulty images
+    tr_files = tr_empty_files + tr_ok_files
+    val_files = val_empty_files + val_ok_files
+
+    print('\nTraining data with empty cells: ', len(tr_empty_files))
+    print('Training data without empty cells: ', len(tr_ok_files))
+
+    print('Validation data with empty cells: ', len(val_empty_files))
+    print('Validation data without empty cells: ', len(val_ok_files))
+
+    data_dict = {'tr_data': tr_files, 'tr_labels': tr_labels, 
+            'val_data': val_files, 'val_labels': val_labels}
+
+    return data_dict
+
+class ImageDataset(Dataset):
+    """PyTorch Dataset class is used for generating training and validation datasets."""
+    def __init__(self, img_paths, img_labels, transform=None, target_transform=None):
+        self.img_paths = img_paths
+        self.img_labels = img_labels
+        self.transform = transform
+        self.target_transform = target_transform
+
+    def __len__(self):
+        return len(self.img_labels)
+
+    def __getitem__(self, idx):
+        img_path = self.img_paths[idx]
+        try:
+            image = Image.open(img_path).convert('RGB')
+            label = self.img_labels[idx]
+        except:
+            # Image is considered damaged if reading the image fails
+            damaged_images.append(img_path)
+            return None
+        if self.transform:
+            image = self.transform(image.convert("RGB"))
+        if self.target_transform:
+            label = self.target_transform(label)
+            
+        return image, label
+
+def initialize_model():
+    """Function for initializing pretrained neural network model (DenseNet121)."""
+    model_ft = models.densenet121(weights=torchvision.models.DenseNet121_Weights.IMAGENET1K_V1)
+    num_ftrs = model_ft.classifier.in_features
+    model_ft.classifier = nn.Linear(num_ftrs, args.num_classes)
+    input_size = 224
+
+    return model_ft, input_size
+
+def collate_fn(batch):
+    """Helper function for creating data batches."""
+    batch = list(filter(lambda x: x is not None, batch))
+ 
+    return torch.utils.data.dataloader.default_collate(batch)
+
+def initialize_dataloaders(data_dict, input_size):
+    """Function for initializing datasets and dataloaders."""
+    # Train and validation datasets 
+    train_dataset = ImageDataset(img_paths=data_dict['tr_data'], img_labels=data_dict['tr_labels'],  transform=RandAug(input_size, args.augment_choice))
+    validation_dataset = ImageDataset(img_paths=data_dict['val_data'], img_labels=data_dict['val_labels'], transform=RandAug(input_size, 'identity'))
+    # Train and validation dataloaders
+    train_dataloader = DataLoader(train_dataset, collate_fn=collate_fn, batch_size=args.batch_size, shuffle=True, num_workers=4)
+    validation_dataloader = DataLoader(validation_dataset, collate_fn=collate_fn, batch_size=args.batch_size, shuffle=True, num_workers=4)
+    
+    return {'train': train_dataloader, 'val': validation_dataloader}
+
+def get_criterion(data_dict):
+    """Function for generating class weights and initializing the loss function."""
+    y = np.asarray(data_dict['tr_labels'])
+    # Class weights are used for compensating the unbalance 
+    # in the number of training data from the two classes
+    class_weights=class_weight.compute_class_weight(class_weight='balanced', classes=np.unique(y), y=y)
+    class_weights=torch.tensor(class_weights, dtype=torch.float).to(args.device)
+    print('\nClass weights: ', class_weights.tolist())
+     # Cross Entropy Loss function
+    criterion = nn.CrossEntropyLoss(weight=class_weights, reduction='mean')
+
+    return criterion
+
+def get_optimizer(model):
+    """Function for initializing the optimizer."""
+    # Model parameters are split into two groups: parameters of the classifier
+    # layer and other model parameters
+    params_1 = [param for name, param in model.named_parameters()
+                if name not in ["classifier.weight", "classifier.bias"]]
+    params_2 = model.classifier.parameters()
+    # 10 x larger learning rate is used when training the parameters 
+    # of the classification layers
+    params_to_update = [
+            {'params': params_1, 'lr': args.lr},
+            {'params': params_2, 'lr': args.lr * 10}
+            ]
+    # Adam optimizer
+    optimizer = torch.optim.Adam(params_to_update, args.lr)
+    # Scheduler reduces learning rate when validation accuracy does not improve for an epoch
+    scheduler = ReduceLROnPlateau(optimizer, mode='max', factor=0.1, patience=0, verbose=True)
+
+    return optimizer, scheduler
+
+def train_model(model, dataloaders, criterion, optimizer, scheduler=None):
+    """Function for model training and validation."""
+    since = time.time()
+    # Lists for saving train and validation metrics for each epoch
+    tr_loss_history = []
+    tr_acc_history = []
+    tr_f1_history = []
+    val_loss_history = []
+    val_acc_history = []
+    val_f1_history = []
+    # Lists for saving learning rates for the 2 parameter groups
+    lr1_history = []
+    lr2_history = []
+    
+    # Best F1 value and best epoch are saved in variables
+    best_f1 = 0
+    best_epoch = 0
+    early_stop = False
+
+    # Train / validation loop
+    for epoch in tqdm(range(args.num_epochs)):
+        # Save learning rates for the epoch
+        lr1_history.append(optimizer.param_groups[0]["lr"])
+        lr2_history.append(optimizer.param_groups[1]["lr"])
+        
+        print('Epoch {}/{}'.format(epoch+1, args.num_epochs))
+        print('-' * 10)
+
+        # Each epoch has a training and validation phase
+        for phase in ['train', 'val']:
+            if phase == 'train':
+                model.train()  # Set model to training mode
+            else:
+                model.eval()   # Set model to evaluate mode
+
+            running_loss = 0.0
+            running_corrects = 0
+            running_f1 = 0.0
+
+            # Iterate over data in batch
+            for inputs, labels in dataloaders[phase]:
+                if dataloaders[phase] is None:
+                    continue
+                else:
+                    inputs = inputs.to(args.device)
+                    labels = labels.long().to(args.device)
+
+                    # Zero the parameter gradients
+                    optimizer.zero_grad()
+
+                    # Track history only in training phase
+                    with torch.set_grad_enabled(phase == 'train'):
+                        # Get model outputs and calculate loss
+                        outputs = model(inputs)
+                        loss = criterion(outputs, labels)
+                        # Model predictions of the image labels for the batch
+                        _, preds = torch.max(outputs, 1)
+
+                        # Backward + optimize only if in training phase
+                        if phase == 'train':
+                            loss.backward()
+                            optimizer.step()
+                    
+                    # Get weighted F1 score for the results
+                    precision_recall_fscore = precision_recall_fscore_support(labels.data.detach().cpu().numpy(), preds.detach().cpu().numpy(), average='weighted', zero_division=0)
+                    f1_score = precision_recall_fscore[2]
+
+                    # update statistics
+                    running_loss += loss.item() * inputs.size(0)
+                    running_corrects += torch.sum(preds == labels.data).cpu()
+                    running_f1 += f1_score
+
+            # Calculate loss, accuracy and F1 score for the epoch
+            epoch_loss = running_loss / len(dataloaders[phase].dataset)
+            epoch_acc = running_corrects.double() / len(dataloaders[phase].dataset)
+            epoch_f1 = running_f1 / len(dataloaders[phase])
+
+            print('\nEpoch {} - {} - Loss: {:.4f} Acc: {:.4f} F1: {:.4f}\n'.format(epoch+1, phase, epoch_loss, epoch_acc, epoch_f1))
+            
+            # Validation step
+            if phase == 'val':
+                val_acc_history.append(epoch_acc)
+                val_loss_history.append(epoch_loss)
+                val_f1_history.append(epoch_f1)
+                if epoch_f1 > best_f1:
+                    print('\nF1 score {:.4f} improved from {:.4f}. Saving the model.\n'.format(epoch_f1, best_f1))
+                    # Model with best F1 score is saved
+                    utils.save_model(model, 224, args.save_model_format, args.save_model_path, args.model_name, args.date)
+                    model = model.to(args.device)
+                    best_f1 = epoch_f1
+                    best_epoch = epoch
+                elif epoch - best_epoch > args.early_stop_threshold:
+                    # terminates the training loop if validation accuracy has not improved
+                    print("Early stopped training at epoch %d" % epoch)
+                    # Set early stopping condition
+                    early_stop = True
+                    break  
+            elif phase == 'train':
+                tr_acc_history.append(epoch_acc)
+                tr_loss_history.append(epoch_loss)
+                tr_f1_history.append(epoch_f1)
+
+        # Break outer loop if early stopping condition is activated
+        if early_stop:
+            break
+        # Take scheduler step
+        if scheduler:
+            scheduler.step(val_f1_history[-1])
+
+    time_elapsed = time.time() - since
+    print('\nTraining complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
+    print('Best validation F1 score: {:.4f}'.format(best_f1))
+    # Returns model with the weights from the best epoch (based on validation accuracy)
+    hist_dict = {'tr_acc': tr_acc_history, 
+                 'val_acc': val_acc_history, 
+                 'val_loss': val_loss_history,
+                 'val_f1': val_f1_history,
+                 'tr_loss': tr_loss_history,
+                 'tr_f1': tr_f1_history,
+                 'lr1': lr1_history,
+                 'lr2': lr2_history}
+
+    return hist_dict
+
+def main():
+    # Set random seed(s)
+    utils.set_seed(args.random_seed)
+    # Load image paths and labels
+    data_dict = get_datapaths()
+    # Initialize the model 
+    model, input_size = initialize_model()
+    # Print the model architecture
+    #print(model_ft)
+    # Send the model to GPU (if available)
+    model = model.to(args.device)
+    print("\nInitializing Datasets and Dataloaders...")
+    dataloaders_dict = initialize_dataloaders(data_dict, input_size)
+    criterion = get_criterion(data_dict)
+    optimizer, scheduler = get_optimizer(model)
+    # Train and evaluate model
+    hist_dict = train_model(model, dataloaders_dict, criterion, optimizer, scheduler)
+    print('Damaged images: ', damaged_images)
+    utils.plot_metrics(hist_dict, args.results_folder, args.date)
+
+if __name__ == '__main__':
+    main()

utils.py

@@ -0,0 +1,107 @@
+import torch
+import onnx
+import onnxruntime
+import os
+import matplotlib.pyplot as plt
+import numpy as np
+import random
+
+def set_seed(random_seed):
+    """Function for setting random seed for the relevant libraries."""
+    np.random.seed(random_seed)
+    random.seed(random_seed)
+    torch.manual_seed(random_seed)
+    torch.cuda.manual_seed(random_seed)
+    # When running on the CuDNN backend, two further options must be set
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+    # Set a fixed value for the hash seed
+    os.environ["PYTHONHASHSEED"] = str(random_seed)
+    print(f"Random seed set as {random_seed}")
+
+def save_model(model, input_size, save_model_format, save_model_path, model_name, date):
+    """Function for saving the model in .pth or .onnx format.
+    Code modified from 
+    https://pytorch.org/tutorials/advanced/super_resolution_with_onnxruntime.html"""
+    if save_model_format == 'onnx':
+        onnx_model_path = os.path.join(save_model_path, model_name + '_' + date + '.onnx')
+        # Random batch size
+        batch_size = 1
+        # Random input to the model (with correct dimensions)
+        x = torch.randn(batch_size, 3, input_size, input_size, requires_grad=True)
+        model = model.to('cpu')
+        torch_out = model(x)
+
+        # Export the model
+        torch.onnx.export(model,                   # model being run
+                        x,                         # model input (or a tuple for multiple inputs)
+                        onnx_model_path,           # where to save the model (can be a file or file-like object)
+                        export_params=True,        # store the trained parameter weights inside the model file
+                        opset_version=10,          # the ONNX version to export the model to
+                        do_constant_folding=True,  # whether to execute constant folding for optimization
+                        input_names = ['input'],   # the model's input names
+                        output_names = ['output'], # the model's output names
+                        dynamic_axes={'input' : {0 : 'batch_size'},    # variable length axes
+                                        'output' : {0 : 'batch_size'}})
+
+        print('ONNX model saved to ', onnx_model_path)
+        # Test transformed model
+        onnx_model = onnx.load(onnx_model_path)
+        onnx.checker.check_model(onnx_model)
+        print('ONNX model checked.')
+
+        def to_numpy(tensor):
+            return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
+
+        onnx_session = onnxruntime.InferenceSession(onnx_model_path)
+        # compute ONNX Runtime output prediction
+        onnx_inputs = {onnx_session.get_inputs()[0].name: to_numpy(x)}
+        onnx_out = onnx_session.run(None, onnx_inputs)
+        # compare ONNX Runtime and PyTorch results
+        np.testing.assert_allclose(to_numpy(torch_out), onnx_out[0], rtol=1e-03, atol=1e-05)
+        print("Exported model has been tested with ONNXRuntime, and the result looks good!\n")
+
+    else:
+        pytorch_model_path = os.path.join(save_model_path, 'densenet_' + date + '.pth')
+        torch.save(model, pytorch_model_path)
+        print('Pytorch model saved to ', pytorch_model_path)
+
+
+def plot_metrics(hist_dict, results_folder, date):
+    """Function for plotting the training and validation results."""
+    epochs = range(1, len(hist_dict['tr_loss'])+1)
+    plt.plot(epochs, hist_dict['tr_loss'], 'g', label='Training loss')
+    plt.plot(epochs, hist_dict['val_loss'], 'b', label='Validation loss')
+    plt.title('Training and Validation loss')
+    plt.xlabel('Epochs')
+    plt.ylabel('Loss')
+    plt.legend()
+    plt.savefig(results_folder + date + '_tr_val_loss.jpg', bbox_inches='tight')
+    plt.close()
+
+    plt.plot(epochs, hist_dict['tr_acc'], 'g', label='Training accuracy')
+    plt.plot(epochs, hist_dict['val_acc'], 'b', label='Validation accuracy')
+    plt.title('Training and Validation accuracy')
+    plt.xlabel('Epochs')
+    plt.ylabel('Accuracy')
+    plt.legend()
+    plt.savefig(results_folder + date +  '_tr_val_acc.jpg', bbox_inches='tight')
+    plt.close()
+
+    plt.plot(epochs, hist_dict['tr_f1'], 'g', label='Training F1 score')
+    plt.plot(epochs, hist_dict['val_f1'], 'b', label='Validation F1 score')
+    plt.title('Training and Validation F1 score')
+    plt.xlabel('Epochs')
+    plt.ylabel('F1 score')
+    plt.legend()
+    plt.savefig(results_folder + date +  '_tr_val_f1.jpg', bbox_inches='tight')
+    plt.close()
+
+    plt.plot(epochs, hist_dict['lr1'], 'g', label='Backbone learning rate')
+    plt.plot(epochs, hist_dict['lr2'], 'b', label='Classifier learning rate')
+    plt.title('Learning rate')
+    plt.xlabel('Epochs')
+    plt.ylabel('Learning rate')
+    plt.legend()
+    plt.savefig(results_folder + date +  '_learning_rate.jpg', bbox_inches='tight')
+    plt.close()