ALiCe P.
ALiCe P.

Reputation: 241

How to fill segmented area within a mask in image segmentation task?

I have an image of a lung and a predicted mask for that picture. They look like this: enter image description here

I want to get an image where I'll see the predicted area in some color. What I've managed to do so far is to draw the contours. It looks like this (in yellow):

enter image description here

My code:

lung = Image.fromarray(images_medseg[0])
if lung.mode != 'RGB':
lung = lung.convert('RGB')
lung.save("main.jpeg")

mask = Image.fromarray((masks_medseg[0] * 255).astype(np.uint8))
if mask.mode != 'RGB':
mask = mask.convert('RGB')
mask.save("segmented.jpeg")

seg  = cv2.imread('segmented.jpeg',cv2.IMREAD_GRAYSCALE)
main = cv2.imread('main.jpeg',cv2.IMREAD_GRAYSCALE)
main = cv2.cvtColor(main,cv2.COLOR_GRAY2BGR)

RGBforLabel = { 1:(0,0,255), 2:(0,255,255) }

# Find external contours
contours, _ = cv2.findContours(seg, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)

# Iterate over all contours
for i,c in enumerate(contours):
# Find mean colour inside this contour by doing a masked mean
mask = np.zeros(seg.shape, np.uint8)
cv2.drawContours(mask,[c],-1,255, -1)
cv2.imwrite(f"mask-{i}.png", mask)
mean,_,_,_ = cv2.mean(seg, mask=mask)
# DEBUG: print(f"i: {i}, mean: {mean}")

# Get appropriate colour for this label
label = 2 if mean > 1.0 else 1
colour = RGBforLabel.get(label)
# DEBUG: print(f"Colour: {colour}")

# Outline contour in that colour on main image, line thickness=1
cv2.drawContours(main,[c],-1,colour,1)



# Save result
cv2.imwrite('result.png',main) 

res = mpimg.imread('/content/result.png')
imgplot = plt.imshow(res)
plt.show()

Could someone please explain what should I do to get something like this: enter image description here

Upvotes: 0

Views: 1792

Answers (1)

fmw42
fmw42

Reputation: 53164

You can use the mask to colorize your image where the mask is white using Numpy in Python/OpenCV. If we assume your mask is binary (0 and 255), then

result = image.copy()
result[mask==255] = [0,255,255]

Or, if that does not work and your mask is 3 channels but binary (0 and 255), then

result = image.copy()
result[np.where((mask==[255,255,255]).all(axis=2))] = [0,255,255]

Upvotes: 4

Related Questions