import numpy as np import pywt import cv2 from PIL import Image def wavelet_blocking_noise_estimation(image: Image.Image, blocksize: int = 8) -> Image.Image: """Estimate local noise using wavelet blocking. Returns 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)