Convert PyTorch Tensor to OpenCV Image using Python

While PyTorch provides tools for building and training deep learning models, it does not provide many image processing functions that are available in OpenCV. Converting PyTorch tensor to OpenCV image allows us to take advantage of advanced image processing functions available in OpenCV. Fortunately, the conversion process is pretty simple. This tutorial shows how to do that using Python.

Prepare environment

• Install the following packages using pip:

pip install torch torchvision

pip install opencv-python

Code

In the following code, we read the image as a PyTorch tensor. It has a shape (C, H, W) where C is the number of channels, H is the height, and W is the width. Next, we convert the tensor to NumPy array, since OpenCV represents images in NumPy array format. We transpose NumPy array to change the shape from (C, H, W) to (H, W, C). This is necessary because OpenCV images have a shape of (H, W, C). OpenCV uses BGR as its default color space, while PyTorch uses RGB. Therefore, we're converting the color space from RGB to BGR. That's it, now the image can be processed using processing functions provided by OpenCV.

import cv2
from torchvision.io import read_image

tensorImg = read_image('test.jpg')
arrImg = tensorImg.numpy().transpose(1, 2, 0)

cvImg = cv2.cvtColor(arrImg, cv2.COLOR_RGB2BGR)