parallel_png_to_jpg.py

import argparse
import glob
import multiprocessing
import os
from functools import partial

import cv2
from tqdm import tqdm


def png_to_jpg(img_png_path, jpg_dir):
    img = cv2.imread(img_png_path)
    img_basename = os.path.splitext(os.path.basename(img_png_path))[0]
    img_jpg_path = os.path.join(jpg_dir, img_basename+'.jpg')
    cv2.imwrite(img_jpg_path, img)
    return img_jpg_path


def process(img_png_paths, jpg_dir, n_workers):
    png_to_jpg_fn = partial(png_to_jpg, jpg_dir=jpg_dir)
    with multiprocessing.Pool(n_workers) as pool:
        for img_jpg_path in tqdm(pool.imap_unordered(png_to_jpg_fn, img_png_paths))


def get_args():
    parser = argparse.ArgumentParser()

    parser.add_argument('--png_dir', type=str, required=True)
    parser.add_argument('--jpg_dir', type=str, required=True)
    parser.add_argument('--n_workers', type=int, default=1)

    args = parser.parse_args()

    assert os.path.exists(args.png_dir) and os.path.isdir(args.png_dir)
    os.makedirs(args.jpg_dir, exist_ok=True)

    return args


if __name__ == '__main__':
    args = get_args()
    img_png_paths = sorted(glob.glob(os.path.join(args.png_dir, '*.png')))
    process(img_png_paths, args.jpg_dir, args.n_workers)

In order to run multiprocessing on a torch model, replace

from multiprocessing import Process, Pool

with

from torch.multiprocessing import Pool, Process, set_start_method
try:
     set_start_method('spawn')
except RuntimeError:
    pass

Reference: https://stackoverflow.com/questions/48822463/how-to-use-pytorch-multiprocessing