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
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
|
'''
'''
import click
from app.settings import types
from app.utils import click_utils
from app.settings import app_cfg as cfg
from app.utils.logger_utils import Logger
log = Logger.getLogger()
identity_sources = ['subdir', 'subdir_head', 'subdir_tail']
@click.command()
@click.option('-i', '--input', 'opt_fp_in', default=None,
help='Override enum input filename CSV')
@click.option('-o', '--output', 'opt_fp_out', default=None,
help='Override enum output filename CSV')
@click.option('-m', '--media', 'opt_dir_media', default=None,
help='Override enum media directory')
@click.option('--data_store', 'opt_data_store',
type=cfg.DataStoreVar,
default=click_utils.get_default(types.DataStore.SSD),
show_default=True,
help=click_utils.show_help(types.Dataset))
@click.option('--dataset', 'opt_dataset',
type=cfg.DatasetVar,
required=True,
show_default=True,
help=click_utils.show_help(types.Dataset))
@click.option('--slice', 'opt_slice', type=(int, int), default=(None, None),
help='Slice list of files')
@click.option('-t', '--threads', 'opt_threads', default=12,
help='Number of threads')
@click.option('-f', '--force', 'opt_force', is_flag=True,
help='Force overwrite file')
@click.option('--identity', 'opt_identity', default=None, type=click.Choice(identity_sources),
help='Identity source, blank for no identity')
@click.option('--recursive/--no-recursive', 'opt_recursive', is_flag=True, default=False,
help='Use glob recursion (slower)')
@click.pass_context
def cli(ctx, opt_fp_in, opt_fp_out, opt_dataset, opt_data_store, opt_dir_media, opt_slice, opt_threads,
opt_identity, opt_force, opt_recursive):
"""Generates sha256, uuid, and identity index CSV file"""
import sys
from glob import glob
from os.path import join
from pathlib import Path
import time
from multiprocessing.dummy import Pool as ThreadPool
import random
import uuid
import pandas as pd
from tqdm import tqdm
from glob import glob
from app.models.data_store import DataStore
from app.utils import file_utils, im_utils
# set data_store
data_store = DataStore(opt_data_store, opt_dataset)
# get filepath out
fp_out = data_store.metadata(types.Metadata.FILE_RECORD) if opt_fp_out is None else opt_fp_out
# exit if exists
if not opt_force and Path(fp_out).exists():
log.error('File exists. Use "-f / --force" to overwite')
return
# ----------------------------------------------------------------
# glob files
fp_in = opt_fp_in if opt_fp_in is not None else data_store.media_images_original()
log.info(f'Globbing {fp_in}')
fp_ims = file_utils.glob_multi(fp_in, ['jpg', 'png'], recursive=opt_recursive)
# fail if none
if not fp_ims:
log.error('No images. Try with "--recursive"')
return
# slice to reduce
if opt_slice:
fp_ims = fp_ims[opt_slice[0]:opt_slice[1]]
log.info('Found {:,} images'.format(len(fp_ims)))
# ----------------------------------------------------------------
# multithread process into SHA256
pbar = tqdm(total=len(fp_ims))
def as_sha256(fp_im):
pbar.update(1)
return file_utils.sha256(fp_im)
# convert to thread pool
sha256s = [] # ?
pool = ThreadPool(opt_threads)
with tqdm(total=len(fp_ims)) as pbar:
sha256s = pool.map(as_sha256, fp_ims)
pbar.close()
# ----------------------------------------------------------------
# convert data to dict
data = []
for sha256, fp_im in zip(sha256s, fp_ims):
fpp_im = Path(fp_im)
subdir = str(fpp_im.parent.relative_to(fp_in))
if opt_identity:
subdirs = subdir.split('/')
if not len(subdirs) > 0:
log.error(f'Could not split subdir: "{subdir}. Try different option for "--identity"')
log.error('exiting')
return
if opt_identity == 'subdir':
identity = subdirs[0] # use first/only part
elif opt_identity == 'subdir_head':
identity = subdirs[0] # use first part of subdir path
elif opt_identity == 'subdir_tail':
identity = subdirs[-1] # use last part of subdir path
else:
identity = ''
data.append({
'subdir': subdir,
'fn': fpp_im.stem,
'ext': fpp_im.suffix.replace('.',''),
'sha256': sha256,
'uuid': uuid.uuid4(),
'identity_key': identity
})
log.info(f'adding identity index using: "{opt_identity}". This may take a while...')
# convert dict to DataFrame
df_records = pd.DataFrame.from_dict(data)
# sort based on identity_key
df_records = df_records.sort_values(by=['identity_key'], ascending=True)
# add new column for identity
df_records['identity_index'] = [-1] * len(df_records)
# populate the identity_index
df_records_identity_groups = df_records.groupby('identity_key')
# enumerate groups to create identity indices
for identity_index, df_records_identity_group_tuple in enumerate(df_records_identity_groups):
identity_key, df_records_identity_group = df_records_identity_group_tuple
for ds_record in df_records_identity_group.itertuples():
df_records.at[ds_record.Index, 'identity_index'] = identity_index
# reset index after being sorted
df_records = df_records.reset_index(drop=True)
df_records.index.name = 'index' # reassign 'index' as primary key column
# write to CSV
file_utils.mkdirs(fp_out)
df_records.to_csv(fp_out)
# done
log.info(f'wrote rows: {len(df_records)} to {fp_out}')
|
