stefanknegt / Probabilistic-Unet-Pytorch

A Probabilistic U-Net for segmentation of ambiguous images implemented in PyTorch
Apache License 2.0
270 stars 67 forks source link

Testing #5

Closed Mikulano closed 5 years ago

Mikulano commented 5 years ago

Hello! Thank you for your code. I don't understand how to test it after training. Maybe do you have some examples of code for testing or recommendations?

stefanknegt commented 5 years ago

I have used the following in the training file. You can then write the _generalized_energydistance function (as explained in the original paper) to properly evaluate the network. In addition, you need to supply all four masks from the dataloader. Good luck!

            iou_score = 0
            ged_score = 0
            for val_step, (patch, _, _, masks, _) in enumerate(val_loader):
                masks = torch.squeeze(masks,0)
                patch = patch.to(device)
                net.forward(patch, segm=None)
                num_preds = 4
                predictions = []
                for i in range(num_preds):
                    mask_pred = net.sample(testing=True)
                    mask_pred = (torch.sigmoid(mask_pred) > 0.5).float()
                    mask_pred = torch.squeeze(mask_pred, 0)
                    predictions.append(mask_pred)
                predictions = torch.cat(predictions, 0)

                iou_score_iter, ged_score_iter = generalized_energy_distance_iou(predictions, masks)
                iou_score += iou_score_iter
                ged_score += ged_score_iter

            ged = ged_score/len(val_dataset)
            iou = iou_score/len(val_dataset)
Mikulano commented 5 years ago

Thank you!