Convert semantic segmentation prediction to opencv Mat in python

vision
1

I am looking for a way to convert prediction(tensor) of semantic segmentation network to OpenCV Mat in python so it can be processed by OpenCV, e.g.:

...
rgb = Variable(rgb).cuda()
semantic = model(rgb.unsqueeze(0))
_,pred = torch.max(semantic, dim=1)
pred = pred*255
#contours, here = cv2.findContours(pred,cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) # this line won't work
torchvision.utils.save_image(pred, 'pred.png') # this line saves prediction as binary image correctly

I would like also the commented line to work since the last line tells me that ‘pred’ is already binary image and should be able to be passed to OpenCV functions like findContours(), I’ve tried like ‘np.transpose(pred.cpu().numpy(), (1,2,0))’ and it doesn’t work. I see some examples in C++ like here, but how can i achieve that in python?

Thanks in advance!

2

OpenCV works with numpy arrays in Python, so you can just call pred.numpy() to get the underlying array.
Your transpose/permute call also seems to be valid, as OpenCV works with channel-last images.

The issue seems to be more related to OpenCV than PyTorch, but what kind of error do you get in cv2.findContours, if you pass the numpy array as the input?

3

Hey, yes you are right it’s more related to OpenCV than Pytorch. I can’t remember the error message anymore because I found the solution myself:

cv2.findContours(pred, ...) --> cv2.findContours(np.uint8(pred), ...)