133 lines
4.2 KiB
Python
133 lines
4.2 KiB
Python
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__
|
|
|
|
class Color(Enum):
|
|
RED = auto()
|
|
GREEN_RIGHT = auto()
|
|
GREEN_BOTTOM = auto()
|
|
BLUE = auto()
|
|
|
|
def __str__(self):
|
|
return self.name.lower()
|
|
|
|
# Among:
|
|
# - `wavelet`
|
|
# - `bilateral`
|
|
# - `tv_chambolle`
|
|
def denoise(imageNpArray, denoiserName):
|
|
skImageRestorationDenoise = getattr(skimage.restoration, f'denoise_{denoiserName}')
|
|
|
|
match denoiserName:
|
|
case 'wavelet':
|
|
imageDenoisedNpArray = skImageRestorationDenoise(imageNpArray, rescale_sigma=True)
|
|
case 'bilateral':
|
|
imageDenoisedNpArray = skImageRestorationDenoise(imageNpArray, sigma_color=0.05, sigma_spatial=15)
|
|
case '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):
|
|
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)
|