# Visualise the test images after training the model on segmentation task

**URL:** <https://discuss.pytorch.org/t/visualise-the-test-images-after-training-the-model-on-segmentation-task/56615>\
**Category:** vision\
**Created:** [September 23, 2019, 8:54am UTC](https://discuss.pytorch.org/t/visualise-the-test-images-after-training-the-model-on-segmentation-task/56615 "2019-09-23T08:54:20Z")\
**Posts on this page:** 1\
**Showing post:** 4

<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:** [September 23, 2019, 11:46pm UTC](https://discuss.pytorch.org/t/visualise-the-test-images-after-training-the-model-on-segmentation-task/56615/4 "2019-09-23T23:46:31Z")

</div>

Assuming you are using the dummy code from [here](https://discuss.pytorch.org/t/training-semantic-segmentation/49275/4), you could do the following:

```python
rev_mapping = {mapping[k]: k for k in mapping}
pred = mask # or e.g. pred = torch.randint(0, 19, (224, 224))
pred_image = torch.zeros(3, pred.size(0), pred.size(1), dtype=torch.uint8)
for k in rev_mapping:
    print(k)
    pred_image[:, pred==k] = torch.tensor(rev_mapping[k]).byte().view(3, 1)

plt.imshow(pred_image.permute(1, 2, 0).numpy())

```

---

_[View the full topic](https://discuss.pytorch.org/t/visualise-the-test-images-after-training-the-model-on-segmentation-task/56615)._