DengPingFan / Polyp-PVT

Polyp-PVT: Polyp Segmentation with Pyramid Vision Transformers, AIR 2023.
196 stars 45 forks source link

evaluation #13

Open LvGuangzu opened 5 months ago

LvGuangzu commented 5 months ago

Hello, I see that you only have the mdice evaluation indicator in Train.py. Could you please share the code for other evaluation indicators during the training of the model?

def test(model, path, dataset):

data_path = os.path.join(path, dataset)
image_root = '{}/images/'.format(data_path)
gt_root = '{}/masks/'.format(data_path)
model.eval()
num1 = len(os.listdir(gt_root))
test_loader = test_dataset(image_root, gt_root, 352)
DSC = 0.0
for i in range(num1):
    image, gt, name = test_loader.load_data()
    gt = np.asarray(gt, np.float32)
    gt /= (gt.max() + 1e-8)
    image = image.cuda()

    res, res1  = model(image)
    # eval Dice
    res = F.upsample(res + res1 , size=gt.shape, mode='bilinear', align_corners=False)
    res = res.sigmoid().data.cpu().numpy().squeeze()
    res = (res - res.min()) / (res.max() - res.min() + 1e-8)
    input = res
    target = np.array(gt)
    N = gt.shape
    smooth = 1
    input_flat = np.reshape(input, (-1))
    target_flat = np.reshape(target, (-1))
    intersection = (input_flat * target_flat)
    dice = (2 * intersection.sum() + smooth) / (input.sum() + target.sum() + smooth)
    dice = '{:.4f}'.format(dice)
    dice = float(dice)
    DSC = DSC + dice

return DSC / num1