Leverage iterativeMean in extract_noise.py

This commit is contained in:
2024-04-26 14:25:08 +02:00
parent 674f87d480
commit 1ea1bde910

View File

@@ -7,7 +7,7 @@ import os
from tqdm import tqdm from tqdm import tqdm
import csv import csv
import rawpy import rawpy
from utils import Color, denoise from utils import Color, denoise, iterativeMean
import matplotlib.pyplot as plt import matplotlib.pyplot as plt
from scipy.ndimage import gaussian_filter from scipy.ndimage import gaussian_filter
@@ -102,7 +102,7 @@ def getImageNpArray(imageFileName, computeExtremes, color):
# `color` is the actual color to estimate PRNU with. # `color` is the actual color to estimate PRNU with.
def treatImage(imageFileName, computeExtremes = False, color = None): def treatImage(imageFileName, computeExtremes = False, color = None):
global mean, numberOfImagesInMean global estimatedPrnuIterativeMean
imageNpArray = getImageNpArray(imageFileName, computeExtremes, color) imageNpArray = getImageNpArray(imageFileName, computeExtremes, color)
if imageNpArray is None: if imageNpArray is None:
return return
@@ -111,11 +111,7 @@ def treatImage(imageFileName, computeExtremes = False, color = None):
else: else:
imageDenoisedNpArray = means[color] imageDenoisedNpArray = means[color]
imageNoiseNpArray = imageNpArray - imageDenoisedNpArray imageNoiseNpArray = imageNpArray - imageDenoisedNpArray
if mean is None: estimatedPrnuIterativeMean.add(imageNoiseNpArray)
mean = imageNoiseNpArray
else:
mean = ((mean * numberOfImagesInMean) + imageNoiseNpArray) / (numberOfImagesInMean + 1)
numberOfImagesInMean += 1
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. # Assuming same intensity scale across color channels.
@@ -135,27 +131,21 @@ def saveNpArray(fileName, npArray):
if denoiser == 'mean': if denoiser == 'mean':
means = {} means = {}
for color in colors: for color in colors:
colorMean = None colorIterativeMean = iterativeMean()
numberOfImagesInColorMean = 0
for imageFileName in tqdm(imagesFileNames, f'Computing mean of {color} colored images'): for imageFileName in tqdm(imagesFileNames, f'Computing mean of {color} colored images'):
imageNpArray = getImageNpArray(imageFileName, False, color) imageNpArray = getImageNpArray(imageFileName, False, color)
imageNpArray = gaussian_filter(imageNpArray, sigma = 5) imageNpArray = gaussian_filter(imageNpArray, sigma = 5)
if colorMean is None: colorIterativeMean.add(imageNpArray)
colorMean = imageNpArray means[color] = colorIterativeMean.mean
else:
colorMean = ((colorMean * numberOfImagesInColorMean) + imageNpArray) / (numberOfImagesInColorMean + 1)
numberOfImagesInColorMean += 1
means[color] = colorMean
fileName = f'mean_{imagesFolderPathFileName}_{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. # 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, colorMean) saveNpArray(fileName, colorIterativeMean.mean)
for color in colors: for color in colors:
mean = None estimatedPrnuIterativeMean = iterativeMean()
numberOfImagesInMean = 0
for imageFileName in tqdm(imagesFileNames, f'Denoising images for color {color}'): for imageFileName in tqdm(imagesFileNames, f'Denoising images for color {color}'):
treatImage(imageFileName, color = color) treatImage(imageFileName, color = color)
npArrayFilePath = f'mean_{imagesFolderPathFileName}_{denoiser}_{color}' npArrayFilePath = f'mean_{imagesFolderPathFileName}_{denoiser}_{color}'
saveNpArray(npArrayFilePath, mean) saveNpArray(npArrayFilePath, estimatedPrnuIterativeMean.mean)