import os

import torch

from torch.utils.data import DataLoader

from torchvision import transforms

from loss import Loss,cal_dice

from dataloaders.la_heart import LAHeart, CenterCrop, ToTensor

from networks.vnet import VNet

def eval_loop(model, criterion, valid_loader, device):

    model.eval()

    running_loss = 0

    dice_valid = 0

    with torch.no_grad():

        for sampled_batch in valid_loader:

            volume_batch, label_batch = sampled_batch['image'], sampled_batch['label']

            volume_batch, label_batch = volume_batch.to(device), label_batch.to(device)

            outputs = model(volume_batch)

            loss = criterion(outputs, label_batch)

            dice = cal_dice(outputs, label_batch)

            print('dice: {:.3f}'.format(dice))

            running_loss += loss.item()

            dice_valid += dice.item()

    loss = running_loss / len(valid_loader)

    dice = dice_valid / len(valid_loader)

    return {'loss': loss, 'dice': dice}

os.environ['CUDA_VISIBLE_DEVICES'] = '0'

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

data_path = '/***、data_set/LASet/data'

patch_size = (112,112,80)

model = VNet(n_channels=1,n_classes=2, normalization='batchnorm').to(device)

# 加载训练模型

weight_path = 'results/VNet.pth'

weight_dict = torch.load(weight_path, map_location=device)

model.load_state_dict(weight_dict)

print('Successfully loading checkpoint.')

criterion = Loss(n_classes=2).to(device)

db_test = LAHeart(base_dir=data_path,

                        split='test',

                        transform=transforms.Compose([

                        CenterCrop(patch_size),

                        ToTensor()

                    ]))

testloader = DataLoader(db_test,batch_size=1, num_workers=4, pin_memory=True)

model.eval()

valid_metrics = eval_loop(model, criterion, testloader, device)

# 这里的dice是测试集中心裁剪的dice

dice = valid_metrics['dice']

print('Average dice: {:.5f}'.format(dice))