# How can I display a test image and display the mask for it based on my trained model?

**URL:** <https://discuss.pytorch.org/t/how-can-i-display-a-test-image-and-display-the-mask-for-it-based-on-my-trained-model/33016>\
**Category:** vision\
**Created:** [December 25, 2018, 12:41pm UTC](https://discuss.pytorch.org/t/how-can-i-display-a-test-image-and-display-the-mask-for-it-based-on-my-trained-model/33016 "2018-12-25T12:41:36Z")\
**Posts on this page:** 20\
**Page:** 1

<div class="post-metadata">

**Author:** ![Neda](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/neda/32/7779_2.png) [@Neda](https://discuss.pytorch.org/u/Neda)\
**Post date:** [December 25, 2018, 12:41pm UTC](https://discuss.pytorch.org/t/how-can-i-display-a-test-image-and-display-the-mask-for-it-based-on-my-trained-model/33016/1 "2018-12-25T12:41:36Z")

</div>

I am wondering how I can test the trained model for semantic segmentation and visualise the mask for the test image. There is an example for classification problem in Pytorch but couldn’t find any obvious example for the segmentation.

I found [this page](https://medium.com/@tsakunelsonz/loading-and-training-a-neural-network-with-custom-dataset-via-transfer-learning-in-pytorch-8e672933469) that test the network, but it’s for classification problem.

I did manipulate it for segmentation application like below but now sure am I doing it right?

```auto
  model.eval()
    total = 0
    correct = 0
    count = 0
    #iterate through test dataset
    for data in (test_loader):
                
                t_image, mask = data
                t_image, mask = t_image.to(device), mask.to(device)
                with torch.no_grad():

                    outputs = model.forward(t_image) 
                    ps = torch.exp(outputs)
                         
                    _, predicted = torch.max(outputs.data, 1) # Find the class index with the maximum value.
                    # We can use the PyTorch.eq() function to do this, which compares the values in two tensors and if they match, returns a 1. If they don’t match, it returns a 0
                    # By summing the output of the .eq() function, we get a count of the number of times the neural network has produced a correct output,
                    # and we take an accumulating sum of these correct predictions so that we can determine the overall accuracy of the network on our test data set.
                    total += mask.size(0)
                    correct += predicted.eq(mask.data).sum().item()
                    count +=1                
                    print("Accuracy of network on test images is ... {:.4f}....count: {}".format(100*correct/total, count ))

```

the first question is why it used only forward here `outputs = model.forward(t_image)`?

And the second question is how can I visualise the output of the test, for example, how can I display a test image and drow the mask for it based on my trained model? Is there any example for this? Thank you in advance.

---

<div class="post-metadata">

**Author:** ![tom](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/tom/32/3162_2.png) [@tom](https://discuss.pytorch.org/u/tom)\
**Post date:** [December 25, 2018, 6:38pm UTC](https://discuss.pytorch.org/t/how-can-i-display-a-test-image-and-display-the-mask-for-it-based-on-my-trained-model/33016/2 "2018-12-25T18:38:10Z")

</div>

> [@Neda](#):
>
> the first question is why it used only forward here `outputs = model.forward(t_image)` ?

Because they haven’t asked here to be told that you should use `output = model(t_image)`. 😉  
More seriously, you would not call `.forward` unless you have a very good reason to not call the model itself (i.e. you tried that, it failed, you know why it failed and why calling `.forward` is better and, personally, I’d ask here about why calling the model won’t work).

> [@](#):
>
> And the second question is how can I visualise the output of the test, for example, how can I display a test image and drow the mask for it based on my trained model? Is there any example for this? Thank you in advance.

I use `pyplot` a lot

```python
# magic line only if you use Jupyter
%matplotlib inline
from matplotlib import pyplot
...
# assuming mask is batch x h x w and we want the first
pyplot.subplot(1, 2, 1) # have two plots in 1 row two columns, first plot
# assuming im is batch x channel x h x w and channel is RGB
pyplot.plot(im[0].detach().cpu().permute(1, 2, 0))
pyplot.subplot(1, 2, 2) # second plot
pyplot.plot(mask[0].detach().cpu())

```

Best regards

Thomas

---

<div class="post-metadata">

**Author:** ![Neda](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/neda/32/7779_2.png) [@Neda](https://discuss.pytorch.org/u/Neda)\
**Post date:** [December 25, 2018, 8:00pm UTC](https://discuss.pytorch.org/t/how-can-i-display-a-test-image-and-display-the-mask-for-it-based-on-my-trained-model/33016/3 "2018-12-25T20:00:46Z")

</div>

@tom thank you. calling the model is working.

In `pyplot.plot(mask[0].detach().cpu())` I think you did plot the current mask (ground truth) which is already in the dataset. is that right?

Maybe I didn’t ask my question clearly. I wanted to test the model and display a test image from test\_loader and visualise a predicted segmentation mask for this test image to find out is model learned anything or not. (for example like Predicted Mask(U-Net) in fig 3 in [this paper](https://www.ncbi.nlm.nih.gov/pmc/articles/PMC5977596/).

---

<div class="post-metadata">

**Author:** ![tom](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/tom/32/3162_2.png) [@tom](https://discuss.pytorch.org/u/tom)\
**Post date:** [December 25, 2018, 8:36pm UTC](https://discuss.pytorch.org/t/how-can-i-display-a-test-image-and-display-the-mask-for-it-based-on-my-trained-model/33016/4 "2018-12-25T20:36:35Z")

</div>

Hi,

sorry, I didn’t adapt it well to your specific code. If just using `predict` in place of mask (based on accuracy involving `predicted.eq(mask.data)`) doesn’t work, what’s the shape of predict?

Best regards

Thomas

---

<div class="post-metadata">

**Author:** ![Neda](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/neda/32/7779_2.png) [@Neda](https://discuss.pytorch.org/u/Neda)\
**Post date:** [December 26, 2018, 2:59pm UTC](https://discuss.pytorch.org/t/how-can-i-display-a-test-image-and-display-the-mask-for-it-based-on-my-trained-model/33016/5 "2018-12-26T14:59:00Z")

</div>

@tom predict size is `torch.Size([1, 240, 320])` and image size is `torch.Size([1, 1, 240, 320])`

---

<div class="post-metadata">

**Author:** ![tom](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/tom/32/3162_2.png) [@tom](https://discuss.pytorch.org/u/tom)\
**Post date:** [December 28, 2018, 1:18pm UTC](https://discuss.pytorch.org/t/how-can-i-display-a-test-image-and-display-the-mask-for-it-based-on-my-trained-model/33016/6 "2018-12-28T13:18:17Z")

</div>

Then the above recipe should work for predict in place of mask.

Best regards

Thomas

---

<div class="post-metadata">

**Author:** ![Neda](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/neda/32/7779_2.png) [@Neda](https://discuss.pytorch.org/u/Neda)\
**Post date:** [December 29, 2018, 3:53pm UTC](https://discuss.pytorch.org/t/how-can-i-display-a-test-image-and-display-the-mask-for-it-based-on-my-trained-model/33016/7 "2018-12-29T15:53:54Z")

</div>

I still couldn’t figure out how can I visualise the output of network after training. This code will visualize the raw output but I don’t know how can I display all dim of image, at the moment will display only one channel in `plt.imshow(outputs[0,0,:,:].detach().cpu())` while the shape is `#print(outputs.shape) # torch.Size([1, 2, 240, 320])` it is the same issue with `plt.imshow(t_image[0,0,:,:].detach().cpu())` while the shape is `#print(t_image.size()) # torch.Size([1, 1, 240, 320])`

```auto
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 outputs
                    # 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()
    plt.subplot(1, 4, 1)
    #print(outputs.shape) # torch.Size([1, 2, 240, 320])
    plt.imshow(outputs[0,0,:,:].detach().cpu())
    plt.title('DL raw output')
    plt.subplot(1, 4, 2)
    plt.imshow(predict.detach().cpu().squeeze())
    plt.title('DL prediction')
    plt.subplot(1, 4, 3)
    plt.imshow(mask.detach().cpu().squeeze())
    plt.title('ground truth')
    plt.subplot(1, 4, 4)
    #print(t_image.size()) # torch.Size([1, 1, 240, 320])
    plt.imshow(t_image[0,0,:,:].detach().cpu())
    plt.title('original input image')

```

@ptrblck Could you please help to sort out this issue? I did try `permute(1, 2, 0))` but couldn’t solve the issue

---

<div class="post-metadata">

**Author:** ![ptrblck](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/32/1823_2.png) [@ptrblck](https://discuss.pytorch.org/u/ptrblck)\
**Post date:** [December 30, 2018, 12:42pm UTC](https://discuss.pytorch.org/t/how-can-i-display-a-test-image-and-display-the-mask-for-it-based-on-my-trained-model/33016/8 "2018-12-30T12:42:51Z")

</div>

I’m not sure, why it’s not working.  
If you would like to visualize both probability maps for the two classes, your code should work:

```python
plt.figure()
plt.imshow(torch.exp(outputs[0,0,:,:]).detach().cpu()) # plot class0
plt.figure()
plt.imshow(torch.exp(outputs[0,1,:,:]).detach().cpu()) # plot class1

```

What is your output currently showing?

---

<div class="post-metadata">

**Author:** ![Neda](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/neda/32/7779_2.png) [@Neda](https://discuss.pytorch.org/u/Neda)\
**Post date:** [December 30, 2018, 3:16pm UTC](https://discuss.pytorch.org/t/how-can-i-display-a-test-image-and-display-the-mask-for-it-based-on-my-trained-model/33016/9 "2018-12-30T15:16:49Z")

</div>

@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?

```auto

     #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')

```

---

<div class="post-metadata">

**Author:** ![InnovArul](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/innovarul/32/5282_2.png) [@InnovArul](https://discuss.pytorch.org/u/InnovArul)\
**Post date:** [December 30, 2018, 3:24pm UTC](https://discuss.pytorch.org/t/how-can-i-display-a-test-image-and-display-the-mask-for-it-based-on-my-trained-model/33016/10 "2018-12-30T15:24:17Z")

</div>

you could try to use grayscale colormap.

> <https://stackoverflow.com/questions/3823752/display-image-as-grayscale-using-matplotlib>

---

<div class="post-metadata">

**Author:** ![Neda](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/neda/32/7779_2.png) [@Neda](https://discuss.pytorch.org/u/Neda)\
**Post date:** [December 30, 2018, 4:38pm UTC](https://discuss.pytorch.org/t/how-can-i-display-a-test-image-and-display-the-mask-for-it-based-on-my-trained-model/33016/11 "2018-12-30T16:38:14Z")

</div>

@InnovArul thank you. Yes cmap = ‘gray’ display the original image as grayscale.

---

<div class="post-metadata">

**Author:** ![Neda](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/neda/32/7779_2.png) [@Neda](https://discuss.pytorch.org/u/Neda)\
**Post date:** [January 4, 2019, 1:50am UTC](https://discuss.pytorch.org/t/how-can-i-display-a-test-image-and-display-the-mask-for-it-based-on-my-trained-model/33016/12 "2019-01-04T01:50:57Z")

</div>

@ptrblck one more question here, if I visualize the probability of output same as your example

```auto
plt.figure()
plt.imshow(torch.exp(outputs[0,0,:,:]).detach().cpu()) # plot class0
plt.figure()
plt.imshow(torch.exp(outputs[0,1,:,:]).detach().cpu()) # plot class1

```

it return a figure which is different with the figure that visualize by following example

```auto
    output = model(t_image)
    prob = torch.exp(output)
    prob_imgs = make_grid(prob.permute(1, 0, 2, 3))
    prob_imgs = prob_imgs.detach()
    plt.imshow(prob_imgs.permute(1, 2, 0))

```

what is cause to be different in these example, is it because of make\_grid? because permute is just swap the first (batch size) and second dimensions (number of classes).  
this image is from second example, for the first example the zeros value in figures will be 1 and 1 converted to zeros. (black and white will swap). which way is correct way to do this? thank you in advance  
 ![train_exp3](https://discuss.pytorch.org/uploads/default/original/2X/a/ae8a8c4e387b82ee5322e2aa2739984004a729e8.png)

---

<div class="post-metadata">

**Author:** ![ptrblck](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/32/1823_2.png) [@ptrblck](https://discuss.pytorch.org/u/ptrblck)\
**Post date:** [January 4, 2019, 10:01pm UTC](https://discuss.pytorch.org/t/how-can-i-display-a-test-image-and-display-the-mask-for-it-based-on-my-trained-model/33016/13 "2019-01-04T22:01:45Z")

</div>

I guess the difference in the visualizations is because `make_grid` created a “color” image (i.e. a tensor with 3 channels), which `matplotlib` apparently treats differently than an image without any channel dimension.  
As long as the values and visualizations are consistent, you could just stick to the approach which you like more.

---

<div class="post-metadata">

**Author:** ![noornk](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/noornk/32/12076_2.png) [@noornk](https://discuss.pytorch.org/u/noornk)\
**Post date:** [June 28, 2019, 2:23am UTC](https://discuss.pytorch.org/t/how-can-i-display-a-test-image-and-display-the-mask-for-it-based-on-my-trained-model/33016/15 "2019-06-28T02:23:03Z")

</div>

Why do we have to plot the mask in two figures? I don’t get this.  
And how are we supposed to calculate the precision of the mask? I mean the area of intersection.

---

<div class="post-metadata">

**Author:** ![ptrblck](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/32/1823_2.png) [@ptrblck](https://discuss.pytorch.org/u/ptrblck)\
**Post date:** [June 28, 2019, 9:47am UTC](https://discuss.pytorch.org/t/how-can-i-display-a-test-image-and-display-the-mask-for-it-based-on-my-trained-model/33016/16 "2019-06-28T09:47:52Z")

</div>

The code plots the probability maps for both classes, not the mask.  
If you are only interested on a single class, you can of course just plot the corresponding prediction.

You could use the [IoU implementation](https://github.com/amdegroot/ssd.pytorch/blob/5b0b77faa955c1917b0c710d770739ba8fbff9b7/layers/box_utils.py#L50) or a dice score.

---

<div class="post-metadata">

**Author:** ![noornk](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/noornk/32/12076_2.png) [@noornk](https://discuss.pytorch.org/u/noornk)\
**Post date:** [July 12, 2019, 2:55am UTC](https://discuss.pytorch.org/t/how-can-i-display-a-test-image-and-display-the-mask-for-it-based-on-my-trained-model/33016/17 "2019-07-12T02:55:27Z")

</div>

Dear @ptrblck sorry for asking so many questions.  
I still have trouble understanding this. The input of the network, for a binary class, is a [1, w, h] image. The prediction is [1, 2, w, h] (assuming the batch size is one), and the mask(ground truth) is [1, w, h]. I have to main questions:

1. How is the loss computed? the loss should get each cell’s value (which 0 or 1), and compare it to the mask values(again 0 or 1).
2. How to get the final predicted mask, in terms of an image? the output is class probability, then how am I suppose to get 0 or 1?

---

<div class="post-metadata">

**Author:** ![ptrblck](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/32/1823_2.png) [@ptrblck](https://discuss.pytorch.org/u/ptrblck)\
**Post date:** [July 12, 2019, 10:40am UTC](https://discuss.pytorch.org/t/how-can-i-display-a-test-image-and-display-the-mask-for-it-based-on-my-trained-model/33016/18 "2019-07-12T10:40:26Z")

</div>

If you are using `nn.CrossEntropyLoss` (or `nn.NLLLoss`) for binary classification, each channel of the output will correspond to the logits (or log probabilities) of the class index.  
Have a look at the [CrossEntropyLoss docs](https://pytorch.org/docs/stable/nn.html#torch.nn.CrossEntropyLoss) to see the applied formula.  
As you can see, the target will be used as an index to get the output logit, which corresponds to this target.

In the case of a binary classification, you can also use `nn.BCEWithLogitsLoss`, which would then expect an output of `[batch_size, 1, h, w]` and a target of the same shape.

If you want to get the prediction in terms of class indices, you can call `preds = torch.argmax(output, 1)` in the former case or use a threshold, if you are using `nn.BCEWithLogitsLoss` e.g.:  
`preds = torch.sigmoid(output) > 0.5`.

---

<div class="post-metadata">

**Author:** ![noornk](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/noornk/32/12076_2.png) [@noornk](https://discuss.pytorch.org/u/noornk)\
**Post date:** [July 12, 2019, 1:31pm UTC](https://discuss.pytorch.org/t/how-can-i-display-a-test-image-and-display-the-mask-for-it-based-on-my-trained-model/33016/19 "2019-07-12T13:31:49Z")

</div>

Thank you so much for your reply.  
My output is in the shape of [batch\_size, 2, h, w], and target is [batch\_size, g, w], so I tried to use nn.CrossEntropyLoss, but I’m getting this error:

> RuntimeError: “host\_softmax” not implemented for ‘Long’

And when my output is [batch\_size, 2, h, w], take for example these two cells: [0, 0, 10, 10] and [0, 1, 10, 10]. Shouldn’t the values complement each other? because the cell is in either class?

---

<div class="post-metadata">

**Author:** ![ptrblck](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/32/1823_2.png) [@ptrblck](https://discuss.pytorch.org/u/ptrblck)\
**Post date:** [July 12, 2019, 1:43pm UTC](https://discuss.pytorch.org/t/how-can-i-display-a-test-image-and-display-the-mask-for-it-based-on-my-trained-model/33016/20 "2019-07-12T13:43:25Z")

</div>

Could you post the code snippet which throws this error?  
Are you passing the model output as float and targets as long to `nn.CrossEntropyLoss`?

> [@noornk](#):
>
> And when my output is [batch\_size, 2, h, w], take for example these two cells: [0, 0, 10, 10] and [0, 1, 10, 10]. Shouldn’t the values complement each other? because the cell is in either class?

Since they represent logits, you won’t see any pattern there.  
However, you could apply a softmax on dim1 and see the corresponding probabilities in each channel for the sake of debugging.

---

<div class="post-metadata">

**Author:** ![noornk](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/noornk/32/12076_2.png) [@noornk](https://discuss.pytorch.org/u/noornk)\
**Post date:** [July 12, 2019, 1:52pm UTC](https://discuss.pytorch.org/t/how-can-i-display-a-test-image-and-display-the-mask-for-it-based-on-my-trained-model/33016/21 "2019-07-12T13:52:47Z")

</div>

Yeah. fixed it! However, now in the middle of training all the outputs become zero.