Make parallel production ready

datasets/raise/extract_noise.py

@@ -7,28 +7,43 @@ from PIL import Image
 import os
 from tqdm import tqdm
 from multiprocessing import Pool
+from threading import Lock
 
-imagesFolderPath = 'tif'
+imagesFolderPath = 'flat-field'
+npArrayFilePath = 'mean.npy'
 
-'''
-for imageFileName in tqdm(os.listdir(imagesFolderPath)):
-	npArrayFilePath = f'mean.npy'
-	if os.path.isfile(npArrayFilePath):
-		continue
+mutex = Lock()
+mean = None#np.zeros((3264, 4928, 3))
+numberOfImagesInMean = 0
+
+imagesFileNames = os.listdir(imagesFolderPath)
+
+pbar = tqdm(total = len(imagesFileNames))
+
+def treatImage(imageFileName):
+	global mean, numberOfImagesInMean, pbar
 	imageFilePath = f'{imagesFolderPath}/{imageFileName}'
 	imagePil = Image.open(imageFilePath)
 	imageNpArray = img_as_float(np.array(imagePil))
+	print('Before mean computation')
 	imageNoiseNpArray = imageNpArray - denoise_tv_chambolle(imageNpArray, weight=0.2, channel_axis=-1)
-	with open(npArrayFilePath, 'wb') as f:
-		np.save(f, imageNoiseNpArray)
-'''
+	print('After mean computation')
+	with mutex:
+		print('Start mutex section')
+		if mean is None:
+			print('Inter A')
+			mean = imageNoiseNpArray
+		else:
+			print('Inter B')
+			mean = ((mean * numberOfImagesInMean) + imageNoiseNpArray) / (numberOfImagesInMean + 1)
+		print('Intermediary 0')
+		numberOfImagesInMean += 1
+		print('Intermediary 1')
+		pbar.update(numberOfImagesInMean)
+		print('End mutex section')
 
-import time
+with Pool(5) as p:
+	p.map(treatImage, imagesFileNames)
 
-def treatImage(imageFilePath):
-	while True:
-		print(imageFilePath)
-		time.sleep(1)
-
-with Pool(22) as p:
-	p.map(treatImage, )
+with open(npArrayFilePath, 'wb') as f:
+	np.save(f, mean)