File size: 2,699 Bytes
a256709 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 |
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)
# initialize the ground truth and output tensor
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)
|