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 3, but how can i achieve that in python?
Thanks in advance!