train_model_utils.py

"""@package docstring 

@file train_model_utils.py 

@brief Group of functions that are supposed to be used directly in the training or evaluation scripts
 
@section libraries_train_model_utils Libraries 
- os
- torch
- numpy
- matplotlib.pyplot
- dataset_utils
- seaborn
- itertools
- datetime
- sklearn.metrics

@section classes_train_model_utils Classes 
- None
 
@section functions_train_model_utils Functions 
- model_statistics(J_history, predicted, ground_truth, pred_vector, gt_vector,
- compute_conf_matrix(predicted, ground_truth)
- compute_top_k_accuracy(pred_vector, gt_vector, topk)
- compute_num_splits_sent(pred_lst)
- compute_multi_thres_performance(pred_lst, gt_lst)
- compute_ROC_curve(pred_vector, gt_vector, pred_num)
- model_simple_metrics(predicted, ground_truth)
- obtain_best_modes(rs, pred)
- obtain_mode(pred)
- one_hot_enc(tensor, num_classes=6)
- print_parameters(model, optimizer)
- save_model_parameters(dir_name, f_name, model)
- save_model(dir_name, f_name, model, optimizer, loss, acc)
- load_model_parameters_stg(model, path, stg, dev)
- load_model_parameters_eval(model, path, dev)
- load_model_stg_12_stg_3(model, path, dev)
- load_model_stg_3_stg_4(model, path, dev)
- load_model_stg_4_stg_5(model, path, dev)
- load_model_stg_5_stg_6(model, path, dev)
- print_current_time()
 
@section global_vars_train_model_utils Global Variables 
- None

@section todo_train_model_utils TODO 
- None
    
@section license License 
MIT License 
Copyright (c) 2022 Raul Kevin do Espirito Santo Viana
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.

@section author_train_model_utils Author(s)
- Created by Raul Kevin Viana
- Last time modified is 2023-01-29 22:22:04.154941
"""

# ==============================================================
# Imports
# ==============================================================

import os
import torch
import numpy as np
import matplotlib.pyplot as plt
import dataset_utils
from datetime import datetime
from sklearn.metrics import classification_report, confusion_matrix, accuracy_score, f1_score, recall_score, \
    precision_score, roc_curve, auc
import seaborn as sns
from itertools import cycle

# ==============================================================
# Functions
# ==============================================================

