eval.py

import argparse
import logging
import os
import random
import sys
import numpy as np
import torch
import torch.backends.cudnn as cudnn
import torch.nn as nn
from torch.utils.data import DataLoader
from tqdm import tqdm
from scipy.ndimage import zoom
from datasets.dataset_us_xray import Ultrasound_dataset, LungXray_dataset
from utils import calculate_metric_percase
from networks.TransUNet_model import TransUNet

cudnn.benchmark = False
cudnn.deterministic = True
random.seed(2021)
np.random.seed(2021)
torch.manual_seed(2021)
torch.cuda.manual_seed(2021)

parser = argparse.ArgumentParser()
parser.add_argument('--root_path', type=str,
                    default='', help='root dir for validation volume data')  
parser.add_argument('--dataset', type=str,
                    default='Ultrasound', help='name of dataset for training')
parser.add_argument('--num_classes', type=int,
                    default=2, help='number of classes including background')
parser.add_argument('--list_dir', type=str, default='./lists/lists_Ultrasound', help='path to dir where train-val split is stored')
parser.add_argument('--model_path', type=str,
                    required=True, help='path to trained model')
parser.add_argument('--img_size', type=int, default=224, help='input patch size of network input')
args = parser.parse_args()

def test_single_volume(image, label, net, classes, patch_size=[256, 256], case=None):
    image, label = image.squeeze(0).cpu().detach().numpy(), label.squeeze(0).cpu().detach().numpy()
    if len(image.shape) == 3:
        prediction = np.zeros_like(label)
        for ind in range(image.shape[0]):
            slice = image[ind, :, :]
            x, y = slice.shape[0], slice.shape[1]
            if x != patch_size[0] or y != patch_size[1]:
                slice = zoom(slice, (patch_size[0] / x, patch_size[1] / y), order=3)
            input = torch.from_numpy(slice).unsqueeze(0).unsqueeze(0).float().cuda()
            net.eval()
            with torch.no_grad():
                outputs = net(input)
                out = torch.argmax(torch.softmax(outputs, dim=1), dim=1).squeeze(0)
                out = out.cpu().detach().numpy()
                if x != patch_size[0] or y != patch_size[1]:
                    pred = zoom(out, (x / patch_size[0], y / patch_size[1]), order=0)
                else:
                    pred = out
                prediction[ind] = pred
    else:
        input = torch.from_numpy(image).unsqueeze(
            0).unsqueeze(0).float().cuda()
        net.eval()
        with torch.no_grad():
            out = torch.argmax(torch.softmax(net(input), dim=1), dim=1).squeeze(0)
            prediction = out.cpu().detach().numpy()
    metric_list = []
    for i in range(1, classes):
        metric_list.append(calculate_metric_percase(prediction == i, label == i))

    return metric_list


def inference(args, model):
    db_test = args.Dataset(base_dir=args.root_path, split="test_vol", list_dir=args.list_dir)
    testloader = DataLoader(db_test, batch_size=1, shuffle=False, num_workers=1)
    logging.info("{} test iterations per epoch".format(len(testloader)))
    model.eval()
    metric_list = 0.0
    for i_batch, sampled_batch in tqdm(enumerate(testloader)):
        h, w = sampled_batch["image"].size()[2:]
        image, label, case_name = sampled_batch["image"], sampled_batch["label"], sampled_batch['case_name'][0]
        metric_i = test_single_volume(image, label, model, classes=args.num_classes, patch_size=[args.img_size, args.img_size],
                                      case=case_name)
        metric_list += np.array(metric_i)
        logging.info('idx %d case %s mean_dice %f mean_hd95 %f' % (i_batch, case_name, np.mean(metric_i, axis=0)[0], np.mean(metric_i, axis=0)[1]))
    metric_list = metric_list / len(db_test)
    for i in range(1, args.num_classes):
        logging.info('Mean class %d mean_dice %f mean_hd95 %f' % (i, metric_list[i-1][0], metric_list[i-1][1]))
    performance = np.mean(metric_list, axis=0)[0]
    mean_hd95 = np.mean(metric_list, axis=0)[1]
    logging.info('Testing performance in best val model: mean_dice : %f mean_hd95 : %f' % (performance, mean_hd95))
    return "Testing Finished!"

if __name__ == "__main__":

    dataset_config = {
        'Ultrasound': {
            'Dataset': Ultrasound_dataset,
            'root_path': '/ssd_scratch/cvit/rupraze/data/ultrasound',
            'list_dir': './lists/lists_Ultrasound',
            'num_classes': 2,
        },
        'LungSeg': {
            'Dataset': LungXray_dataset,
            'root_path': '/ssd_scratch/cvit/rupraze/data/lungs_seg_dataset',
            'list_dir': './lists/lists_CovidLungSeg',
            'num_classes': 2,
        },
    }
    dataset_name = args.dataset
    args.num_classes = dataset_config[dataset_name]['num_classes']
    args.root_path = dataset_config[dataset_name]['root_path']
    args.Dataset = dataset_config[dataset_name]['Dataset']
    args.list_dir = dataset_config[dataset_name]['list_dir']

    args.exp = dataset_name + str(args.img_size)

    net = TransUNet(num_classes=args.num_classes).cuda()
    if not os.path.exists(args.model_path):
        print("Model path doesn't exist \nExiting eval process.")
        exit()
    net.load_state_dict(torch.load(args.model_path))
    snapshot_name = dataset_name + ' \t ' + os.path.splitext(os.path.basename(args.model_path))[0]

    log_folder = './test_log/test_log_' + args.exp
    os.makedirs(log_folder, exist_ok=True)
    logging.basicConfig(filename=log_folder + '/'+snapshot_name+".txt", level=logging.INFO, format='[%(asctime)s.%(msecs)03d] %(message)s', datefmt='%H:%M:%S')
    logging.getLogger().addHandler(logging.StreamHandler(sys.stdout))
    logging.info(str(args))
    logging.info(snapshot_name)

    inference(args, net)