|
import argparse |
|
import os |
|
import ruamel_yaml as yaml |
|
import torch |
|
import torch.nn as nn |
|
import torch.nn.functional as F |
|
from torch.utils.data import DataLoader |
|
|
|
from models.resunet import ModelResUNet_ft |
|
|
|
from dataset.dataset_siim_acr import SIIM_ACR_Dataset |
|
from metric import mIoU, dice |
|
|
|
|
|
def test(args, config): |
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
print("Total CUDA devices: ", torch.cuda.device_count()) |
|
torch.set_default_tensor_type("torch.FloatTensor") |
|
|
|
test_dataset = SIIM_ACR_Dataset(config["test_file"], is_train=False) |
|
test_dataloader = DataLoader( |
|
test_dataset, |
|
batch_size=config["test_batch_size"], |
|
num_workers=4, |
|
pin_memory=True, |
|
sampler=None, |
|
shuffle=True, |
|
collate_fn=None, |
|
drop_last=True, |
|
) |
|
|
|
model = ModelResUNet_ft( |
|
res_base_model="resnet50", out_size=1, imagenet_pretrain=False |
|
) |
|
model = nn.DataParallel( |
|
model, device_ids=[i for i in range(torch.cuda.device_count())] |
|
) |
|
model = model.to(device) |
|
|
|
print("Load model from checkpoint:", args.model_path) |
|
checkpoint = torch.load(args.model_path, map_location="cpu") |
|
state_dict = checkpoint["model"] |
|
model.load_state_dict(state_dict) |
|
|
|
|
|
gt = torch.FloatTensor() |
|
gt = gt.cuda() |
|
pred = torch.FloatTensor() |
|
pred = pred.cuda() |
|
|
|
print("Start testing") |
|
model.eval() |
|
for i, sample in enumerate(test_dataloader): |
|
image = sample["image"] |
|
mask = sample["seg"].float().to(device) |
|
gt = torch.cat((gt, mask), 0) |
|
input_image = image.to(device, non_blocking=True) |
|
with torch.no_grad(): |
|
pred_mask = model(input_image) |
|
pred_mask = F.sigmoid(pred_mask) |
|
pred = torch.cat((pred, pred_mask), 0) |
|
dice_score, dice_neg, dice_pos, num_neg, num_pos = dice(pred, gt) |
|
IoU_score = mIoU(pred, gt) |
|
print("Dice score is", dice_score) |
|
print("IoU score is", IoU_score) |
|
return dice_score, IoU_score |
|
|
|
|
|
if __name__ == "__main__": |
|
parser = argparse.ArgumentParser() |
|
parser.add_argument("--config", default="Path/To/Res_train.yaml") |
|
parser.add_argument("--checkpoint", default="") |
|
parser.add_argument("--model_path", default="Path/To/best_valid.pth") |
|
parser.add_argument("--device", default="cuda") |
|
parser.add_argument("--gpu", type=str, default="0", help="gpu") |
|
args = parser.parse_args() |
|
|
|
config = yaml.load(open(args.config, "r"), Loader=yaml.Loader) |
|
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu |
|
torch.cuda.current_device() |
|
torch.cuda._initialized = True |
|
|
|
test(args, config) |
|
|
