How can I save prediction matrix as a png image?

I've got a binary segmentation model and want to save its' output inside of a png file. How to do it? How to make it 3 channel image?

segmentationimagepng
3 votesJO220.00
1 Answers
NN138.00
3

There are couple of ways you can do it. You can do it with OpenCV or Pillow

Solution 1

import cv2
import numpy as np

prediction = model(image)
pred = np.dstack([prediction, prediction, prediction])
pred = (pred * 255).astype(np.uint8)

cv2.imwrite('pred.png', pred)

Solution 2

import cv2
from PIL import Image
import numpy as np

prediction = model(x_image)
pred = np.dstack([prediction, prediction, prediction])
pred = (pred * 255).astype(np.uint8)

img = Image.fromarray(pred, 'RGB')
img.save('pred.png')
Reply
Couldn't find what you were looking for?and we will find an expert to answer.
How helpful was this page?