diff options
Diffstat (limited to 'megapixels/commands/cv/resize.py')
| -rw-r--r-- | megapixels/commands/cv/resize.py | 128 |
1 files changed, 128 insertions, 0 deletions
diff --git a/megapixels/commands/cv/resize.py b/megapixels/commands/cv/resize.py new file mode 100644 index 00000000..f535c8b6 --- /dev/null +++ b/megapixels/commands/cv/resize.py @@ -0,0 +1,128 @@ +""" +Crop images to prepare for training +""" + +import click +import cv2 as cv +from PIL import Image, ImageOps, ImageFilter + +from app.settings import types +from app.utils import click_utils +from app.settings import app_cfg as cfg + +""" +Filter Q-Down Q-Up Speed +NEAREST ⭐⭐⭐⭐⭐ +BOX ⭐ ⭐⭐⭐⭐ +BILINEAR ⭐ ⭐ ⭐⭐⭐ +HAMMING ⭐⭐ ⭐⭐⭐ +BICUBIC ⭐⭐⭐ ⭐⭐⭐ ⭐⭐ +LANCZOS ⭐⭐⭐⭐ ⭐⭐⭐⭐ ⭐ +""" +methods = { + 'lanczos': Image.LANCZOS, + 'bicubic': Image.BICUBIC, + 'hamming': Image.HAMMING, + 'bileaner': Image.BILINEAR, + 'box': Image.BOX, + 'nearest': Image.NEAREST + } +centerings = { + 'tl': (0.0, 0.0), + 'tc': (0.5, 0.0), + 'tr': (0.0, 0.0), + 'lc': (0.0, 0.5), + 'cc': (0.5, 0.5), + 'rc': (1.0, 0.5), + 'bl': (0.0, 1.0), + 'bc': (1.0, 0.5), + 'br': (1.0, 1.0) +} + +@click.command() +@click.option('-i', '--input', 'opt_dir_in', required=True, + help='Input directory') +@click.option('-o', '--output', 'opt_dir_out', required=True, + help='Output directory') +@click.option('-e', '--ext', 'opt_glob_ext', + default='png', type=click.Choice(['jpg', 'png']), + help='File glob ext') +@click.option('--size', 'opt_size', + type=(int, int), default=(256, 256), + help='Output image size (square)') +@click.option('--method', 'opt_scale_method', + type=click.Choice(methods.keys()), + default='lanczos', + help='Scaling method to use') +@click.option('--equalize', 'opt_equalize', is_flag=True, + help='Equalize historgram') +@click.option('--sharpen', 'opt_sharpen', is_flag=True, + help='Unsharp mask') +@click.option('--center', 'opt_center', default='cc', type=click.Choice(centerings.keys()), + help='Crop focal point') +@click.option('--slice', 'opt_slice', type=(int, int), default=(None, None), + help='Slice the input list') +@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): + """Crop, mirror images""" + + import os + from os.path import join + from pathlib import Path + from glob import glob + from tqdm import tqdm + + from app.utils import logger_utils, file_utils, im_utils + + # ------------------------------------------------- + # init + + log = logger_utils.Logger.getLogger() + + centering = centerings[opt_center] + + # ------------------------------------------------- + # process here + + # 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) + + 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) + + +def eq_hist_yuv(im): + im_yuv = cv.cvtColor(im, cv.COLOR_BGR2YUV) + im_yuv[:,:,0] = cv.equalizeHist(im_yuv[:,:,0]) + return cv.cvtColor(im_yuv, cv.COLOR_YUV2BGR) |