def model_statistics(J_history, predicted, ground_truth, pred_vector, gt_vector,
                        f1_list, recall_list, precision_list,
                        accuracy_list, train_or_val="train"):
    """!
    @brief Evaluates model with metrics, such as accuracy and f1_score. This version plots the
           evolution of the metrics: f1-score, recall, precision, accuracy.

    @param [in] J_history: Loss function values over iterations
    @param [in] predicted: List of predictions made by the model with single value
    @param [in] ground_truth: List of the ground-truths with single value
    @param [in] pred_vector: List of predictions made by the model with vectors values
    @param [in] gt_vector: List of the ground-truths with vectors values
    @param [in] train_or_val: String that is used to write on the image files names
    @param [out] f1: F1 score
    @param [out] recall: Recall score
    @param [out] precision: Precision score
    @param [out] accuracy: Accuracy score
    """

    # Report
    report = classification_report(ground_truth, predicted)
    try:
        f = open("classification_report_" + train_or_val + ".txt", "x")
    except:
        os.remove("classification_report_" + train_or_val + ".txt")
        f = open("classification_report_" + train_or_val + ".txt", "x")

    f.write(report)
    f.close()

    # Plot Loss function throughout iterations
    plt.figure()
    plt.plot(J_history, label='loss')
    plt.xlabel("Batch number")
    plt.title("Loss: " + train_or_val)
    plt.legend()
    plt.grid()
    name = "Loss_" + train_or_val + ".png"
    plt.savefig(name)
    plt.clf()

    # Plot confusion matrix
    labels = ["Non-Split", "QT", "HBT", "VBT", "HTT", "VTT"]

    # Draw confusion matrix
    sns.heatmap(confusion_matrix(ground_truth, predicted), annot=True, fmt='d', cmap='Blues', yticklabels=labels,
                xticklabels=labels)
    name = "confusion_matrix_" + train_or_val + ".png"
    plt.title("Confusion Matrix: " + train_or_val)
    plt.savefig(name)
    plt.clf()

    # Plot metrics
    plt.figure()
    plt.plot(f1_list, 'go-', label='F1-score')
    plt.plot(recall_list, 'bo-', label='Recall')
    plt.plot(accuracy_list,'ro-',label='Accuracy')
    plt.plot(precision_list, 'yo-',label='Precision')
    plt.grid()
    plt.legend()
    plt.xlabel("Epochs")
    plt.title("Metrics: " + train_or_val)
    name = "metrics_evolution" + train_or_val + ".png"
    plt.savefig(name)
    plt.clf()

    # ROC Curves
    # Compute ROC curve and ROC area for each class
    fpr = dict()
    tpr = dict()
    roc_auc = dict()

    # Convert to numpy
    y, y_score = np.array(gt_vector), np.array(pred_vector)

    # Number of classes
    n_classes = np.array(predicted).max() + 1

    # Obtain ROC curve values and area
    for i in range(n_classes):
        fpr[i], tpr[i], _ = roc_curve(y[:, i], y_score[:, i])
        roc_auc[i] = auc(fpr[i], tpr[i])

    # Compute micro-average ROC curve and ROC area
    fpr["micro"], tpr["micro"], _ = roc_curve(y.ravel(), y_score.ravel())
    roc_auc["micro"] = auc(fpr["micro"], tpr["micro"])

    # Plot ROC curves
    # First aggregate all false positive rates
    all_fpr = np.unique(np.concatenate([fpr[i] for i in range(n_classes)]))

    # Then interpolate all ROC curves at this points
    mean_tpr = np.zeros_like(all_fpr)
    for i in range(n_classes):
        mean_tpr += np.interp(all_fpr, fpr[i], tpr[i])

    # Finally average it and compute AUC
    mean_tpr /= n_classes

    fpr["macro"] = all_fpr
    tpr["macro"] = mean_tpr
    roc_auc["macro"] = auc(fpr["macro"], tpr["macro"])

    # # Plot all ROC curves
    plt.figure()
    lw = 2
    colors = cycle(["aqua", "darkorange", "cornflowerblue"])
    for i, color in zip(range(n_classes), colors):
        plt.plot(
            fpr[i],
            tpr[i],
            color=color,
            lw=lw,
            label="ROC curve of class {0} (area = {1:0.2f})".format(i, roc_auc[i]),
        )

    plt.plot([0, 1], [0, 1], "k--", lw=lw)
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel("False Positive Rate")
    plt.ylabel("True Positive Rate")
    plt.title("Some extension of Receiver operating characteristic to multiclass")
    plt.legend(loc="lower right")
    plt.grid()
    name = "ROC_curve" + train_or_val + ".png"
    plt.savefig(name)
    plt.clf()

def right_size(CUs):
    """
    @brief Verify if the CU as the right size: height as to be lower than width

    @param [in] CUs: Feature maps
    @param [out] Boolean value indicating the right size
    """
    
    return False if CUs.shape[-2] > CUs.shape[-1] else True

def compute_conf_matrix(predicted, ground_truth):
    """!
    @brief Computes the confusion matrix

    @param [in] predicted: List of predictions made by the model with single value
    @param [in] ground_truth: List of the ground-truths with single value
    @param [out] accuracy: Accuracy score
    """
    # Plot confusion matrix
    labels = ["Non-Split", "QT", "HBT", "VBT", "HTT", "VTT"]

    # Draw confusion matrix
    c_ma = confusion_matrix(ground_truth, predicted)

    # Obtain amount of examples for each class
    amount_lst = []
    for k in range(len(c_ma)):
        
        amount_lst.append(c_ma[k].sum())

    map = sns.heatmap(c_ma / np.array([amount_lst]).T, annot=True, fmt='.2f', cmap='Blues', yticklabels=labels,
                                        xticklabels=labels)

    map.set(xlabel="Predicted", ylabel="Ground-truth")

    return map.get_figure()

