Hide and retrieve ASCII messages on image bit planes

steganography.py

import numpy as np

def encode(bgr_img, message, bit_plane):
    height, width, depth = bgr_img.shape
    max_bytes = height * width * depth // 8
    max_bits = max_bytes * 8 # we can only store whole byte words

    if len(message) < max_bytes:
        message += b'\0' # mark the end of the message

    message_bits = np.unpackbits(np.frombuffer(message, dtype=np.uint8))
    if message_bits.size > max_bits:
        message_bits = message_bits[:max_bits]
    
    r_message = message_bits[0::3]; r_length = r_message.size
    g_message = message_bits[1::3]; g_length = g_message.size
    b_message = message_bits[2::3]; b_length = b_message.size
    
    # NOTE OpenCV uses BGR order
    r = bgr_img[..., 2].ravel()
    g = bgr_img[..., 1].ravel()
    b = bgr_img[..., 0].ravel()

    # hide message in bit_plane
    mask = 1 << bit_plane
    r[:r_length] = (r[:r_length] & ~mask) | (r_message << bit_plane)
    g[:g_length] = (g[:g_length] & ~mask) | (g_message << bit_plane)
    b[:b_length] = (b[:b_length] & ~mask) | (b_message << bit_plane)
    
    return np.dstack((b.reshape((height, width)), 
                      g.reshape((height, width)), 
                      r.reshape((height, width))))
    
def decode(bgr_img, bit_plane):
    mask = 1 << bit_plane
    __img = (bgr_img & mask) >> bit_plane

    # NOTE OpenCV uses BGR order
    b_message = __img[..., 0].ravel()
    g_message = __img[..., 1].ravel()
    r_message = __img[..., 2].ravel()
    
    # retrieve message from bit_plane
    message_bits = np.dstack((r_message, g_message, b_message)).ravel()
    max_bits = message_bits.size
    while max_bits % 8 != 0:
        max_bits -= 1 # we can only store whole byte words
    
    message = ''.join([chr(byte) for byte in np.packbits(message_bits[:max_bits])])

    # look for a '\0' marking the end of the message
    end = message.find('\0')
    if end != -1:
        message = message[:end]
    
    return message