How to do multiprocessing of image augmentations of large quantity?

Question

I'm trying to do image augmentations for which I'm getting images from a large number of folders Doing this in a sequel takes lots of time, so I'm running the same script in different terminals in order to complete augmentations very quickly by providing the start and end value of the list index as shown in the below code.

do_augmentations(args):
   total_count = len(all_folders)
   split_no = total_count //2
   start = split_no
   if split_no  == 0:
    split_no = 1
   end = total_count
   for folder in all_folders[start:end]:
    allImgs = list(paths.list_images(folder))
    count = len(allImgs)
    for img in allImages:
        augmentations(img)
        cv2.imwrite(img)
def main():
   all_folders= os.walk(folderpath)
   do_augmentations(all_folders)

I was wondering if we could use multiple CPU cores with multithreading and multiprocessing packages in Python in parallel rather than sequentially because it takes a long time. Here I'm mentioning the start and end value of the folder number to run separately in multiple terminals to run faster . I tried using a multiprocessing library to implement this in parallel, but it runs in the same sequential manner as before. Below is the code of my approach to solving this.

from multiprocessing import pool
from multiprocessing.dummy import Pool as ThreadPool
do_augmentations(args):
   all_folers = args[0]
   bpath = args[1]
   for folder in all_folders:
    allImgs = list(paths.list_images(folder))
    count = len(allImgs)
    for img in allImages:
        augmentations(img)
        cv2.imwrite(img)
def main():
   bpath = 'img_foldr'
   all_folders= os.walk(folderpath)
   pool = ThreadPool(4)
   pool.map(do_augmentations,[[all_folders,bpath]])

When running this, it does image processing on one folder at a time in a loop instead of in parallel for many folders simultaneously. I don't understand what I'm doing wrong. Any help or suggestions to solve this will be very helpful.

Update: I tried answer given by Jan Wilamowski as below

def augment(image):
   img = do_augmentation(image)
   cv2.imwrite(img)

def main():
   all_folders  = os.walk(imagefolder)
   all_images = chain(paths.list_images(folder) for folder in all_folders)
   pool = Pool(4)
   pool.map(augment,all_images)

I get error as below

Traceback (most recent call last):
File "aug_img.py", line 424, in
main()
File "aug_img.py", line 346, in main
pool.map(augment,all_images)
File "C:\Users\mathew\anaconda3\envs\retinanet\lib\multiprocessing\pool.py", line 364, in map return self._map_async(func, iterable, mapstar, chunksize).get()
File "C:\Users\mathew\anaconda3\envs\retinanet\lib\multiprocessing\pool.py", line 771, in get raise self._value
File "C:\Users\mathew\anaconda3\envs\retinanet\lib\multiprocessing\pool.py", line 537, in _handle_tasks put(task)
File "C:\Users\mathew\anaconda3\envs\retinanet\lib\multiprocessing\connection.py", line 206, in send
self._send_bytes(_ForkingPickler.dumps(obj))
File "C:\Users\mathew\anaconda3\envs\retinanet\lib\multiprocessing\reduction.py", line 51, in dumps cls(buf, protocol).dump(obj)
TypeError: cannot pickle 'generator' object

user_15744300 · Answer

Have your function work on a single folder and pass the folder list to pool.map(). Also, use a process pool to avoid problems with the GIL (as pointed out by several commenters):

from multiprocessing import Pool

do_augmentations(folder):
    allImgs = list(paths.list_images(folder))
    count = len(allImgs)
    for img in allImages:
        augmentations(img)
        cv2.imwrite(img)

def main():
    bpath = 'img_foldr'
    all_folders = os.walk(folderpath)
    pool = Pool(4)
    pool.map(do_augmentations, all_folders)

You could also break it down further and have your function work on a single image, giving more consistent performance:

from itertools import chain
from imutils import paths

def augment(image):
    augmentations(image)
    cv2.imwrite(image)

def main():
    all_images = paths.list_images(folderpath)
    pool = Pool(4)
    pool.map(augment, all_images)

Note however that disk I/O can be a bottleneck so don't expect a linear performance improvement.

user_5597354 · Answer

Let me give you my simple multiprocessing recipy for any task not just augmentation.

This is a top down view.

import os 
import multiprocessing
from workers import batch_function

no_of_cpus = os.cpu_count()
all_folders= os.walk(folderpath)
input_list = get_all_the_files_from_folder(all_folders) #This is specific to your file structures

mp_dict= split_input_list(number_of_splits=int(no_of_cpus), input_list= input_list)

pool = multiprocessing.Pool()
results = pool.map(workers.batch_function, mp_dict) # Call Image Augmentations

First I would make a list of all the data that needs to be preprocessed and then split it in the number of process I want by means of a function split_input_list.

If you don't need to return anything from the batch function you don't need the results variable but in essence you will have a list of results from each process that you can iterate by means of for res in results: no matter what you return in batch_function.

def split_input_list(process_number, d):
    dict_list = []
    pn = process_number
    for i in range(pn - 1):
        start = len(d) // pn * i
        finish = len(d) // pn * (i + 1)

        split_dict = dict(list(d.items())[start:finish])
        print(len(split_dict))
        dict_list.append(split_dict)

    last_dict = dict(list(d.items())[finish:])
    dict_list.append(last_dict)

    print(len(last_dict))
    return dict_list

Then in a separate workers.py file I usually have multiple batch_function that can accomplish certain tasks. In this case for augmentations I would do something similar to :

def batch_function(mp_list):

    local_split = mp_list[0]

    for k,_ in local_split.items():
        augmentations(k)

    ...
    #return what_you_need

Also if you don't have an impressive amount of RAM and a CPU with 32 cores expect some crashing from lack of memory.

How to do multiprocessing of image augmentations of large quantity?

2 Answers