def compute_top_k_accuracy(pred_vector, gt_vector, topk):
    """!
    @brief Computes the top k accuracy score

    @param [in] predicted: List of predictions made by the model with probabilities for each split (pytorch tensor)
    @param [in] ground_truth: List of the ground-truths with single value (pytorch tensor)
    @param [in] topk: Number of best accuricies to choose
    @param [out] accuracy: Accuracy score
    """
    # Initialize variables
    n_entries = gt_vector.shape[0]
    res = 0

    # Get top predictions
    top_pred = torch.topk(pred_vector, topk)
    idxs = top_pred[1]  # Indexs

    for n in range(n_entries):

        if gt_vector[n] in idxs[n]:
            res += 1

    # Compute accuracy
    res = res / n_entries  
        
    return res

def compute_num_splits_sent(pred_lst):
    """!
    @brief Computes the num of splits that would be analyzed by the encoder

    @param [in] predicted: List of predictions made by the model with probabilities values
    @param [out] res: Mean of number of splits sent
    """
    # Initialize variables
    n_entries = len(pred_lst)
    
    temp  = [len(p) for p in pred_lst]
    res = sum(temp)/n_entries
        
    return res

def compute_multi_thres_performance(pred_lst, gt_lst):
    """!
    @brief Computes multi-threshold performance

    @param [in] predicted: List of predictions made by the model with integer value
    @param [in] ground_truth: List of the ground-truths with single value 
    @param [out] res: Accuracy score
    """
    # Initialize variables
    n_entries = len(gt_lst)
    
    temp  = [True for l, v in zip(pred_lst, gt_lst) if v in l]  # If the ground truth is within the predictions, add 1 to the list
    res = len(temp)

    # Compute accuracy
    res = res / n_entries  
        
    return res

def compute_ROC_curve(pred_vector, gt_vector, pred_num):
    """!
    @brief Computes ROC curve

    @param [in] pred_vector: List of predictions vectors (one-hot encoded)
    @param [in] gt_vector: List of the ground-truths vectors (one-hot encoded)
    @param [in] pred_num: List of the predicitons with numbers corresponding to partitions
    @return [out] figure: Figure with the ROC curve
    """
    # ROC Curves
    # Compute ROC curve and ROC area for each class
    fpr = dict()
    tpr = dict()
    roc_auc = dict()

    # Convert to numpy
    y, y_score = np.array(gt_vector), np.array(pred_vector)

    # Number of classes
    n_classes = np.array(pred_num).max() + 1

    # Obtain ROC curve values and area
    for i in range(n_classes):
        fpr[i], tpr[i], _ = roc_curve(y[:, i], y_score[:, i])
        roc_auc[i] = auc(fpr[i], tpr[i])

    # Compute micro-average ROC curve and ROC area
    fpr["micro"], tpr["micro"], _ = roc_curve(y.ravel(), y_score.ravel())
    roc_auc["micro"] = auc(fpr["micro"], tpr["micro"])

    # Plot ROC curves
    # First aggregate all false positive rates
    all_fpr = np.unique(np.concatenate([fpr[i] for i in range(n_classes)]))

    # Then interpolate all ROC curves at this points
    mean_tpr = np.zeros_like(all_fpr)
    for i in range(n_classes):
        mean_tpr += np.interp(all_fpr, fpr[i], tpr[i])

    # Finally average it and compute AUC
    mean_tpr /= n_classes

    fpr["macro"] = all_fpr
    tpr["macro"] = mean_tpr
    roc_auc["macro"] = auc(fpr["macro"], tpr["macro"])

    # # Plot all ROC curves
    figure = plt.figure()
    lw = 2
    colors = cycle(["b", "y", "k", "r", "m", "g"])
    for i, color in zip(range(n_classes), colors):
        plt.plot(
            fpr[i],
            tpr[i],
            color=color,
            lw=lw,
            label="ROC curve of class {0} (area = {1:0.2f})".format(i, roc_auc[i]),
        )

    plt.plot([0, 1], [0, 1], "k--", lw=lw)
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel("False Positive Rate")
    plt.ylabel("True Positive Rate")
    plt.title("Some extension of Receiver operating characteristic to multiclass")
    plt.legend(loc="lower right")
    plt.grid()

    return figure

