How to crop image based on binary mask

Question

I am using torch with some semantic segmentation algorithms to produce a binary mask of the segmented images. I would then like to crop the images based on that mask. To be clea

Answer 1

Here is a solution relying only on numpy:

def get_segment_crop(img,tol=0, mask=None):
    if mask is None:
        mask = img > tol
    return img[np.ix_(mask.any(1), mask.any(0))]

now execute get_segment_crop(rgb, mask=segment_mask) where rgb is an ndarray of shape (w,h,c) and segment_mask is a boolean ndarray (i.e. containing True/False entries) of shape (w,h), given that w=width, h=height.