Robust_image_source_identification_on_modern_smartphones/datasets/raise/utils.py

from enum import Enum, auto
import skimage.restoration
import numpy as np
import rawpy
from tqdm import tqdm
from PIL import Image
from skimage import img_as_float
from datetime import datetime
import builtins as __builtin__
from scipy.ndimage import gaussian_filter
import os

class Color(Enum):
    RED = auto()
    GREEN_RIGHT = auto()
    GREEN_BOTTOM = auto()
    BLUE = auto()

    def __str__(self):
        return self.name.lower()

class Denoiser(Enum):
    WAVELET = auto()
    BILATERAL = auto()
    TV_CHAMBOLLE = auto()
    MEAN = auto()

    def __str__(self):
        return self.name.lower()

# Among:
# - `wavelet`
# - `bilateral`
# - `tv_chambolle`
def denoise(imageNpArray, denoiser):
    skImageRestorationDenoise = getattr(skimage.restoration, f'denoise_{denoiser}')

    match denoiser:
        case Denoiser.WAVELET:
            imageDenoisedNpArray = skImageRestorationDenoise(imageNpArray, rescale_sigma=True)
        case Denoiser.BILATERAL:
            imageDenoisedNpArray = skImageRestorationDenoise(imageNpArray, sigma_color=0.05, sigma_spatial=15)
        case Denoiser.TV_CHAMBOLLE:
            imageDenoisedNpArray = skImageRestorationDenoise(imageNpArray, weight=0.2)
    return imageDenoisedNpArray

class iterativeMean:
    mean = None
    numberOfElementsInMean = 0

    def add(self, element):
        if self.mean is None:
            self.mean = element
        else:
            self.mean = ((self.mean * self.numberOfElementsInMean) + element) / (self.numberOfElementsInMean + 1)
        self.numberOfElementsInMean += 1

RAW_IMAGE_FILE_EXTENSIONS = [
    'arw',
    'nef',
]

def getRawColorChannel(raw, color):
    colorDesc = raw.color_desc.decode('ascii')
    assert colorDesc == 'RGBG'
    assert np.array_equal(raw.raw_pattern, np.array([[0, 1], [3, 2]], dtype = np.uint8))
    # RG
    # GB
    rawImageVisible = raw.raw_image_visible.copy()

    redRawImageVisible = rawImageVisible[::2, ::2]
    greenRightRawImageVisible = rawImageVisible[::2, 1::2]

    greenBottomRawImageVisible = rawImageVisible[1::2, ::2]
    blueRawImageVisible = rawImageVisible[1::2, 1::2]

    match color:
        case Color.RED:
            imageNpArray = redRawImageVisible
        case Color.GREEN_RIGHT:
            imageNpArray = greenRightRawImageVisible
        case Color.GREEN_BOTTOM:
            imageNpArray = greenBottomRawImageVisible
        case Color.BLUE:
            imageNpArray = blueRawImageVisible
    return imageNpArray

def isARawImage(imageFilePath):
    return any([imageFilePath.lower().endswith(f'.{rawImageFileExtension}') for rawImageFileExtension in RAW_IMAGE_FILE_EXTENSIONS])

def getColorChannel(imageFilePath, color):
    if isARawImage(imageFilePath):
        numpyFilePath = f'{imageFilePath}.{color}.npy'
        if os.path.isfile(numpyFilePath):
            imageNpArray = np.load(numpyFilePath)
        else:
            with rawpy.imread(imageFilePath) as raw:
                imageNpArray = getRawColorChannel(raw, color)
    else:
        imagePil = Image.open(imageFilePath)
        imageNpArray = img_as_float(np.array(imagePil))
    return imageNpArray

def escapeFilePath(filePath):
    return filePath.replace('/', '_')


def getNewIndex(index, offset):
    newIndex = (index - offset) * 2 + offset
    return newIndex

def silentTqdm(data, desc = None):
    return data

def mergeSingleColorChannelImagesAccordingToBayerFilter(singleColorChannelImages):
    colorImage = singleColorChannelImages[list(Color)[0]]
    multipleColorsImage = np.empty([dimension * 2 for dimension in colorImage.shape], dtype = np.float64)
    '''
    Assume Bayer Filter:
    RG
    GB
    '''
    multipleColorsImage[::2,::2] = singleColorChannelImages[Color.RED]
    multipleColorsImage[1::2,::2] = singleColorChannelImages[Color.GREEN_RIGHT]
    multipleColorsImage[::2,1::2] = singleColorChannelImages[Color.GREEN_BOTTOM]
    multipleColorsImage[1::2,1::2] = singleColorChannelImages[Color.BLUE]
    return multipleColorsImage

def saveNpArray(fileName, npArray):
    with open(f'{fileName}.npy', 'wb') as f:
        np.save(f, npArray)

def rescaleRawImageForDenoiser(imageNpArray, minColor, maxColor):
    imageNpArray = (imageNpArray - minColor) / (maxColor - minColor)
    return imageNpArray

def updateExtremes(imageNpArray, minColor, maxColor):
    colorRawImageVisibleMin = imageNpArray.min()
    colorRawImageVisibleMax = imageNpArray.max()
    if minColor is None or colorRawImageVisibleMin < minColor:
        minColor = colorRawImageVisibleMin
    if maxColor is None or colorRawImageVisibleMax > maxColor:
        maxColor = colorRawImageVisibleMax
    return minColor, maxColor

def print(*toPrint):
    __builtin__.print(datetime.now(), *toPrint)

def getColorMeans(imagesFileNames, colors, singleColorChannelCropResolution = None):
    colorMeans = {}
    for color in colors:
        colorIterativeMean = iterativeMean()
        for imageFileName in tqdm(imagesFileNames, f'Computing mean of {str(color).replace("_", " ")} colored images'):
            imageNpArray = getImageNpArray(imageFileName, False, color, Denoiser.MEAN)
            if singleColorChannelCropResolution is not None:
                imageNpArray = getImageCrop(imageNpArray, singleColorChannelCropResolution)
            imageNpArray = gaussian_filter(imageNpArray, sigma = 5)
            colorIterativeMean.add(imageNpArray)
        colorMeans[color] = colorIterativeMean.mean
    return colorMeans

def getImageNpArray(imageFilePath, computeExtremes, color, denoiser):
    global minColor, maxColor
    imageNpArray = getColorChannel(imageFilePath, color)

    if computeExtremes:
        minColor, maxColor = updateExtremes(imageNpArray, minColor, maxColor)
        return

    if isARawImage(imageFilePath) and denoiser != Denoiser.MEAN:
        imageNpArray = rescaleRawImageForDenoiser(imageNpArray, minColor, maxColor)
        # Pay attention to range of values expected by the denoiser.
        # Indeed if provide the thousands valued raw image, then the denoiser only returns values between 0 and 1 and making the difference between both look pointless.
    return imageNpArray

def getImageCrop(image, cropResolution):
    imageCrop = image[:cropResolution[0], :cropResolution[1]]
    return imageCrop