def model_simple_metrics(predicted, ground_truth):
    """!
    @brief Evaluates model with metrics 4 metrics, such as accuracy, f1_score, recall and precision

    @param [in] predicted: List of predictions made by the model with single value
    @param [in] ground_truth: List of the ground-truths with single value
    @param [out] f1: F1 score
    @param [out] recall: Recall score
    @param [out] precision: Precision score
    @param [out] accuracy: Accuracy score
    """

    # Compute metrics
    accuracy = accuracy_score(ground_truth, predicted)
    f1 = f1_score(ground_truth, predicted, average='weighted')
    recall = recall_score(ground_truth, predicted, average='weighted')
    precision = precision_score(ground_truth, predicted, average='weighted')

    return f1, recall, precision, accuracy

def obtain_best_modes(rs, pred):
    """!
    @brief Converts a prediction into a specific number that corresponds to the best way to split (non-split, quad tree,
           binary vert tree...)

    @param [in] rs: Thresholds
    @param [in] pred: Predicted values from the model with one-hot encoding
    @param [out] mode: Predicted values with the number of the mode
    """

    # Obtain maximum value for each prediction
    y_max = torch.reshape(torch.max(pred, dim=1)[0], shape=(-1, 1))
    # Obtain lowest possible value for each prediction
    y_max_rs = y_max * rs

    # Obtain logic truth where preds>=y_max_rs
    search_RD_logic = pred >= y_max_rs

    # Obtain the indexs that are true
    search_RD = torch.nonzero(torch.squeeze(search_RD_logic.int()), as_tuple=False)
    
    return search_RD

def obtain_mode(pred):
    """!
    @brief Converts a prediction into a specific number that corresponds to the best way to split (non-split, quad tree,
           binary vert tree...)

    @param [in] pred: Predicted values from the model with one-hot encoding
    @param [out] mode: Predicted values with the number of the mode
    """

    mode = torch.argmax(pred, dim=1)

    return mode

def one_hot_enc(tensor, num_classes=6):
    """!
    @brief Implements one-hot encoding to a specific tensor with the set of split modes

    @param [in] tensor: Tensor with a set of split modes
    @param [in] num_classes: Number classes in the tensor
    @param [out] new_tensor: Tensor with one-hot encoding implemented
    """

    new_tensor = torch.reshape(torch.nn.functional.one_hot(tensor, num_classes=num_classes), shape=(-1, num_classes))

    return new_tensor

def print_parameters(model, optimizer):
    """!
    @brief Prints the parameters from the state dictionaries of the model and optimizer

    @param [in] model: Model that the parameters will be printed
    @param [in] optimizer: Optimizer that the parameters will be printed
    """

    print()
    print()

    # Print model's state_dict
    print("Model's state_dict:")
    for param_tensor in model.state_dict():
        print(param_tensor, "\t", model.state_dict()[param_tensor].size())

    print()

    # Print optimizer's state_dict
    print("Optimizer's state_dict:")
    for var_name in optimizer.state_dict():
        print(var_name, "\t", optimizer.state_dict()[var_name])

    print()
    print()


def save_model_parameters(dir_name, f_name, model):
    """!
    @brief Saves only the model parameters to a specific folder

    @param [in] dir_name: Name of the directory where the parameters will be saved
    @param [in] f_name: Name of the file that the parameters will be saved on
    @param [in] model: Model which the parameters will be saved
    """

    # Check to see if directory already exists
    CURR_DIR = os.getcwd()  # Get current dir
    try:
        os.mkdir(CURR_DIR + '/' + dir_name)

    except:
        pass
    
    file_path = dir_name + '/' + f_name + '.pth'
    torch.save(model.state_dict(), file_path)
    print("Saved PyTorch Model State to", file_path)


