forensics/wavelet.py

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

def noise_estimation(image: Image.Image, blocksize: int = 8) -> Image.Image:
    """Estimate local noise using wavelet blocking. Returns a PIL image of the noise map.

    Args:
        image (Image.Image, string: filepath): The input image to analyze.
        blocksize (int): The size of the blocks to use for wavelet blocking.

    Returns:
        Image.Image: A PIL image of the noise map.
    """    
    im = np.array(image.convert('L'))
    y = np.double(im)
    cA1, (cH, cV, cD) = pywt.dwt2(y, 'db8')
    cD = cD[:cD.shape[0] // blocksize * blocksize, :cD.shape[1] // blocksize * blocksize]
    block = np.zeros((cD.shape[0] // blocksize, cD.shape[1] // blocksize, blocksize ** 2))
    for ii in range(0, cD.shape[0] - blocksize + 1, blocksize):
        for jj in range(0, cD.shape[1] - blocksize + 1, blocksize):
            block_elements = cD[ii:ii+blocksize, jj:jj+blocksize]
            block[ii // blocksize, jj // blocksize, :] = block_elements.flatten()
    noise_map = np.median(np.abs(block), axis=2) / 0.6745
    noise_map_8u = cv2.normalize(noise_map, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
    resized_noise_map = cv2.resize(noise_map_8u, (im.shape[1], im.shape[0]), interpolation=cv2.INTER_NEAREST)
    return Image.fromarray(resized_noise_map)