Load lazily images (#62)

datasets/raise/attribute_source_camera.py

@@ -24,7 +24,7 @@ PREDICT_ONLY_ON_WHOLE_TRAINING_SET = False
 
 setting = ','.join([escapeFilePath(imageCameraFolder) for imageCameraFolder in IMAGES_CAMERAS_FOLDER]) + f'_{DENOISER}'
 
-imagesCamerasFileNames = {camera: os.listdir(imageCameraFolder) for camera, imageCameraFolder in IMAGES_CAMERAS_FOLDER.items()}
+imagesCamerasFileNames = {camera: [imageCameraFile for imageCameraFile in os.listdir(imageCameraFolder) if imageCameraFile.endswith('.NEF') or imageCameraFile.endswith('.ARW')] for camera, imageCameraFolder in IMAGES_CAMERAS_FOLDER.items()}
 # Fix randomness for reproducibility.
 random.seed(0)
 # Randomize order to not have a bias (chronological for instance) when split to make training and testing sets.
@@ -104,15 +104,20 @@ from utils import silentTqdm
 returnSingleColorChannelImage = lambda singleColorChannelImage, _minColor, _maxColor: singleColorChannelImage
 
 # Assume to have `{min,max}Color` hardcoded.
-print('Load training images to memory')
+# Can just load to memory `getSingleColorChannelImages`, see [Robust_image_source_identification_on_modern_smartphones/issues/62#issuecomment-1861](https://gitea.lemnoslife.com/Benjamin_Loison/Robust_image_source_identification_on_modern_smartphones/issues/62#issuecomment-1861).
 rescaleIfNeeded = rescaleRawImageForDenoiser
 cameraTrainingImages = {}
-for cameraTrainingImageIndex in tqdm(range(numberOfTrainingImages), 'Camera training image index'):
-    for cameraIndex, camera in enumerate(tqdm(IMAGES_CAMERAS_FOLDER, 'Camera')):
+for cameraTrainingImageIndex in tqdm(range(numberOfTrainingImages), 'Load to memory camera training image'):
+    for cameraIndex, camera in enumerate(IMAGES_CAMERAS_FOLDER):
         singleColorChannelImages = getSingleColorChannelImages(camera, cameraTrainingImageIndex)
         multipleColorsImage = getMultipleColorsImage(singleColorChannelImages)
         cameraTrainingImages[camera] = cameraTrainingImages.get(camera, []) + [multipleColorsImage]
-print('Training images loaded to memory')
+
+singleColorChannelTestingImages = {camera: [] for camera in IMAGES_CAMERAS_FOLDER}
+for camera in IMAGES_CAMERAS_FOLDER:
+    for cameraTestingImageIndex in tqdm(range(numberOfTestingImages), 'Load to memory camera testing image'):
+        singleColorChannelImages = getSingleColorChannelImages(camera, numberOfTrainingImages + cameraTestingImageIndex)
+        singleColorChannelTestingImages[camera] += [singleColorChannelImages]
 
 # 2 loops:
 # - the first one is about computing `{min,max}Color`
@@ -128,7 +133,7 @@ for computeExtremes in tqdm(([True] if minColor is None or maxColor is None else
                 for cameraTestingImageIndex in tqdm(range(numberOfTestingImages), 'Camera testing image index'):
                     print(f'{camera=} {numberOfTrainingImages + cameraTestingImageIndex=}')
 
-                    singleColorChannelImages = getSingleColorChannelImages(camera, numberOfTrainingImages + cameraTestingImageIndex)
+                    singleColorChannelImages = singleColorChannelTestingImages[camera][cameraTestingImageIndex]#getSingleColorChannelImages(camera, numberOfTrainingImages + cameraTestingImageIndex)
                     multipleColorsImage = getMultipleColorsImage(singleColorChannelImages)
 
                     imagePrnuEstimateNpArray = getImagePrnuEstimateNpArray(singleColorChannelImages, multipleColorsImage, camera)
@@ -171,7 +176,7 @@ for computeExtremes in tqdm(([True] if minColor is None or maxColor is None else
                             if DENOISER != Denoiser.MEAN:
                                 cameraTestingImageNoise = cameraTestingImagesNoise[actualCamera][cameraTestingImageIndex]
                             else:
-                                singleColorChannelImages = getSingleColorChannelImages(camera, numberOfTrainingImages + cameraTestingImageIndex)
+                                singleColorChannelImages = singleColorChannelTestingImages[camera][cameraTestingImageIndex]#getSingleColorChannelImages(camera, numberOfTrainingImages + cameraTestingImageIndex)
                                 multipleColorsImage = getMultipleColorsImage(singleColorChannelImages)
                                 cameraTestingImageNoise = getImagePrnuEstimateNpArray(singleColorChannelImages, multipleColorsImage, camera)
 

datasets/raise/utils.py

@@ -8,6 +8,7 @@ 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()
@@ -89,8 +90,12 @@ def isARawImage(imageFilePath):
 
 def getColorChannel(imageFilePath, color):
     if isARawImage(imageFilePath):
-        with rawpy.imread(imageFilePath) as raw:
-            imageNpArray = getRawColorChannel(raw, color)
+        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))
@@ -145,7 +150,7 @@ def getColorMeans(imagesFileNames, colors, singleColorChannelCropResolution = No
     colorMeans = {}
     for color in colors:
         colorIterativeMean = iterativeMean()
-        for imageFileName in tqdm(imagesFileNames, f'Computing mean of {color.replace("_", " ")} colored images'):
+        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)