def save_model(dir_name, f_name, model, optimizer, loss, acc):
    """!
    @brief Saves the parameters of the model and of the optimizer, and also the loss and the accuracy. These are saved
           into the folder specified by the user.

    @param [in] dir_name: Name of the directory where the parameters will be saved
    @param [in] f_name: Name of the file that the parameters will be saved on
    @param [in] model: Model which the parameters will be saved
    @param [in] optimizer: Optimizer which the parameters will be saved
    @param [in] loss: Loss value
    @param [in] acc: Accuracy value
    """

    # Check to see if directory already exists
    curr_dir = os.getcwd()  # Get current dir
    try:
        os.mkdir(curr_dir + '/' + dir_name)

    except:
        pass

    torch.save({
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'loss': loss,
        'acc': acc
    }, dir_name + '/' + f_name + '.tar')


def load_model_parameters_stg(model, path, stg, dev):
    """!
    @brief Loads all stages but make sure that the stage number 'stg' has the same parameters has the previous

    @param [in] model: Model which the parameters will be loaded
    @param [in] path: Path/Folder containing the files that are supposed to be loaded
    @param [in] stg: Integer containing the last stage number to load
    @param [in] dev: Device do load the model to
    @param [out] model: Model loaded tuple
    """

    # Get files names
    files = dataset_utils.get_files_from_folder(path, endswith=".pth")
    files.sort()

    # Load state dict to each model
    for i in range(len(files)):
        # Load
        file_path = path + '/' + files[i]

        if type(dev) is int:
            m = torch.load(file_path, map_location="cuda:"+str(dev))
            print("Loading model to GPU number", dev)

        else:
            m = torch.load(file_path, map_location="cpu")
            print("Loading model to CPU")
        
        model[i].load_state_dict(m)

        # To avoid inconsistent inference results
        model[i].eval()

    # Specific stage must have the same parameters of the previous stage to it
    file_path = path + '/' + files[stg-2]

    if type(dev) is int:
        m = torch.load(file_path, map_location="cuda:"+str(dev))
        print("Loading model to GPU number", dev)

    else:
        m = torch.load(file_path, map_location="cpu")
        print("Loading model to CPU")
    
    model[stg-1].load_state_dict(m)

    # To avoid inconsistent inference results
    model[stg-1].eval()

    return model

def load_model_parameters_eval(model, path, dev):
    """!
    @brief Loads all stages, meant to be used with the eval_model script

    @param [in] model: Model which the parameters will be loaded
    @param [in] path: Path/Folder containing the files that are supposed to be loaded
    @param [in] dev: Device do load the model to
    @param [out] model: Model loaded tuple
    """

    # Get files names
    files = dataset_utils.get_files_from_folder(path, endswith=".pth")
    files.sort()

    # Load state dict to each model
    for i in range(len(files)):
        # Load
        file_path = path + '/' + files[i]

        if type(dev) is int:
            m = torch.load(file_path, map_location="cuda:"+str(dev))
            print("Loading model to GPU number", dev)

        else:
            m = torch.load(file_path, map_location="cpu")
            print("Loading model to CPU")
        
        model[i].load_state_dict(m)

        # To avoid inconsistent inference results
        model[i].eval()

    return model

def load_model_stg_12_stg_3(model, path, dev):
    """!
    @brief THis function makes it possible to load parameters from the first and second stage to the third

    @param [in] model: Model which the parameters will be loaded, with 2 models (one for the first and second stage, and another for the third stage)
    @param [in] path: Path/Folder containing the files that are supposed to be loaded
    @param [in] dev: Device to load the model to
    """

    # Get files names
    files = dataset_utils.get_files_from_folder(path, endswith=".pth")
    files.sort()
    
    # Load state dict to each model
    file_path = path + "/" + files[0]
    stg_1_2 = torch.load(file_path, map_location="cuda:"+str(dev))
    if type(dev) is int:
        stg_1_2 = torch.load(file_path, map_location="cuda:"+str(dev))
        print("Loading model to GPU number", dev)

    else:
        stg_1_2 = torch.load(file_path, map_location="cpu")
        print("Loading model to CPU")

    model[0].load_state_dict(stg_1_2)

    # Load specific layers
    with torch.no_grad():

        # Conditional convolution
        model[1].simple_conv[0].weight.copy_(stg_1_2["simple_conv_stg2.0.weight"])
        model[1].simple_conv[0].bias.copy_(stg_1_2["simple_conv_stg2.0.bias"])
        model[1].simple_conv[1].weight.copy_(stg_1_2["simple_conv_stg2.1.weight"])

        model[1].simple_conv_no_activation[0].weight.copy_(stg_1_2["simple_conv_no_activation_stg2.0.weight"])
        model[1].simple_conv_no_activation[0].bias.copy_(stg_1_2["simple_conv_no_activation_stg2.0.bias"])

        model[1].simple_conv2[0].weight.copy_(stg_1_2["simple_conv2_stg2.0.weight"])
        model[1].simple_conv2[0].bias.copy_(stg_1_2["simple_conv2_stg2.0.bias"])
        model[1].simple_conv2[1].weight.copy_(stg_1_2["simple_conv2_stg2.1.weight"])

        model[1].simple_conv_no_activation2[0].weight.copy_(stg_1_2["simple_conv_no_activation2_stg2.0.weight"])
        model[1].simple_conv_no_activation2[0].bias.copy_(stg_1_2["simple_conv_no_activation2_stg2.0.bias"])
        
        model[1].activation_PRelu.weight.copy_(stg_1_2["activation_PRelu_stg2.weight"])
        model[1].activation_PRelu2.weight.copy_(stg_1_2["activation_PRelu2_stg2.weight"])        

    # To avoid inconsistent inference results
    for m in model:

        m.eval()

    return model

def load_model_stg_3_stg_4(model, path, dev):
    """!
    @brief This function makes it possible to load parameters from the third stage to the fourth

    @param [in] model: Model which the parameters will be loaded, with 2 models (one for the first and second stage, and another for the third stage)
    @param [in] path: Path/Folder containing the files that are supposed to be loaded
    @param [in] dev: Device to load the model to
    """

    # Get files names
    files = dataset_utils.get_files_from_folder(path, endswith=".pth")

    # Load Stage 1 and 2
    for k in range(len(files)):

        if '0' in files[k]:
            break
    
    # Load state dict to each model
    file_path = path + "/" + files[k]
    stg_1_2 = torch.load(file_path, map_location="cuda:"+str(dev))
    if type(dev) is int:
        stg_1_2 = torch.load(file_path, map_location="cuda:"+str(dev))
        print("Loading model to GPU number", dev)

    else:
        stg_1_2 = torch.load(file_path, map_location="cpu")
        print("Loading model to CPU")

    model[0].load_state_dict(stg_1_2)

    # Load Stage 3
    for k in range(len(files)):

        if '1' in files[k]:
            break
    
    # Load state dict to each model
    file_path = path + "/" + files[k]

    stg_3 = torch.load(file_path, map_location="cuda:"+str(dev))
    if type(dev) is int:
        stg_3 = torch.load(file_path, map_location="cuda:"+str(dev))
        print("Loading model to GPU number", dev)

    else:
        stg_3 = torch.load(file_path, map_location="cpu")
        print("Loading model to CPU")

    model[-2].load_state_dict(stg_3)

    model[-1].load_state_dict(stg_3)

    # To avoid inconsistent inference results
    for m in model:

        m.eval()

    return model


def load_model_stg_4_stg_5(model, path, dev):
    """!
    @brief This function makes it possible to load parameters from the fourth stage to the fith

    @param [in] model: Model which the parameters will be loaded, with 2 models (one for the first and second stage, and another for the third stage)
    @param [in] path: Path/Folder containing the files that are supposed to be loaded
    @param [in] dev: Device to load the model to
    """

    # Get files names
    files = dataset_utils.get_files_from_folder(path, endswith=".pth")

    # Load Stage 1 and 2
    for k in range(len(files)):
        if '0' in files[k]:
            break
    
    # Load state dict to each model
    file_path = path + "/" + files[k]
    stg_1_2 = torch.load(file_path, map_location="cuda:"+str(dev))
    if type(dev) is int:
        stg_1_2 = torch.load(file_path, map_location="cuda:"+str(dev))
        print("Loading model to GPU number", dev)

    else:
        stg_1_2 = torch.load(file_path, map_location="cpu")
        print("Loading model to CPU")

    model[0].load_state_dict(stg_1_2)

    # Load Stage 3
    for k in range(len(files)):
        if '1' in files[k]:
            break
    
    # Load state dict to each model
    file_path = path + "/" + files[k]

    stg_3 = torch.load(file_path, map_location="cuda:"+str(dev))
    if type(dev) is int:
        stg_3 = torch.load(file_path, map_location="cuda:"+str(dev))
        print("Loading model to GPU number", dev)

    else:
        stg_3 = torch.load(file_path, map_location="cpu")
        print("Loading model to CPU")

    model[1].load_state_dict(stg_3)


    # Load Stage 4
    for k in range(len(files)):

        if '2' in files[k]:
            break
    
    # Load state dict to each model
    file_path = path + "/" + files[k]

    stg_4 = torch.load(file_path, map_location="cuda:"+str(dev))
    if type(dev) is int:
        stg_4 = torch.load(file_path, map_location="cuda:"+str(dev))
        print("Loading model to GPU number", dev)

    else:
        stg_4 = torch.load(file_path, map_location="cpu")
        print("Loading model to CPU")

    model[2].load_state_dict(stg_4)

    # Load stage 5
    model[3].load_state_dict(stg_4)

    # To avoid inconsistent inference results
    for m in model:
        m.eval()

    return model


