forensics/bitplane.py

import cv2 as cv
import numpy as np
from PIL import Image

def bit_plane_extractor(
    image: Image.Image,
    channel: str = "Luminance",
    bit: int = 0,
    filter_type: str = "Disabled"
) -> Image.Image:
    """Extract and visualize a bit plane from a selected channel of the image.

    Args:
        image (Image.Image, string: filepath): The input image to analyze.
        channel (str, optional): The channel to extract. Defaults to "Luminance".
        bit (int, optional): The bit to extract. Defaults to 0.
        filter_type (str, optional): The type of filter to apply. Defaults to "Disabled".
    """
    img = np.array(image.convert("RGB"))
    if channel == "Luminance":
        img = cv.cvtColor(img, cv.COLOR_RGB2GRAY)
    elif channel == "RGB Norm":
        b, g, r = cv.split(img.astype(np.float64))
        img = np.sqrt(np.power(b, 2) + np.power(g, 2) + np.power(r, 2)).astype(np.uint8)
    else:
        idx = {"Red": 0, "Green": 1, "Blue": 2}[channel]
        img = img[:, :, idx]
    plane = cv.bitwise_and(np.full_like(img, 2 ** bit), img)
    plane = cv.normalize(plane, None, 0, 255, cv.NORM_MINMAX).astype(np.uint8)
    if filter_type == "Median":
        plane = cv.medianBlur(plane, 3)
    elif filter_type == "Gaussian":
        plane = cv.GaussianBlur(plane, (3, 3), 0)
    return Image.fromarray(plane)