summaryrefslogtreecommitdiff
path: root/megapixels/commands/processor/crop.py
blob: 778be0c45de9b1fd7586abca88a076328275ca58 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
"""
Crop images to prepare for training
"""

import click
from PIL import Image, ImageOps, ImageFilter, ImageDraw

from app.settings import types
from app.utils import click_utils
from app.settings import app_cfg as cfg

@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_ext',
  default='jpg', type=click.Choice(['jpg', 'png']),
  help='File glob ext')
@click.option('--size', 'opt_size', 
  type=(int, int), default=(256, 256),
  help='Output image size')
@click.option('-t', '--crop-type', 'opt_crop_type',
  default='center', type=click.Choice(['center', 'mirror', 'face', 'person', 'none']),
  help='Force fit image center location')
@click.pass_context
def cli(ctx, opt_dir_in, opt_dir_out, opt_ext, opt_size, opt_crop_type):
  """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
  
  # -------------------------------------------------
  # process here

  log = logger_utils.Logger.getLogger()
  log.info('crop images')

  # get list of files to process
  fp_ims = glob(join(opt_dir_in, '*.{}'.format(opt_ext)))
  log.debug('files: {}'.format(len(fp_ims)))
  
  # ensure output dir exists
  file_utils.mkdirs(opt_dir_out)

  for fp_im in tqdm(fp_ims):
    im = process_crop(fp_im, opt_size, opt_crop_type)
    fp_out = join(opt_dir_out, Path(fp_im).name)
    im.save(fp_out)

  
def process_crop(fp_im, opt_size, crop_type):
  im  = Image.open(fp_im)
  if crop_type == 'center':
    im = crop_square_fit(im, opt_size)
  elif crop_type == 'mirror':
    im = mirror_crop_square(im, opt_size)
  return im

def crop_square_fit(im, size, center=(0.5, 0.5)):
  return ImageOps.fit(im, size, method=Image.BICUBIC, centering=center)

def mirror_crop_square(im, size):
  # force to even dims
  if im.size[0] % 2 or im.size[1] % 2:
    im = ImageOps.fit(im, ((im.size[0] // 2) * 2, (im.size[1] // 2) * 2))

  # create new square image
  min_size, max_size = (min(im.size), max(im.size))
  orig_w, orig_h = im.size
  margin = (max_size - min_size) // 2
  w, h = (max_size, max_size)
  im_new = Image.new('RGB', (w, h), color=(0, 0, 0))
  
  #crop (l, t, r, b)
  if orig_w > orig_h:
    # landscape, mirror expand T/B
    im_top = ImageOps.mirror(im.crop((0, 0, margin, w)))
    im_bot = ImageOps.mirror(im.crop((orig_h - margin, 0, orig_h, w)))
    im_new.paste(im_top, (0, 0))
    im_new.paste(im, (margin, 0, orig_h + margin, w))
    im_new.paste(im_bot, (h - margin, 0))
  elif orig_h > orig_w:
    # portrait, mirror expand L/R
    im_left = ImageOps.mirror(im.crop((0, 0, margin, h)))
    im_right = ImageOps.mirror(im.crop((orig_w - margin, 0, orig_w, h)))
    im_new.paste(im_left, (0, 0))
    im_new.paste(im, (margin, 0, orig_w + margin, h))
    im_new.paste(im_right, (w - margin, 0))

  return im_new.resize(size)

  
def center_crop_face():
  pass

def center_crop_person():
  pass