Cropping image from a binary mask

Question

I am trying to use DeepLab v3 to detect object and mask where the actual object is. DeepLab model produces a resized_im (3D) and a mask seg_map (2D) of 0 and non-zero values, 0 means it's the background.

Currently, it is only possible to plot an image with an overlay mask on the object. I want to crop the object out of the resized_im with transparent background. Is there any advice for the work?

You can play around with the notebook here: https://colab.research.google.com/drive/138dTpcYfne40hqrb13n_36okSGYhrJnz?usp=sharing&hl=en#scrollTo=p47cYGGOQE1W&forceEdit=true&sandboxMode=true

I also tried here: How to crop image based on binary mask but none seems to work on my case

Answer 1

You just need to convert your segmentation mask to boolean numpy array, then multiply image by it. Don't forget that your image has 3 channels while mask has only 1. It may look something like that:

# seg_map - segmentation mask from network, resized_im - your input image
mask = np.greater(seg_map, 0) # get only non-zero positive pixels/labels
mask = np.expand_dims(mask, axis=-1) # (H, W) -> (H, W, 1)
mask = np.concatenate((mask, mask, mask), axis=-1) # (H, W, 1) -> (H, W, 3), (don't like it, so if you know how to do it better, please let me know)
crops = resized_im * mask # apply mask on image

You can use different logical numpy function if you want to choose certain labels, for example:

mask = np.equal(seg_map, 5) # to get only objects with label 5