diff --git a/datasets/raise/extract_noise.py b/datasets/raise/extract_noise.py index 264d14a..ad6aad2 100755 --- a/datasets/raise/extract_noise.py +++ b/datasets/raise/extract_noise.py @@ -7,19 +7,19 @@ import csv from utils import Color, denoise, iterativeMean, escapeFilePath, saveNpArray, getColorMeans, getImageNpArray import matplotlib.pyplot as plt -imagesFolderPath = 'rafael/arw' -imagesFolderPathFileName = escapeFilePath(imagesFolderPath) +IMAGES_FOLDER_PATH = 'rafael/230424' +imagesFolderPathFileName = escapeFilePath(IMAGES_FOLDER_PATH) # Among: # `denoise` possible denoisers and `mean`. -denoiser = 'mean' +DENOISER = 'mean' -raiseNotFlatFields = False +RAISE_NOT_FLAT_FIELDS = False # `[Color.RED, Color.GREEN_RIGHT, ...]` or `Color` or `[None]` for not raw images. -colors = [None] +COLORS = [None] -imagesFileNames = os.listdir(imagesFolderPath + ('/png' if raiseNotFlatFields else '')) +imagesFileNames = os.listdir(IMAGES_FOLDER_PATH + ('/png' if RAISE_NOT_FLAT_FIELDS else '')) -if raiseNotFlatFields: +if RAISE_NOT_FLAT_FIELDS: files = {} with open('RAISE_all.csv') as csvfile: @@ -28,7 +28,7 @@ if raiseNotFlatFields: file = row['File'] + '.png' files[file] = row - imagesFileNames = [imageFileName for imageFileName in tqdm(imagesFileNames, 'Filtering images') if files[imageFileName]['Device'] == 'Nikon D7000' and Image.open(f'{imagesFolderPath}/png/{imageFileName}').size == (4946, 3278)] + imagesFileNames = [imageFileName for imageFileName in tqdm(imagesFileNames, 'Filtering images') if files[imageFileName]['Device'] == 'Nikon D7000' and Image.open(f'{IMAGES_FOLDER_PATH}/png/{imageFileName}').size == (4946, 3278)] # Among: # - `None` @@ -47,11 +47,11 @@ minColor = None maxColor = None def getImageFilePath(imageFileName): - if raiseNotFlatFields: + if RAISE_NOT_FLAT_FIELDS: imageFileName = imageFileName.replace('.png', '.NEF') - imageFilePath = f'{imagesFolderPath}/nef/{imageFileName}' + imageFilePath = f'{IMAGES_FOLDER_PATH}/nef/{imageFileName}' else: - imageFilePath = f'{imagesFolderPath}/{imageFileName}' + imageFilePath = f'{IMAGES_FOLDER_PATH}/{imageFileName}' return imageFilePath # `color` is the actual color to estimate PRNU with. @@ -61,17 +61,17 @@ def treatImage(imageFileName, computeExtremes = False, color = None): imageNpArray = getImageNpArray(imageFilePath, computeExtremes, color) if imageNpArray is None: return - if denoiser != 'mean': - imageDenoisedNpArray = denoise(imageNpArray, denoiser) + if DENOISER != 'mean': + imageDenoisedNpArray = denoise(imageNpArray, DENOISER) else: imageDenoisedNpArray = colorMeans[color] imageNoiseNpArray = imageNpArray - imageDenoisedNpArray estimatedPrnuIterativeMean.add(imageNoiseNpArray) -if (minColor is None or maxColor is None) and denoiser != 'mean': +if (minColor is None or maxColor is None) and DENOISER != 'mean': # Assuming same intensity scale across color channels. for imageFileName in tqdm(imagesFileNames, 'Computing extremes of images'): - for color in colors: + for color in COLORS: treatImage(imageFileName, computeExtremes = True, color = color) # To skip this step next time. @@ -79,19 +79,19 @@ if (minColor is None or maxColor is None) and denoiser != 'mean': print(f'{minColor=}') print(f'{maxColor=}') -if denoiser == 'mean': - colorMeans = getColorMeans(imagesFileNames, colors) +if DENOISER == 'mean': + colorMeans = getColorMeans(imagesFileNames, COLORS) for color in Color: colorMeans[color] = colorMeans[color] fileName = f'mean_{imagesFolderPathFileName}_{color}' # Then use `merge_single_color_channel_images_according_to_bayer_filter.py` to consider all color channels, instead of saving this single color channel as an image. saveNpArray(fileName, colorMeans[color]) -for color in colors: +for color in COLORS: estimatedPrnuIterativeMean = iterativeMean() for imageFileName in tqdm(imagesFileNames, f'Denoising images for color {color}'): treatImage(imageFileName, color = color) - npArrayFilePath = f'mean_{imagesFolderPathFileName}_{denoiser}_{color}' + npArrayFilePath = f'mean_{imagesFolderPathFileName}_{DENOISER}_{color}' saveNpArray(npArrayFilePath, estimatedPrnuIterativeMean.mean)