summaryrefslogtreecommitdiff
path: root/megapixels/commands
diff options
context:
space:
mode:
Diffstat (limited to 'megapixels/commands')
-rw-r--r--megapixels/commands/cv/_old_files_to_face_rois.py (renamed from megapixels/commands/cv/faces_to_csv.py)4
-rw-r--r--megapixels/commands/cv/embeddings.py100
-rw-r--r--megapixels/commands/cv/files_to_rois.py (renamed from megapixels/commands/cv/faces_to_csv_indexed.py)4
-rw-r--r--megapixels/commands/cv/rois_to_pose.py (renamed from megapixels/commands/cv/face_pose_to_csv.py)60
-rw-r--r--megapixels/commands/cv/rois_to_vecs.py109
-rw-r--r--megapixels/commands/datasets/file_meta.py84
-rw-r--r--megapixels/commands/datasets/filter_poses.py76
-rw-r--r--megapixels/commands/datasets/lookup.py44
-rw-r--r--megapixels/commands/datasets/sha256.py55
9 files changed, 385 insertions, 151 deletions
diff --git a/megapixels/commands/cv/faces_to_csv.py b/megapixels/commands/cv/_old_files_to_face_rois.py
index 1fd47571..d92cbd74 100644
--- a/megapixels/commands/cv/faces_to_csv.py
+++ b/megapixels/commands/cv/_old_files_to_face_rois.py
@@ -12,8 +12,8 @@ from app.settings import app_cfg as cfg
color_filters = {'color': 1, 'gray': 2, 'all': 3}
@click.command()
-@click.option('-i', '--input', 'opt_dirs_in', required=True, multiple=True,
- help='Input directory')
+@click.option('-i', '--input', 'opt_fp_files', required=True,
+ help='Input file meta CSV')
@click.option('-o', '--output', 'opt_fp_out', required=True,
help='Output CSV')
@click.option('-e', '--ext', 'opt_ext',
diff --git a/megapixels/commands/cv/embeddings.py b/megapixels/commands/cv/embeddings.py
deleted file mode 100644
index 9cb26ae7..00000000
--- a/megapixels/commands/cv/embeddings.py
+++ /dev/null
@@ -1,100 +0,0 @@
-"""
-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 directory')
-@click.option('-r', '--records', 'opt_fp_records', required=True,
- help='Input directory')
-@click.option('-m', '--media', 'opt_fp_media', required=True,
- help='Image directory')
-@click.option('-o', '--output', 'opt_fp_out', required=True,
- help='Output CSV')
-@click.option('--size', 'opt_size',
- type=(int, int), default=(300, 300),
- help='Output image size')
-@click.option('-g', '--gpu', 'opt_gpu', default=0,
- help='GPU index')
-@click.option('--slice', 'opt_slice', type=(int, int), default=(None, None),
- help='Slice list of files')
-@click.option('-f', '--force', 'opt_force', is_flag=True,
- help='Force overwrite file')
-@click.option('-j', '--jitters', 'opt_jitters', default=cfg.DLIB_FACEREC_JITTERS,
- help='Number of jitters')
-@click.option('-p', '--padding', 'opt_padding', default=cfg.DLIB_FACEREC_PADDING,
- help='Percentage padding')
-@click.pass_context
-def cli(ctx, opt_fp_in, opt_fp_records, opt_fp_out, opt_fp_media, opt_size, opt_gpu,
- opt_slice, opt_jitters, opt_padding, opt_force):
- """Converts frames with faces to CSV of rows"""
-
- import sys
- import os
- from os.path import join
- from pathlib import Path
-
- from tqdm import tqdm
- import numpy as np
- import dlib # must keep a local reference for dlib
- import cv2 as cv
- import dlib
- import pandas as pd
-
- from app.utils import logger_utils, file_utils, im_utils
- from app.models.bbox import BBox
- from app.processors import face_recognition
-
- # -------------------------------------------------
- # init here
-
- log = logger_utils.Logger.getLogger()
-
- if not opt_force and Path(opt_fp_out).exists():
- log.error('File exists. Use "-f / --force" to overwite')
- return
-
- # init dlib FR
- facerec = face_recognition.RecognitionDLIB()
-
- # load data
- df_rois = pd.read_csv(opt_fp_in)
- df_records = pd.read_csv(opt_fp_records)
-
- if opt_slice:
- df_rois = df_rois[opt_slice[0]:opt_slice[1]]
- log.info('Processing {:,} rows'.format(len(df_rois)))
- nrows = len(df_rois)
-
- # face vecs
- vecs = []
-
- for roi_idx, row in tqdm(df_rois.iterrows(), total=nrows):
- # make image path
- record_id = int(row['id'])
- df = df_records.iloc[record_id]
- fp_im = join(opt_fp_media, df['subdir'], '{}.{}'.format(df['fn'], df['ext']))
- # load image
- im = cv.imread(fp_im)
- # make bbox
- xywh = [row['x'], row['y'], row['w'] , row['h']]
- bbox = BBox.from_xywh(*xywh)
- # scale to actual image size
- dim = (row['image_width'], row['image_height'])
- bbox_dim = bbox.to_dim(dim)
- # compute vec
- vec = facerec.vec(im, bbox_dim, jitters=opt_jitters, padding=opt_padding)
- vec_str = ','.join([repr(x) for x in vec])
- vecs.append( {'id': row['id'], 'vec': vec_str})
-
- # save data
- file_utils.mkdirs(opt_fp_out)
- df_vecs = pd.DataFrame.from_dict(vecs)
- df_vecs.to_csv(opt_fp_out, index=False)
- log.info('saved {:,} lines to {}'.format(len(df_vecs), opt_fp_out)) \ No newline at end of file
diff --git a/megapixels/commands/cv/faces_to_csv_indexed.py b/megapixels/commands/cv/files_to_rois.py
index ef958f89..1aaf991c 100644
--- a/megapixels/commands/cv/faces_to_csv_indexed.py
+++ b/megapixels/commands/cv/files_to_rois.py
@@ -100,7 +100,7 @@ def cli(ctx, opt_fp_in, opt_dir_media, opt_fp_out, opt_size, opt_detector_type,
data = []
for df_file in tqdm(df_files.itertuples(), total=len(df_files)):
- fp_im = join(opt_dir_media, df_file.subdir, '{}.{}'.format(df_file.fn, df_file.ext))
+ fp_im = join(opt_dir_media, str(df_file.subdir), f'{df_file.fn}.{df_file.ext}')
im = cv.imread(fp_im)
# filter out color or grayscale iamges
@@ -115,7 +115,7 @@ def cli(ctx, opt_fp_in, opt_dir_media, opt_fp_out, opt_size, opt_detector_type,
continue
try:
- bboxes = detector.detect(im, opt_size=opt_size, opt_pyramids=opt_pyramids, opt_largest=opt_largest)
+ bboxes = detector.detect(im, size=opt_size, pyramids=opt_pyramids, largest=opt_largest)
except Exception as e:
log.error('could not detect: {}'.format(fp_im))
log.error('{}'.format(e))
diff --git a/megapixels/commands/cv/face_pose_to_csv.py b/megapixels/commands/cv/rois_to_pose.py
index ca7489de..3877cecf 100644
--- a/megapixels/commands/cv/face_pose_to_csv.py
+++ b/megapixels/commands/cv/rois_to_pose.py
@@ -1,18 +1,15 @@
"""
-Crop images to prepare for training
+Converts ROIs to pose: yaw, roll, pitch
"""
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
-color_filters = {'color': 1, 'gray': 2, 'all': 3}
-
@click.command()
-@click.option('-f', '--files', 'opt_fp_files', required=True,
+@click.option('-i', '--input', 'opt_fp_files', required=True,
help='Input ROI CSV')
@click.option('-r', '--rois', 'opt_fp_rois', required=True,
help='Input ROI CSV')
@@ -27,9 +24,11 @@ color_filters = {'color': 1, 'gray': 2, 'all': 3}
help='Slice list of files')
@click.option('-f', '--force', 'opt_force', is_flag=True,
help='Force overwrite file')
+@click.option('-d', '--display', 'opt_display', is_flag=True,
+ help='Display image for debugging')
@click.pass_context
def cli(ctx, opt_fp_files, opt_fp_rois, opt_dir_media, opt_fp_out, opt_size,
- opt_slice, opt_force):
+ opt_slice, opt_force, opt_display):
"""Converts ROIs to pose: roll, yaw, pitch"""
import sys
@@ -58,6 +57,7 @@ def cli(ctx, opt_fp_files, opt_fp_rois, opt_dir_media, opt_fp_out, opt_size,
face_pose = FacePoseDLIB()
face_landmarks = LandmarksDLIB()
+ # load datra
df_files = pd.read_csv(opt_fp_files)
df_rois = pd.read_csv(opt_fp_rois)
@@ -70,32 +70,54 @@ def cli(ctx, opt_fp_files, opt_fp_rois, opt_dir_media, opt_fp_out, opt_size,
# -------------------------------------------------
# process here
-
- df_roi_groups = df_rois.groupby('index')
- log.debug('processing {:,} groups'.format(len(df_roi_groups)))
+ df_img_groups = df_rois.groupby('image_index')
+ log.debug('processing {:,} groups'.format(len(df_img_groups)))
poses = []
- #for df_roi_group in tqdm(df_roi_groups.itertuples(), total=len(df_roi_groups)):
- for df_roi_group_idx, df_roi_group in tqdm(df_roi_groups):
+ # iterate
+ #for df_roi_group_idx, df_roi_group in tqdm(df_roi_groups):
+ for image_index, df_img_group in tqdm(df_img_groups):
# make fp
- image_index = df_roi_group.image_index.values[0]
+ #image_index = df_roi_group.image_index.values[0]
pds_file = df_files.iloc[image_index]
fp_im = join(opt_dir_media, pds_file.subdir, '{}.{}'.format(pds_file.fn, pds_file.ext))
im = cv.imread(fp_im)
# get bbox
- x = df_roi_group.x.values[0]
- y = df_roi_group.y.values[0]
- w = df_roi_group.w.values[0]
- h = df_roi_group.h.values[0]
+ x = df_img_group.x.values[0]
+ y = df_img_group.y.values[0]
+ w = df_img_group.w.values[0]
+ h = df_img_group.h.values[0]
dim = im.shape[:2][::-1]
bbox = BBox.from_xywh(x, y, w, h).to_dim(dim)
# get pose
landmarks = face_landmarks.landmarks(im, bbox)
- pose = face_pose.pose(landmarks, dim)
- pose['image_index'] = image_index
- poses.append(pose)
+ pose_data = face_pose.pose(landmarks, dim, project_points=opt_display)
+ pose_degrees = pose_data['degrees'] # only keep the degrees data
+
+ # use the project point data if display flag set
+ if opt_display:
+ pts_im = pose_data['points_image']
+ pts_model = pose_data['points_model']
+ pt_nose = pose_data['point_nose']
+ dst = im.copy()
+ face_pose.draw_pose(dst, pts_im, pts_model, pt_nose)
+ face_pose.draw_degrees(dst, pose_degrees)
+ # display to cv window
+ cv.imshow('', dst)
+ while True:
+ k = cv.waitKey(1) & 0xFF
+ if k == 27 or k == ord('q'): # ESC
+ cv.destroyAllWindows()
+ sys.exit()
+ elif k != 255:
+ # any key to continue
+ break
+
+ # add image index and append to result CSV data
+ pose_degrees['image_index'] = image_index
+ poses.append(pose_degrees)
# save date
diff --git a/megapixels/commands/cv/rois_to_vecs.py b/megapixels/commands/cv/rois_to_vecs.py
new file mode 100644
index 00000000..525f4404
--- /dev/null
+++ b/megapixels/commands/cv/rois_to_vecs.py
@@ -0,0 +1,109 @@
+"""
+Converts ROIs to face vector
+"""
+
+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_files', required=True,
+ help='Input file meta CSV')
+@click.option('-r', '--rois', 'opt_fp_rois', required=True,
+ help='Input ROI CSV')
+@click.option('-m', '--media', 'opt_dir_media', required=True,
+ help='Input media directory')
+@click.option('-o', '--output', 'opt_fp_out', required=True,
+ help='Output CSV')
+@click.option('--size', 'opt_size',
+ type=(int, int), default=(300, 300),
+ help='Output image size')
+@click.option('-j', '--jitters', 'opt_jitters', default=cfg.DLIB_FACEREC_JITTERS,
+ help='Number of jitters')
+@click.option('-p', '--padding', 'opt_padding', default=cfg.DLIB_FACEREC_PADDING,
+ help='Percentage padding')
+@click.option('--slice', 'opt_slice', type=(int, int), default=(None, None),
+ help='Slice list of files')
+@click.option('-f', '--force', 'opt_force', is_flag=True,
+ help='Force overwrite file')
+@click.option('-g', '--gpu', 'opt_gpu', default=0,
+ help='GPU index')
+@click.pass_context
+def cli(ctx, opt_fp_files, opt_fp_rois, opt_dir_media, opt_fp_out, opt_size,
+ opt_slice, opt_force, opt_gpu, opt_jitters, opt_padding):
+ """Converts face ROIs to vectors"""
+
+ import sys
+ import os
+ from os.path import join
+ from pathlib import Path
+ from glob import glob
+
+ from tqdm import tqdm
+ import numpy as np
+ import dlib # must keep a local reference for dlib
+ import cv2 as cv
+ import pandas as pd
+
+ from app.models.bbox import BBox
+ from app.utils import logger_utils, file_utils, im_utils
+ from app.processors import face_recognition
+
+
+ # -------------------------------------------------
+ # init here
+
+ log = logger_utils.Logger.getLogger()
+
+ # init face processors
+ facerec = face_recognition.RecognitionDLIB()
+
+ # load data
+ df_file_meta = pd.read_csv(opt_fp_files)
+ df_rois = pd.read_csv(opt_fp_rois)
+
+ if not opt_force and Path(opt_fp_out).exists():
+ log.error('File exists. Use "-f / --force" to overwite')
+ return
+
+ if opt_slice:
+ df_rois = df_rois[opt_slice[0]:opt_slice[1]]
+
+ # -------------------------------------------------
+ # process here
+ df_img_groups = df_rois.groupby('image_index')
+ log.debug('processing {:,} groups'.format(len(df_img_groups)))
+
+ vecs = []
+
+ for image_index, df_img_group in tqdm(df_img_groups):
+ # make fp
+ roi_index = df_img_group.index.values[0]
+ file_meta = df_file_meta.iloc[image_index] # locate image meta
+ fp_im = join(opt_dir_media, file_meta.subdir, '{}.{}'.format(file_meta.fn, file_meta.ext))
+ im = cv.imread(fp_im)
+ # get bbox
+ x = df_img_group.x.values[0]
+ y = df_img_group.y.values[0]
+ w = df_img_group.w.values[0]
+ h = df_img_group.h.values[0]
+ imw = df_img_group.image_width.values[0]
+ imh = df_img_group.image_height.values[0]
+ dim = im.shape[:2][::-1]
+ # get face vector
+ dim = (imw, imh)
+ bbox_dim = BBox.from_xywh(x, y, w, h).to_dim(dim) # convert to int real dimensions
+ # compute vec
+ # padding=opt_padding not yet implemented in 19.16 but merged in master
+ vec = facerec.vec(im, bbox_dim, jitters=opt_jitters)
+ vec_str = ','.join([repr(x) for x in vec]) # convert to string for CSV
+ vecs.append( {'roi_index': roi_index, 'image_index': image_index, 'vec': vec_str})
+
+
+ # save date
+ file_utils.mkdirs(opt_fp_out)
+ df = pd.DataFrame.from_dict(vecs)
+ df.index.name = 'index'
+ df.to_csv(opt_fp_out) \ No newline at end of file
diff --git a/megapixels/commands/datasets/file_meta.py b/megapixels/commands/datasets/file_meta.py
new file mode 100644
index 00000000..e1456f44
--- /dev/null
+++ b/megapixels/commands/datasets/file_meta.py
@@ -0,0 +1,84 @@
+"""
+Begin with this file to process folder of images
+- Converts folders and subdirectories into CSV with file attributes split
+"""
+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()
+
+@click.command()
+@click.option('-i', '--input', 'opt_fp_in', required=True,
+ help='Input directory')
+@click.option('-o', '--output', 'opt_fp_out', required=True,
+ help='Output file for file meta CSV')
+@click.option('--slice', 'opt_slice', type=(int, int), default=(None, None),
+ help='Slice list of files')
+@click.option('--recursive/--no-recursive', 'opt_recursive', is_flag=True, default=False,
+ help='Use glob recursion (slower)')
+@click.option('-t', '--threads', 'opt_threads', default=4,
+ help='Number of threads')
+@click.option('-f', '--force', 'opt_force', is_flag=True,
+ help='Force overwrite file')
+@click.pass_context
+def cli(ctx, opt_fp_in, opt_fp_out, opt_slice, opt_recursive, opt_threads, opt_force):
+ """Multithreading test"""
+
+ 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 pandas as pd
+ from tqdm import tqdm
+ from glob import glob
+
+ from app.utils import file_utils, im_utils
+
+
+ if not opt_force and Path(opt_fp_out).exists():
+ log.error('File exists. Use "-f / --force" to overwite')
+ return
+
+ fp_ims = []
+ log.info(f'Globbing {opt_fp_in}')
+ for ext in ['jpg', 'png']:
+ if opt_recursive:
+ fp_glob = join(opt_fp_in, '**/*.{}'.format(ext))
+ fp_ims += glob(fp_glob, recursive=True)
+ else:
+ fp_glob = join(opt_fp_in, '*.{}'.format(ext))
+ fp_ims += glob(fp_glob)
+
+ if not fp_ims:
+ log.warn('No images. Try with "--recursive"')
+ return
+
+ if opt_slice:
+ fp_ims = fp_ims[opt_slice[0]:opt_slice[1]]
+
+ log.info('Processing {:,} images'.format(len(fp_ims)))
+
+
+ # convert data to dict
+ data = []
+ for i, fp_im in enumerate(tqdm(fp_ims)):
+ fpp_im = Path(fp_im)
+ subdir = str(fpp_im.parent.relative_to(opt_fp_in))
+ data.append( {
+ 'subdir': subdir,
+ 'fn': fpp_im.stem,
+ 'ext': fpp_im.suffix.replace('.','')
+ })
+
+ # save to CSV
+ file_utils.mkdirs(opt_fp_out)
+ df = pd.DataFrame.from_dict(data)
+ df.index.name = 'index'
+ df.to_csv(opt_fp_out) \ No newline at end of file
diff --git a/megapixels/commands/datasets/filter_poses.py b/megapixels/commands/datasets/filter_poses.py
new file mode 100644
index 00000000..304eeff2
--- /dev/null
+++ b/megapixels/commands/datasets/filter_poses.py
@@ -0,0 +1,76 @@
+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()
+
+@click.command()
+@click.option('-i', '--input', 'opt_fp_in', required=True,
+ help='Input directory')
+@click.option('-o', '--output', 'opt_fp_out', required=True,
+ help='Output directory')
+@click.option('--slice', 'opt_slice', type=(int, int), default=(None, None),
+ help='Slice list of files')
+@click.option('-f', '--force', 'opt_force', is_flag=True,
+ help='Force overwrite file')
+@click.option('--yaw', 'opt_yaw', type=(float, float), default=(-25,25),
+ help='Yaw (min, max)')
+@click.option('--roll', 'opt_roll', type=(float, float), default=(-15,15),
+ help='Roll (min, max)')
+@click.option('--pitch', 'opt_pitch', type=(float, float), default=(-10,10),
+ help='Pitch (min, max)')
+@click.option('--drop', 'opt_drop', type=click.Choice(['valid', 'invalid']), default='invalid',
+ help='Drop valid or invalid poses')
+@click.pass_context
+def cli(ctx, opt_fp_in, opt_fp_out, opt_slice, opt_yaw, opt_roll, opt_pitch,
+ opt_drop, opt_force):
+ """Filter out exaggerated poses"""
+
+ 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 pandas as pd
+ from tqdm import tqdm
+ from glob import glob
+
+ from app.utils import file_utils, im_utils
+
+
+ if not opt_force and Path(opt_fp_out).exists():
+ log.error('File exists. Use "-f / --force" to overwite')
+ return
+
+ df_poses = pd.read_csv(opt_fp_in).set_index('index')
+
+ if opt_slice:
+ df_poses = df_poses[opt_slice[0]:opt_slice[1]]
+
+ log.info('Processing {:,} rows'.format(len(df_poses)))
+
+ # extend a new temporary column
+ df_poses['valid'] = [0] * len(df_poses)
+
+ # filter out extreme poses
+ for ds_pose in tqdm(df_poses.itertuples(), total=len(df_poses)):
+ if ds_pose.yaw > opt_yaw[0] and ds_pose.yaw < opt_yaw[1] \
+ and ds_pose.roll > opt_roll[0] and ds_pose.roll < opt_roll[1] \
+ and ds_pose.pitch > opt_pitch[0] and ds_pose.pitch < opt_pitch[1]:
+ df_poses.at[ds_pose.Index, 'valid'] = 1
+
+ # filter out valid/invalid
+ drop_val = 0 if opt_drop == 'valid' else 0 # drop 0's if drop == valid, else drop 1's
+ df_poses_filtered = df_poses.drop(df_poses[df_poses.valid == int()].index, axis=0)
+
+ # drop temp column
+ df_poses_filtered = df_poses_filtered.drop('valid', axis=1)
+
+ # save filtered poses
+ df_poses_filtered.to_csv(opt_fp_out)
+ log.info('Saved {:,} rows'.format(len(df_poses_filtered))) \ No newline at end of file
diff --git a/megapixels/commands/datasets/lookup.py b/megapixels/commands/datasets/lookup.py
new file mode 100644
index 00000000..11f54957
--- /dev/null
+++ b/megapixels/commands/datasets/lookup.py
@@ -0,0 +1,44 @@
+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()
+
+lookup_types = ['image', 'identity']
+
+@click.command()
+@click.option('-i', '--input', 'opt_fp_in', required=True,
+ help='Input CSV file')
+@click.option('-t', '--type', 'opt_type', default='image',
+ type=click.Choice(lookup_types),
+ help='Type of lookup')
+@click.option('-d', '--dataset', 'opt_dataset', required=True,
+ type=cfg.DatasetVar,
+ default=click_utils.get_default(types.Dataset.LFW),
+ show_default=True,
+ help=click_utils.show_help(types.Dataset))
+@click.option('--index', 'opt_index', required=True,
+ help='Index to lookup')
+@click.pass_context
+def cli(ctx, opt_fp_in, opt_fp_out, opt_index):
+ """Display image info"""
+
+ from glob import glob
+ from os.path import join
+ from pathlib import Path
+ import time
+
+ import pandas as pd
+ from tqdm import tqdm
+
+ from app.utils import file_utils, im_utils
+
+ # lookup and index and display all information
+ df = pd.read_csv(opt_fp_in).set_index('index')
+
+
+
+ \ No newline at end of file
diff --git a/megapixels/commands/datasets/sha256.py b/megapixels/commands/datasets/sha256.py
index c04fb504..4c734073 100644
--- a/megapixels/commands/datasets/sha256.py
+++ b/megapixels/commands/datasets/sha256.py
@@ -10,18 +10,18 @@ log = Logger.getLogger()
@click.command()
@click.option('-i', '--input', 'opt_fp_in', required=True,
help='Input directory')
-@click.option('-o', '--output', 'opt_fp_out',
+@click.option('-m', '--media', 'opt_dir_media', required=True,
+ help='Input media directory')
+@click.option('-o', '--output', 'opt_fp_out', required=True,
help='Output directory')
@click.option('--slice', 'opt_slice', type=(int, int), default=(None, None),
help='Slice list of files')
-@click.option('--recursive/--no-recursive', 'opt_recursive', is_flag=True, default=False,
- help='Use glob recursion (slower)')
@click.option('-t', '--threads', 'opt_threads', default=4,
help='Number of threads')
@click.option('-f', '--force', 'opt_force', is_flag=True,
help='Force overwrite file')
@click.pass_context
-def cli(ctx, opt_fp_in, opt_fp_out, opt_slice, opt_recursive, opt_threads, opt_force):
+def cli(ctx, opt_fp_in, opt_dir_media, opt_fp_out, opt_slice, opt_threads, opt_force):
"""Multithreading test"""
from glob import glob
@@ -42,47 +42,46 @@ def cli(ctx, opt_fp_in, opt_fp_out, opt_slice, opt_recursive, opt_threads, opt_f
log.error('File exists. Use "-f / --force" to overwite')
return
- fp_ims = []
- for ext in ['jpg', 'png']:
- if opt_recursive:
- fp_glob = join(opt_fp_in, '**/*.{}'.format(ext))
- fp_ims += glob(fp_glob, recursive=True)
- else:
- fp_glob = join(opt_fp_in, '*.{}'.format(ext))
- fp_ims += glob(fp_glob)
+ df_files = pd.read_csv(opt_fp_in).set_index('index')
if opt_slice:
- fp_ims = fp_ims[opt_slice[0]:opt_slice[1]]
+ df_files = df_files[opt_slice[0]:opt_slice[1]]
- log.info('Processing {:,} images'.format(len(fp_ims)))
+ log.info('Processing {:,} images'.format(len(df_files)))
- pbar = tqdm(total=100)
+
+ # prepare list of images to multithread into sha256s
+ file_objs = []
+ for ds_file in df_files.itertuples():
+ fp_im = join(opt_dir_media, str(ds_file.subdir), f"{ds_file.fn}.{ds_file.ext}")
+ file_objs.append({'fp': fp_im, 'index': ds_file.Index})
+
+ # convert to thread pool
+ pbar = tqdm(total=len(file_objs))
- def as_sha256(fp_im):
+ def as_sha256(file_obj):
pbar.update(1)
- return file_utils.sha256(fp_im)
+ file_obj['sha256'] = file_utils.sha256(file_obj['fp'])
+ return file_obj
# multithread pool
+ pool_file_objs = []
st = time.time()
pool = ThreadPool(opt_threads)
- with tqdm(total=len(fp_ims)) as pbar:
- sha256s = pool.map(as_sha256, fp_ims)
+ with tqdm(total=len(file_objs)) as pbar:
+ pool_file_objs = pool.map(as_sha256, file_objs)
pbar.close()
-
+
# convert data to dict
data = []
- for i, fp_im in enumerate(fp_ims):
- fpp_im = Path(fp_im)
- subdir = str(fpp_im.parent.relative_to(opt_fp_in))
- sha256 = sha256s[i]
+ for pool_file_obj in pool_file_objs:
data.append( {
- 'sha256': sha256,
- 'subdir': subdir,
- 'fn': fpp_im.stem,
- 'ext': fpp_im.suffix.replace('.','')
+ 'sha256': pool_file_obj['sha256'],
+ 'index': pool_file_obj['index']
})
# save to CSV
+ file_utils.mkdirs(opt_fp_out)
df = pd.DataFrame.from_dict(data)
df.to_csv(opt_fp_out, index=False)