test.py

import argparse
import json
import os
import random
import yaml
import numpy as np
import torch
from torch.backends import cudnn
from torch.utils.data import DataLoader
from tqdm import tqdm

from Dataset.StaticDataset import StaticDataset
from utils.ConfusionMatrix import BoundingBoxConfusionMatrix
from utils.dataset_utils import DatasetType
from utils.utils import create_test_folder, postprocess, calc_map, \
    collect_metrics, random_color, state_2_csv, \
    save_config, testing_load_checkpoint, test_results_2_csv
from utils.vis_utils import initialize_visualizations, visualize, save_confusion_matrix, save_run_plots, \
    save_map_bar_plot, save_recall_plot, save_img_predictions, save_precision_plot,\
    save_map_plot

def parse_args():
    """
    Collects arguments passed to the train.py.

    :return: An arguments object for us in rest of code.
    :rtype: argparse.Namespace
    """
    parser = argparse.ArgumentParser()

    parser.add_argument('--config_file',
                        default="configs/config-sample-test.json",
                        type=str,
                        metavar='N',
                        help='config file for training and generation')

    args = parser.parse_args()
    return args

def seed_and_settings(args=None):
    """
    initializes seeds and settings.

    :param args: arguments for the training.
    :type args: argparse.Namespace
    :return: device to use for training, running configuration, initialized state dict, initialized plot dict.
    :rtype: list
    """

    os.environ["NO_ALBUMENTATIONS_UPDATE"] = "1"

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    with open(args.config_file, 'r') as f_config_json:
        config = json.load(f_config_json)

    config["config_path"] = args.config_file
    config["number_classes"] = len(config["generation"]["object_label"].values())
    config["training"] = {}
    config["training"]["nr_dataset_partitions"] = 4
    # seed everything

    seed_value = config["test"]["seed"]

    if seed_value != -1:
        random.seed(seed_value)
        np.random.seed(seed_value)
        torch.manual_seed(seed_value)
        os.environ['PYTHONHASHSEED'] = str(seed_value)

    if torch.cuda.is_available():
        torch.cuda.set_device(config["test"]["gpu_ordinal_for_testing"])
        cudnn.enabled = True
        cudnn.benchmark = True

        if seed_value != -1:
            torch.cuda.manual_seed(seed_value)
            torch.cuda.manual_seed_all(seed_value)

    plot_dict = initialize_visualizations(config)

    return device, config, plot_dict

def initialize_dataset(config):
    """
    Initializes the dataset and dataloader used for testing.

    :param config: Running configuration
    :type config: dict
    :return: Dataloader used for testing
    :rtype: DataLoader
    """


    with open(config["test"]["dataset_yaml"], "r", encoding="utf-8") as yaml_file:
        dataset_yaml = yaml.safe_load(yaml_file)

    test_path = os.path.join(dataset_yaml["path"].replace("../", ""), dataset_yaml["test"])

    test_dataset = StaticDataset(img_path=test_path,
                                 img_size=config["test"]["img_size"],
                                 dataset_type=DatasetType.TEST,
                                 augmentation_probabilities=None
                                 )

    test_dataloader = DataLoader(dataset=test_dataset,
                              batch_size=config["test"]["batch_size"],
                              num_workers=4,
                              collate_fn=getattr(test_dataset, "collate_fn", None),
                              persistent_workers=True,
                              pin_memory=True
                              )

    config["label_color"] = {}

    for obj_name in dataset_yaml["names"].values():
        config["label_color"][obj_name] = random_color()

    return test_dataloader

def save_plots(state_dict, saved_config, config):
    """
    Saves plots related to the testing and other results from during the model training. Representing both results
    from the current testing and results from during the previous model training.

    :param state_dict: Current state of the model.
    :type state_dict: dict
    :param saved_config: Saved config from the training.
    :type saved_config: dict
    :param config: Current config used for testing.
    :type config: dict
    :return: None
    :rtype: None
    """
    state_2_csv(state_dict, config["test_path"], config)
    save_confusion_matrix(state_dict["test_conf_matrix"], config["test_path"], config)
    save_run_plots(state_dict, config["test_path"], config)
    save_map_bar_plot(state_dict, config["test_path"], config)
    save_map_plot(state_dict["eval_metrics"], "Test Object mAP@50", config["test_path"], config)
    save_recall_plot(state_dict, config["test_path"], config)
    save_precision_plot(state_dict, config["test_path"], config)
    test_results_2_csv(state_dict, config)
    save_config(saved_config, config["test_path"])

def main():

    # Parse input and setup testing settings.
    args = parse_args()
    device, config, plot_dict = seed_and_settings(args=args)

    if config["test"]["weights"] == "" or config["test"]["weights"] == None:
        raise TypeError("Model weights are missing! Please specify model path to run the testing")

    # Load model for testing.
    model, state_dict, config, saved_config = testing_load_checkpoint(config, device)

    # Initialize testing dataset.
    test_dataloader = initialize_dataset(config)

    # Create testing folder for saving files.
    config = create_test_folder(config)
    model.eval()

    confusion_matrix = BoundingBoxConfusionMatrix(config["number_classes"],
                                                  config["evaluation"]["confusion_matrix"]["iou_thres"],
                                                  config["evaluation"]["confusion_matrix"]["conf_thres"])

    preds_list = []
    targets_list = []
    state_dict["eval_metrics"] = state_dict["val_metrics"]

    # Run test on the test dataset.
    with torch.no_grad():
        for batch in tqdm(test_dataloader):

            # Run test images through model
            batch["img"] = batch["img"].to(device)
            preds = model(batch["img"])
            preds = postprocess(preds, config["test"]["postprocess_conf_thres"], config["test"]["postprocess_iou_thres"])

            # Collect testing metrics
            batch_preds, batch_targets = collect_metrics(preds, batch, config["test"]["img_size"])
            preds_list.extend(batch_preds)
            targets_list.extend(batch_targets)


        save_img_predictions(preds_list, targets_list, config["test_path"], config)
        confusion_matrix.update_confusion_matrix(preds_list, targets_list)
        state_dict["test_conf_matrix"] = confusion_matrix.confusion_matrix

        state_dict["current_test_metrics"] = calc_map(preds_list, targets_list, config["evaluation"]["map_conf_thres"])

        # Visualize testing results
        visualize(state_dict, plot_dict, show=False)
        # Save plots from testing
        save_plots(state_dict, saved_config, config)

if __name__ == '__main__':

    main()