RLE: rle_encode / rle_decode

rle

def rle_encode(mask):
    """Encodes a mask in Run Length Encoding (RLE).
    Returns a string of space-separated values.
    """
    assert mask.ndim == 2, "Mask must be of shape [Height, Width]"
    # Flatten it column wise
    m = mask.T.flatten()
    # Compute gradient. Equals 1 or -1 at transition points
    g = np.diff(np.concatenate([[0], m, [0]]), n=1)
    # 1-based indicies of transition points (where gradient != 0)
    rle = np.where(g != 0)[0].reshape([-1, 2]) + 1
    # Convert second index in each pair to lenth
    rle[:, 1] = rle[:, 1] - rle[:, 0]
    return " ".join(map(str, rle.flatten()))


def rle_decode(rle, shape=(101,101)):
    """Decodes an RLE encoded list of space separated
    numbers and returns a binary mask."""
    try:
        rle = list(map(int, rle.split()))
        rle = np.array(rle, dtype=np.int32).reshape([-1, 2])
        rle[:, 1] += rle[:, 0]
        rle -= 1
        mask = np.zeros([shape[0] * shape[1]], np.bool)
        for s, e in rle:
            assert 0 <= s < mask.shape[0]
            assert 1 <= e <= mask.shape[0], "shape: {}  s {}  e {}".format(shape, s, e)
            mask[s:e] = 1
        # Reshape and transpose
        mask = mask.reshape([shape[1], shape[0]]).T
        return mask
    except:
        return np.zeros(shape)