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
|
"""
Crop images to prepare for training
"""
import click
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_fp_in', required=True,
help='Input CSV')
@click.option('--media', 'opt_dir_media', required=True,
help='Input image/video directory')
@click.option('-o', '--output', 'opt_dir_out', required=True,
help='Output directory for extracted ROI images')
@click.option('--size', 'opt_size',
type=(int, int), default=(300, 300),
help='Output image size')
@click.option('--slice', 'opt_slice', type=(int, int), default=(None, None),
help='Slice list of files')
@click.option('--padding', 'opt_padding', default=0,
help='Facial padding')
@click.option('--ext', 'opt_ext_out', default='jpg', type=click.Choice(['jpg', 'png']),
help='Output image type')
@click.pass_context
def cli(ctx, opt_fp_in, opt_dir_media, opt_dir_out, opt_size, opt_slice,
opt_padding, opt_ext_out):
"""Extrace ROIs to images"""
import os
from os.path import join
from pathlib import Path
from glob import glob
from tqdm import tqdm
import numpy as np
from PIL import Image, ImageOps, ImageFilter, ImageDraw
import cv2 as cv
import pandas as pd
from app.utils import logger_utils, file_utils, im_utils
from app.models.bbox import BBox
# -------------------------------------------------
# process here
log = logger_utils.Logger.getLogger()
df_rois = pd.read_csv(opt_fp_in)
if opt_slice:
df_rois = df_rois[opt_slice[0]:opt_slice[1]]
log.info('Processing {:,} rows'.format(len(df_rois)))
file_utils.mkdirs(opt_dir_out)
df_rois_grouped = df_rois.groupby(['fn']) # group by fn/filename
groups = df_rois_grouped.groups
for group in groups:
# get image
group_rows = df_rois_grouped.get_group(group)
row = group_rows.iloc[0]
fp_im = join(opt_dir_media, '{fn}{ext}'.format(**row)) #TODO change to ext
im = Image.open(fp_im)
for idx, roi in group_rows.iterrows():
log.info('{}'.format(roi['fn']))
# get bbox to im dimensions
xywh = [roi['x'], roi['y'], roi['w'] , roi['h']]
bbox = BBox.from_xywh(*xywh)
dim = im.size
bbox_dim = bbox.to_dim(dim)
# expand
bbox_dim_exp = bbox_dim.expand_dim(opt_padding, dim)
# crop
x1y2 = bbox_dim_exp.pt_tl + bbox_dim_exp.pt_br
im_crop = im.crop(box=x1y2)
# save
idx_zpad = file_utils.zpad(idx, zeros=3)
fp_im_out = join(opt_dir_out, '{}_{}.{}'.format(roi['fn'], idx_zpad, opt_ext_out))
im_crop.save(fp_im_out)
|
