update exerciseC

exercises/exerciseC/process_images.py

@@ -3,17 +3,11 @@ import sys
 from multiprocessing import Pool as ProcessPool
 import time
 
-def process_image(fname):
-    n_threads = os.getenv('OMP_NUM_THREADS', '(unset)')
-    print(f"Worker {fname=} OMP_NUM_THREADS={n_threads}")
+def process_image(input_tuple):
 
-    # An image is an array with width, height and three (RGB) color channels
-    # (Sometimes there is a transparency channel too: RGBA)
-    im = Image.open(fname)
-    try:
-        A = np.median(im, axis=2)[::4, ::4]
-    except:
-        A = np.array(im)[::4, ::4]
+    fname, A = input_tuple
+    n_threads = os.getenv('OMP_NUM_THREADS', '(unset)')
+    print(f"Worker {fname=} OMP_NUM_THREADS={n_threads}", flush=True)
 
     # Decompose image
     U, S, Vh = np.linalg.svd(A)
@@ -36,10 +30,10 @@ if __name__ == '__main__':
     fnames = sys.argv[3:]
 
     # Check that the output folders exist, or create them if needed
-    if not os.path.isdir('processed_images'): os.mkdir('processed_images')
     if not os.path.isdir('timings'): os.mkdir('timings')
+    if not os.path.isdir('processed_images'): os.mkdir('processed_images')
 
-    print(f"Controller with {n_processes} processes and {n_threads} threads / worker")
+    print(f"Controller with {n_processes} processes and {n_threads} threads / worker", flush=True)
 
     # The environment that is set in the parent is inherited by child workers,
     # we need to set the variable before numpy is imported!
@@ -50,19 +44,28 @@ if __name__ == '__main__':
     import numpy as np
     from PIL import Image
 
-    # Time the execution
+    # I/O Load the images
+    image_arrays = []
+    for fname in fnames:
+        im = Image.open(fname)
+        A = np.array(im)
+        image_arrays.append((fname, A))
+
+    # Time the execution of the pool map
     start_time = time.time()
     with ProcessPool(n_processes) as p:
-        new_images = p.map(process_image, fnames)
+        new_images = p.map(process_image, image_arrays)
     elapsed_time = time.time() - start_time
 
+    # I/O save the processed images
     for im, fname in zip(new_images, fnames):
         im = Image.fromarray(im)
         im.save(fname.replace('images', 'processed_images'))
 
-    print(f'{n_processes} processes and {n_threads} threads and {len(fnames)} jobs: {elapsed_time}')
+    print(f'{n_processes} processes and {n_threads} threads and {len(fnames)} jobs: {elapsed_time}\n',
+          flush=True)
 
-    # IO: Save the timing to a unique txt file
+    # I/O: Save the timing to a unique txt file
     filename = f'timings/{n_processes:02}_processes_{n_threads:02}_threads.txt'
     with open(filename, 'w') as file:
         file.write(f'{n_processes} {n_threads} {elapsed_time:.6f}')