From 1ea1bde9102a1bdea8c472a4c0e06ffb9ddf21f7 Mon Sep 17 00:00:00 2001
From: Benjamin Loison <benjamin_loison@users.noreply.gitea.lemnoslife.com>
Date: Fri, 26 Apr 2024 14:25:08 +0200
Subject: [PATCH] Leverage `iterativeMean` in `extract_noise.py`

---
 datasets/raise/extract_noise.py | 28 +++++++++-------------------
 1 file changed, 9 insertions(+), 19 deletions(-)

diff --git a/datasets/raise/extract_noise.py b/datasets/raise/extract_noise.py
index 14cf151..b2a15c7 100755
--- a/datasets/raise/extract_noise.py
+++ b/datasets/raise/extract_noise.py
@@ -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)