def load_model_stg_5_stg_6(model, path, dev):
    """!
    @brief This function makes it possible to load parameters from the fourth stage to the fith

    @param [in] model: Model which the parameters will be loaded, with 2 models (one for the first and second stage, and another for the third stage)
    @param [in] path: Path/Folder containing the files that are supposed to be loaded
    @param [in] dev: Device to load the model to
    """

    # Get files names
    files = dataset_utils.get_files_from_folder(path, endswith=".pth")

    # Load Stage 1 and 2
    for k in range(len(files)):
        if '0' in files[k]:
            break
    
    # Load state dict to each model
    file_path = path + "/" + files[k]
    stg_1_2 = torch.load(file_path, map_location="cuda:"+str(dev))
    if type(dev) is int:
        stg_1_2 = torch.load(file_path, map_location="cuda:"+str(dev))
        print("Loading model to GPU number", dev)

    else:
        stg_1_2 = torch.load(file_path, map_location="cpu")
        print("Loading model to CPU")

    model[0].load_state_dict(stg_1_2)

    # Load Stage 3
    for k in range(len(files)):
        if '1' in files[k]:
            break
    
    # Load state dict to each model
    file_path = path + "/" + files[k]

    stg_3 = torch.load(file_path, map_location="cuda:"+str(dev))
    if type(dev) is int:
        stg_3 = torch.load(file_path, map_location="cuda:"+str(dev))
        print("Loading model to GPU number", dev)

    else:
        stg_3 = torch.load(file_path, map_location="cpu")
        print("Loading model to CPU")

    model[1].load_state_dict(stg_3)


    # Load Stage 4
    for k in range(len(files)):
        if '2' in files[k]:
            break
    
    # Load state dict to each model
    file_path = path + "/" + files[k]

    stg_4 = torch.load(file_path, map_location="cuda:"+str(dev))
    if type(dev) is int:
        stg_4 = torch.load(file_path, map_location="cuda:"+str(dev))
        print("Loading model to GPU number", dev)

    else:
        stg_4 = torch.load(file_path, map_location="cpu")
        print("Loading model to CPU")

    model[2].load_state_dict(stg_4)

    # Load Stage 5
    for k in range(len(files)):
        if '3' in files[k]:
            break
    
    # Load state dict to each model
    file_path = path + "/" + files[k]

    stg_5 = torch.load(file_path, map_location="cuda:"+str(dev))
    if type(dev) is int:
        stg_5 = torch.load(file_path, map_location="cuda:"+str(dev))
        print("Loading model to GPU number", dev)

    else:
        stg_5 = torch.load(file_path, map_location="cpu")
        print("Loading model to CPU")

    model[3].load_state_dict(stg_5)

    # Load stage 6
    model[4].load_state_dict(stg_5)

    # To avoid inconsistent inference results
    for m in model:
        m.eval()

    return model

def print_current_time():
    """!
    @brief Prints current time
    """
    now = datetime.now()
    current_time = now.strftime("%H:%M:%S")
    print("Current Time =", current_time)