diff options
Diffstat (limited to 'megapixels/commands/cv/resize.py')
| -rw-r--r-- | megapixels/commands/cv/resize.py | 73 |
1 files changed, 47 insertions, 26 deletions
diff --git a/megapixels/commands/cv/resize.py b/megapixels/commands/cv/resize.py index f535c8b6..dcd621b3 100644 --- a/megapixels/commands/cv/resize.py +++ b/megapixels/commands/cv/resize.py @@ -62,9 +62,11 @@ centerings = { help='Crop focal point') @click.option('--slice', 'opt_slice', type=(int, int), default=(None, None), help='Slice the input list') +@click.option('-t', '--threads', 'opt_threads', default=8, + help='Number of threads') @click.pass_context def cli(ctx, opt_dir_in, opt_dir_out, opt_glob_ext, opt_size, opt_scale_method, - opt_equalize, opt_sharpen, opt_center, opt_slice): + opt_equalize, opt_sharpen, opt_center, opt_slice, opt_threads): """Crop, mirror images""" import os @@ -72,6 +74,8 @@ def cli(ctx, opt_dir_in, opt_dir_out, opt_glob_ext, opt_size, opt_scale_method, from pathlib import Path from glob import glob from tqdm import tqdm + from multiprocessing.dummy import Pool as ThreadPool + from functools import partial from app.utils import logger_utils, file_utils, im_utils @@ -80,46 +84,63 @@ def cli(ctx, opt_dir_in, opt_dir_out, opt_glob_ext, opt_size, opt_scale_method, log = logger_utils.Logger.getLogger() - centering = centerings[opt_center] # ------------------------------------------------- # process here + def pool_resize(fp_im, opt_size, scale_method, centering): + # Threaded image resize function + try: + pbar.update(1) + try: + im = Image.open(fp_im).convert('RGB') + im.verify() + except Exception as e: + log.warn('Could not open: {}'.format(fp_im)) + log.error(e) + return False + + im = ImageOps.fit(im, opt_size, method=scale_method, centering=centering) + + if opt_equalize: + im_np = im_utils.pil2np(im) + im_np_eq = eq_hist_yuv(im_np) + im_np = cv.addWeighted(im_np_eq, 0.35, im_np, 0.65, 0) + im = im_utils.np2pil(im_np) + + if opt_sharpen: + im = im.filter(ImageFilter.UnsharpMask) + + fp_out = join(opt_dir_out, Path(fp_im).name) + im.save(fp_out) + return True + except: + return False + + centering = centerings[opt_center] + scale_method = methods[opt_scale_method] + # get list of files to process fp_ims = glob(join(opt_dir_in, '*.{}'.format(opt_glob_ext))) if opt_slice: fp_ims = fp_ims[opt_slice[0]:opt_slice[1]] log.info('processing {:,} files'.format(len(fp_ims))) - # set scale method - scale_method = methods[opt_scale_method] - + # ensure output dir exists file_utils.mkdirs(opt_dir_out) - # resize and save images - for fp_im in tqdm(fp_ims): - try: - im = Image.open(fp_im).convert('RGB') - im.verify() - except Exception as e: - log.warn('Could not open: {}'.format(fp_im)) - log.error(e) - continue - - im = ImageOps.fit(im, opt_size, method=scale_method, centering=centering) + # setup multithreading + pbar = tqdm(total=len(fp_ims)) + pool_resize = partial(pool_resize, opt_size=opt_size, scale_method=scale_method, centering=centering) + #result_list = pool.map(prod_x, data_list) + pool = ThreadPool(opt_threads) + with tqdm(total=len(fp_ims)) as pbar: + results = pool.map(pool_resize, fp_ims) + pbar.close() - if opt_equalize: - im_np = im_utils.pil2np(im) - im_np_eq = eq_hist_yuv(im_np) - im_np = cv.addWeighted(im_np_eq, 0.35, im_np, 0.65, 0) - im = im_utils.np2pil(im_np) + log.info('Resized: {} / {} images'.format(results.count(True), len(fp_ims))) - if opt_sharpen: - im = im.filter(ImageFilter.UnsharpMask) - - fp_out = join(opt_dir_out, Path(fp_im).name) - im.save(fp_out) def eq_hist_yuv(im): |