96 lines
3.5 KiB
Python
Executable File
96 lines
3.5 KiB
Python
Executable File
#!/usr/bin/env python
|
|
|
|
import numpy as np
|
|
import os
|
|
from tqdm import tqdm
|
|
import csv
|
|
from utils import Color, denoise, iterativeMean, escapeFilePath, saveNpArray, getColorMeans, getImageNpArray, Denoiser
|
|
import matplotlib.pyplot as plt
|
|
|
|
IMAGES_FOLDER_PATH = 'rafael/230424'
|
|
imagesFolderPathFileName = escapeFilePath(IMAGES_FOLDER_PATH)
|
|
DENOISER = Denoiser.WAVELET
|
|
|
|
RAISE_NOT_FLAT_FIELDS = False
|
|
# `[Color.RED, Color.GREEN_RIGHT, ...]` or `Color` or `[None]` for not raw images.
|
|
COLORS = [None]
|
|
|
|
imagesFileNames = os.listdir(IMAGES_FOLDER_PATH + ('/png' if RAISE_NOT_FLAT_FIELDS else ''))
|
|
|
|
if RAISE_NOT_FLAT_FIELDS:
|
|
files = {}
|
|
|
|
with open('RAISE_all.csv') as csvfile:
|
|
reader = csv.DictReader(csvfile)
|
|
for row in tqdm(list(reader), 'CSV parsing'):
|
|
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'{IMAGES_FOLDER_PATH}/png/{imageFileName}').size == (4946, 3278)]
|
|
|
|
# Among:
|
|
# - `None`
|
|
# - `'sky'`
|
|
# - `'wall'`
|
|
type_ = None
|
|
if type_ is not None:
|
|
ranges = {
|
|
'sky': range(2_699, 2_807),
|
|
'wall': range(2_807, 2_912),
|
|
}
|
|
imagesFileNames = [f'DSC0{imageIndex}.ARW' for imageIndex in ranges[type_]]
|
|
imagesFolderPathFileName += f'_{type_}'
|
|
|
|
minColor = None
|
|
maxColor = None
|
|
|
|
def getImageFilePath(imageFileName):
|
|
if RAISE_NOT_FLAT_FIELDS:
|
|
imageFileName = imageFileName.replace('.png', '.NEF')
|
|
imageFilePath = f'{IMAGES_FOLDER_PATH}/nef/{imageFileName}'
|
|
else:
|
|
imageFilePath = f'{IMAGES_FOLDER_PATH}/{imageFileName}'
|
|
return imageFilePath
|
|
|
|
# `color` is the actual color to estimate PRNU with.
|
|
def treatImage(imageFileName, computeExtremes = False, color = None):
|
|
global estimatedPrnuIterativeMean
|
|
imageFilePath = getImageFilePath(imageFileName)
|
|
imageNpArray = getImageNpArray(imageFilePath, computeExtremes, color, DENOISER)
|
|
if imageNpArray is None:
|
|
return
|
|
if DENOISER != 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 != Denoiser.MEAN:
|
|
# Assuming same intensity scale across color channels.
|
|
for imageFileName in tqdm(imagesFileNames, 'Computing extremes of images'):
|
|
for color in COLORS:
|
|
treatImage(imageFileName, computeExtremes = True, color = color)
|
|
|
|
# To skip this step next time.
|
|
# Maybe thanks to `rawpy.RawPy` fields, possibly stating device maximal value, can avoid doing so to some extent.
|
|
print(f'{minColor=}')
|
|
print(f'{maxColor=}')
|
|
|
|
if DENOISER == 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:
|
|
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, estimatedPrnuIterativeMean.mean)
|