@ptrblck thank you for the probability map. You helped me before for the probability map of output, and it’s done. I wanted to visualise the prediction results from the test_loader, the mask (ground truth), and an original input image. I am wondering why my original input image doesn’t look grayscale image. Am I doing something wrong?

     #test model
    model.eval()
    total = 0
    test_loss = 0
    correct = 0
    count = 0
    #iterate through test dataset
    for vi, data in enumerate(test_loader):
                
                t_image, mask = data
                #print(t_image.shape) # torch.Size([1, 1, 240, 320])
                t_image, mask = t_image.to(device), mask.to(device)
                with torch.no_grad():

                    outputs = model(t_image)
                    #print(outputs.shape) # torch.Size([1, 2, 240, 320])
                    test_loss += criterion(outputs, mask).item() / len(test_loader)

                    pr = torch.exp(outputs) # get the exp of output and will give the probability map of outpus
                    # The outputs are energies for the 2 classes.
                    # Higher the energy for a class, the more the network thinks that the image is of the particular class. So, let’s get the index of the highest energy:     
                    _, predicted = torch.max(outputs.data, 1)  
                    
                    total += mask.size(0)
                    correct += predicted.eq(mask.data).sum().item()
                    accuracy = 100 * correct / total
                    predict = predicted.eq(mask.data)
                    #print(predict.shape) # torch.Size([1, 240, 320])
                    count +=1                
                    print(count, "Test Loss: {:.3f}".format(test_loss), "Test Accuracy: %d %%" % (accuracy))

    
    plt.figure(figsize=(11,4))
    plt.subplot(1, 3, 1)
    plt.imshow(predict.detach().cpu().squeeze())
    plt.title('DL prediction')
    plt.subplot(1, 3, 2)
    plt.imshow(mask.detach().cpu().squeeze())
    plt.title('ground truth')
    plt.subplot(1, 3, 3)
    #print(t_image.size()) # torch.Size([1, 1, 240, 320])
    plt.imshow(t_image.detach().cpu().squeeze())
    plt.title('original input image')