Vectorized Binary Mask to RLE Encoding (MSCOCO)

binary_mask_to_rle.py

# Given a numpy binary mask, the following function converts it to RLE encoding that MSCOCO defines
# This implementation is highly vectorized and about 100X faster than naive python loop implmentations

def binary_mask_to_rle(binary_mask):
    # Flatten the binary mask
    flat_mask = binary_mask.ravel(order='F')

    # Calculate the differences between consecutive elements in the flattened binary mask
    diffs = np.diff(np.concatenate(([0], flat_mask, [0])))

    # Find the indices where the differences are non-zero
    run_starts = np.where(diffs != 0)[0]

    # Calculate the lengths of the runs
    run_lengths = np.diff(np.concatenate(([0] ,run_starts, [flat_mask.size])))

    if run_lengths[-1] == 0:
        run_lengths = run_lengths[:-1]

    # Create the RLE dictionary
    rle = {'counts': run_lengths.tolist(), 'size': list(binary_mask.shape)}

    return rle