Created
September 20, 2020 08:14
-
-
Save mayankgrwl97/34e6ef1091881cb2045bf9aa7dbbf382 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
In order to run multiprocessing on a torch model, replace
with
Reference: https://stackoverflow.com/questions/48822463/how-to-use-pytorch-multiprocessing