test.py

import os
import json
from pathlib import Path

from torch.utils.data import DataLoader

from utils.evaluation_runner import evaluate, _build_test_datasets, load_trained_model
from utils.results_sorted import *
from Configs import Config
config = Config()

def main():
    config.require_data()
    config.prepare_output_dirs()
    results = {}
    datasets = _build_test_datasets(config)
    for class_name in config.data_name_list:
        results[class_name] = {}
        dataset = datasets[class_name]
        data_loader = DataLoader(dataset, batch_size=config.batch_size, shuffle=False, num_workers=config.num_workers)
        seghead_names = os.listdir(os.path.join(config.seghead_dir, class_name))
        seghead_names.sort(key=lambda x: int(x.split('.')[0]))
        seghead_paths = [os.path.join(config.seghead_dir, class_name, weight_name) for weight_name in seghead_names]
        for seghead_path in seghead_paths[-config.val_num:]:
            iteration = int(Path(seghead_path).stem)
            results[class_name][iteration] = {}
            encoder, fdp, seghead = load_trained_model(config, seghead_path, class_name)
            I_AUROC, P_AUROC, P_AP, P_PRO = evaluate(
                config, data_loader, encoder, fdp, seghead, class_name,
                checkpoint_name=Path(seghead_path).stem,
            )
            results[class_name][iteration]["I_AUROC"] = I_AUROC; results[class_name][iteration]["P_AUROC"] = P_AUROC
            results[class_name][iteration]["P_AP"] = P_AP; results[class_name][iteration]["P_PRO"] = P_PRO
            print(f"{config.mode}-class setting val on {class_name}, iteration--{iteration}: "
                  f"I-AUROC:{I_AUROC}, P-AUROC:{P_AUROC}, P-AP:{P_AP}, P-PRO:{P_PRO}")
    if not os.path.exists(os.path.dirname(config.save_result_path)):
        os.makedirs(os.path.dirname(config.save_result_path))
    with open(config.save_result_path, 'w', encoding='utf-8') as f:
        json.dump(results, f, indent=4, ensure_ascii=False)
if __name__ == '__main__':
    main()