Files
Robust_image_source_identif…/datasets/raise/split_and_compare_prnus_of_subgroups.py
2024-04-30 03:49:45 +02:00

71 lines
3.2 KiB
Python
Executable File

#!/usr/bin/env python
import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm
from utils import denoise, iterativeMean, getColorChannel, escapeFilePath, Color, mergeSingleColorChannelImagesAccordingToBayerFilter, rescaleRawImageForDenoiser, updateExtremes, saveNpArray
import sys
import os
import random
sys.path.insert(0, '../../algorithms/distance/')
from rms_diff import rmsDiffNumpy
NUMBER_OF_SUBGROUPS = 2
DENOISER = 'wavelet'
IMAGES_FOLDER = 'flat-field/NEF'
setting = escapeFilePath(IMAGES_FOLDER) + f'_{DENOISER}'
imagesFileNames = os.listdir(IMAGES_FOLDER)
random.seed(0)
# To not have a bias (chronological for instance) when split to make subgroups.
random.shuffle(imagesFileNames)
numberOfImagesPerSubgroup = len(imagesFileNames) // NUMBER_OF_SUBGROUPS
subgroupsIterativeMean = [iterativeMean() for _ in range(NUMBER_OF_SUBGROUPS)]
rmss = []
minColor = None
maxColor = None
returnSingleColorChannelImage = lambda singleColorChannelImage, _minColor, _maxColor: singleColorChannelImage
for computeExtremes in tqdm(([True] if minColor is None or maxColor is None else []) + [False], 'Compute extremes'):
rescaleIfNeeded = returnSingleColorChannelImage if computeExtremes else rescaleRawImageForDenoiser
for subgroupImageIndex in tqdm(range(numberOfImagesPerSubgroup), 'Subgroup image index'):
for subgroupIndex in tqdm(range(NUMBER_OF_SUBGROUPS), 'Subgroup'):
imageIndex = (subgroupIndex * NUMBER_OF_SUBGROUPS) + subgroupImageIndex
imageFileName = imagesFileNames[imageIndex]
imageFilePath = f'{IMAGES_FOLDER}/{imageFileName}'
singleColorChannelImages = {color: rescaleIfNeeded(getColorChannel(imageFilePath, color), minColor, maxColor) for color in Color}
multipleColorsImage = mergeSingleColorChannelImagesAccordingToBayerFilter(singleColorChannelImages)
if computeExtremes:
minColor, maxColor = updateExtremes(multipleColorsImage, minColor, maxColor)
continue
singleColorChannelDenoisedImages = {color: denoise(singleColorChannelImages[color], DENOISER) for color in Color}
multipleColorsDenoisedImage = mergeSingleColorChannelImagesAccordingToBayerFilter(singleColorChannelDenoisedImages)
imagePrnuEstimateNpArray = multipleColorsImage - multipleColorsDenoisedImage
subgroupIterativeMean = subgroupsIterativeMean[subgroupIndex]
subgroupIterativeMean.add(imagePrnuEstimateNpArray)
if subgroupIndex == NUMBER_OF_SUBGROUPS - 1:
assert NUMBER_OF_SUBGROUPS == 2
rms = rmsDiffNumpy(subgroupIterativeMean.mean, subgroupsIterativeMean[1 - subgroupIndex].mean)
rmss += [rms]
for subgroupIndex in range(NUMBER_OF_SUBGROUPS):
plt.imsave(f'{setting}_estimated_prnu_subgroup_{subgroupIndex}.png', (subgroupsIterativeMean[subgroupIndex].mean))
plt.title(f'RMS between both subgroups estimated PRNUs with {DENOISER} denoiser for a given number of images among them')
plt.xlabel('Number of images of each subgroup')
plt.ylabel('RMS between both subgroups estimated PRNUs')
plt.plot(rmss)
saveNpArray(f'{setting}_rmss', rmss)
plt.savefig(f'{setting}_rms_between_estimated_prnu_of_2_subgroups.svg')
#plt.show()