Leverage iterativeMean in extract_noise.py

This commit is contained in:
Benjamin Loison 2024-04-26 14:25:08 +02:00
parent 674f87d480
commit 1ea1bde910
Signed by: Benjamin_Loison
SSH Key Fingerprint: SHA256:BtnEgYTlHdOg1u+RmYcDE0mnfz1rhv5dSbQ2gyxW8B8

View File

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