Clean datasets/raise/extract_noise.py

datasets/raise/extract_noise.py

@@ -6,6 +6,7 @@ import numpy as np
 from PIL import Image
 import os
 from tqdm import tqdm
+import csv
 
 imagesFolderPath = 'raise_nef_png'
 denoiser = 'wavelet'
@@ -16,35 +17,16 @@ denoise = getattr(skimage.restoration, f'denoise_{denoiser}')
 mean = None
 numberOfImagesInMean = 0
 
-import csv
-import magic
-import re
-
-deviceFiles = {}
-resolutionFiles = {}
 files = {}
-resolutionFiles = {}
 
 with open('RAISE_all.csv') as csvfile:
     reader = csv.DictReader(csvfile)
-    for row in tqdm(list(reader)):
+    for row in tqdm(list(reader), 'CSV parsing'):
         file = row['File'] + '.png'
-        device = row['Device']
-        resolution = row['Image Size']
-        deviceFiles[file] = device
-        #resolutionFiles[file] = resolution
-        image = Image.open(f'{imagesFolderPath}/{file}')
-        #t = magic.from_file(f'{imagesFolderPath}/{file}.png')
-        resolution = image.size
-        resolutionFiles[file] = resolution#re.search('(\d+) x (\d+)', t).groups()
-        #print(device, resolution)
-        #files[f'{file}.png'] = row
+        files[file] = row
 
 imagesFileNames = os.listdir(imagesFolderPath)
-#print(len(imagesFileNames))
-# imageFileName.replace('.png', '')
-imagesFileNames = [imageFileName for imageFileName in imagesFileNames if deviceFiles[imageFileName] == 'Nikon D7000' and resolutionFiles[imageFileName] == (4946, 3278)]
-#print(len(imagesFileNames))
+imagesFileNames = [imageFileName for imageFileName in tqdm(imagesFileNames, 'Filtering images') if files[imageFileName]['Device'] == 'Nikon D7000' and Image.open(f'{imagesFolderPath}/{imageFileName}').size == (4946, 3278)]
 
 def treatImage(imageFileName):
     global mean, numberOfImagesInMean
@@ -60,7 +42,7 @@ def treatImage(imageFileName):
         mean = ((mean * numberOfImagesInMean) + imageNoiseNpArray) / (numberOfImagesInMean + 1)
     numberOfImagesInMean += 1
 
-for imageFileName in tqdm(imagesFileNames):
+for imageFileName in tqdm(imagesFileNames, 'Denoising images'):
     treatImage(imageFileName)
 
 with open(npArrayFilePath, 'wb') as f: