diff options
| author | adamhrv <adam@ahprojects.com> | 2019-01-03 12:51:31 +0100 |
|---|---|---|
| committer | adamhrv <adam@ahprojects.com> | 2019-01-03 12:51:31 +0100 |
| commit | f215db6e84071077082d14f8366ae1cf1aea500f (patch) | |
| tree | 33e4573eb618f21685809cf567fdf196ff673a91 /megapixels/commands | |
| parent | 5340bee951c18910fd764241945f1f136b5a22b4 (diff) | |
fix roi index, clean up pose, roi, records, vector
Diffstat (limited to 'megapixels/commands')
| -rw-r--r-- | megapixels/commands/cv/face_pose.py | 87 | ||||
| -rw-r--r-- | megapixels/commands/cv/face_pose_mt.py | 138 | ||||
| -rw-r--r-- | megapixels/commands/cv/face_roi.py | 21 | ||||
| -rw-r--r-- | megapixels/commands/cv/face_vector.py | 44 | ||||
| -rw-r--r-- | megapixels/commands/cv/face_vector_mt.py | 118 | ||||
| -rw-r--r-- | megapixels/commands/datasets/file_record.py (renamed from megapixels/commands/datasets/records.py) | 89 | ||||
| -rw-r--r-- | megapixels/commands/datasets/msceleb.py | 66 | ||||
| -rw-r--r-- | megapixels/commands/datasets/msceleb_names.py | 57 | ||||
| -rw-r--r-- | megapixels/commands/demo/face_search.py | 4 |
9 files changed, 524 insertions, 100 deletions
diff --git a/megapixels/commands/cv/face_pose.py b/megapixels/commands/cv/face_pose.py index c37d006f..9979ad34 100644 --- a/megapixels/commands/cv/face_pose.py +++ b/megapixels/commands/cv/face_pose.py @@ -1,5 +1,9 @@ """ Converts ROIs to pose: yaw, roll, pitch +pitch: looking down or up in yes gesture +roll: tilting head towards shoulder +yaw: twisting head left to right in no gesture + """ import click @@ -17,7 +21,7 @@ from app.settings import app_cfg as cfg help='Override enum media directory') @click.option('--data_store', 'opt_data_store', type=cfg.DataStoreVar, - default=click_utils.get_default(types.DataStore.SSD), + default=click_utils.get_default(types.DataStore.HDD), show_default=True, help=click_utils.show_help(types.Dataset)) @click.option('--dataset', 'opt_dataset', @@ -91,50 +95,57 @@ def cli(ctx, opt_fp_in, opt_fp_out, opt_dir_media, opt_data_store, opt_dataset, # store poses and convert to DataFrame poses = [] - # iterate + # iterate groups with file/record index as key for record_index, df_img_group in tqdm(df_img_groups): # make fp ds_record = df_record.iloc[record_index] fp_im = data_store.face(ds_record.subdir, ds_record.fn, ds_record.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] - 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_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['record_index'] = record_index - poses.append(pose_degrees) + # iterate image group dataframe with roi index as key + for roi_index, df_img in df_img_group.iterrows(): + # get bbox + x, y, w, h = df_img.x, df_img.y, df_img.w, df_img.h + dim = (ds_record.width, ds_record.height) + #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_data = face_pose.pose(landmarks, dim) + #pose_degrees = pose_data['degrees'] # only keep the degrees data + #pose_degrees['points_nose'] = pose_data + # use the project point data if display flag set + if opt_display: + dst = im.copy() + face_pose.draw_pose(dst, pose_data['point_nose'], pose_data['points']) + face_pose.draw_degrees(dst, pose_data) + # 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_data['roi_index'] = roi_index + for k, v in pose_data['points'].items(): + pose_data[f'point_{k}_x'] = v[0][0] / dim[0] + pose_data[f'point_{k}_y'] = v[0][1] / dim[1] + pose_data.pop('points') + pose_data['point_nose_x'] = pose_data['point_nose'][0] / dim[0] + pose_data['point_nose_y'] = pose_data['point_nose'][1] / dim[1] + pose_data.pop('point_nose') + poses.append(pose_data) - # save date + # create dataframe file_utils.mkdirs(fp_out) df = pd.DataFrame.from_dict(poses) + # save date df.index.name = 'index' - df.to_csv(fp_out)
\ No newline at end of file + df.to_csv(fp_out) + # save script + cmd_line = ' '.join(sys.argv) + file_utils.write_text(cmd_line, '{}.sh'.format(fp_out))
\ No newline at end of file diff --git a/megapixels/commands/cv/face_pose_mt.py b/megapixels/commands/cv/face_pose_mt.py new file mode 100644 index 00000000..8fef2c2c --- /dev/null +++ b/megapixels/commands/cv/face_pose_mt.py @@ -0,0 +1,138 @@ +""" +Converts ROIs to pose: yaw, roll, pitch +""" + +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', 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('--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('-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_in, opt_fp_out, opt_dir_media, opt_data_store, opt_dataset, opt_size, + opt_slice, opt_force, opt_display): + """Converts ROIs to pose: roll, yaw, pitch""" + + 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.face_landmarks import LandmarksDLIB + from app.processors.face_pose import FacePoseDLIB + from app.models.data_store import DataStore + + # ------------------------------------------------- + # init here + + log = logger_utils.Logger.getLogger() + + # set data_store + data_store = DataStore(opt_data_store, opt_dataset) + + # get filepath out + fp_out = data_store.metadata(types.Metadata.FACE_POSE) if opt_fp_out is None else opt_fp_out + if not opt_force and Path(fp_out).exists(): + log.error('File exists. Use "-f / --force" to overwite') + return + + # init face processors + face_pose = FacePoseDLIB() + face_landmarks = LandmarksDLIB() + + # load filepath data + fp_record = data_store.metadata(types.Metadata.FILE_RECORD) + df_record = pd.read_csv(fp_record).set_index('index') + # load ROI data + fp_roi = data_store.metadata(types.Metadata.FACE_ROI) + df_roi = pd.read_csv(fp_roi).set_index('index') + # slice if you want + if opt_slice: + df_roi = df_roi[opt_slice[0]:opt_slice[1]] + # group by image index (speedup if multiple faces per image) + df_img_groups = df_roi.groupby('record_index') + log.debug('processing {:,} groups'.format(len(df_img_groups))) + + # store poses and convert to DataFrame + poses = [] + + # iterate + for record_index, df_img_group in tqdm(df_img_groups): + # make fp + ds_record = df_record.iloc[record_index] + fp_im = data_store.face(ds_record.subdir, ds_record.fn, ds_record.ext) + im = cv.imread(fp_im) + for roi_id, df_img in df_img_group.iterrows(): + # get bbox + x, y, w, h = df_img.x, df_img.y, df_img.w, df_img.h + 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_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['record_index'] = record_index + poses.append(pose_degrees) + + + # save date + file_utils.mkdirs(fp_out) + df = pd.DataFrame.from_dict(poses) + df.index.name = 'index' + df.to_csv(fp_out)
\ No newline at end of file diff --git a/megapixels/commands/cv/face_roi.py b/megapixels/commands/cv/face_roi.py index a08566a8..a09a1ce0 100644 --- a/megapixels/commands/cv/face_roi.py +++ b/megapixels/commands/cv/face_roi.py @@ -20,7 +20,7 @@ color_filters = {'color': 1, 'gray': 2, 'all': 3} help='Override enum media directory') @click.option('--data_store', 'opt_data_store', type=cfg.DataStoreVar, - default=click_utils.get_default(types.DataStore.SSD), + default=click_utils.get_default(types.DataStore.HDD), show_default=True, help=click_utils.show_help(types.Dataset)) @click.option('--dataset', 'opt_dataset', @@ -52,10 +52,12 @@ color_filters = {'color': 1, 'gray': 2, 'all': 3} help='Filter to keep color or grayscale images (color = keep color') @click.option('--largest/--all-faces', 'opt_largest', is_flag=True, default=True, help='Only keep largest face') +@click.option('--zone', 'opt_zone', default=(0.0, 0.0), type=(float, float), + help='Face center must be located within zone region (0.5 = half width/height)') @click.pass_context def cli(ctx, opt_fp_in, opt_dir_media, opt_fp_out, opt_data_store, opt_dataset, opt_size, opt_detector_type, opt_gpu, opt_conf_thresh, opt_pyramids, opt_slice, opt_display, opt_force, opt_color_filter, - opt_largest): + opt_largest, opt_zone): """Converts frames with faces to CSV of ROIs""" import sys @@ -130,7 +132,7 @@ def cli(ctx, opt_fp_in, opt_dir_media, opt_fp_out, opt_data_store, opt_dataset, continue try: - bboxes = detector.detect(im, size=opt_size, pyramids=opt_pyramids, largest=opt_largest) + bboxes = detector.detect(im, size=opt_size, pyramids=opt_pyramids, largest=opt_largest, zone=opt_zone) except Exception as e: log.error('could not detect: {}'.format(fp_im)) log.error('{}'.format(e)) @@ -142,17 +144,17 @@ def cli(ctx, opt_fp_in, opt_dir_media, opt_fp_out, opt_data_store, opt_dataset, 'x': bbox.x, 'y': bbox.y, 'w': bbox.w, - 'h': bbox.h, - 'image_width': im.shape[1], - 'image_height': im.shape[0]} + 'h': bbox.h + } data.append(roi) + if len(bboxes) == 0: + log.warn(f'no faces in: {fp_im}') # debug display if opt_display and len(bboxes): im_md = im_utils.resize(im, width=min(1200, opt_size[0])) for bbox in bboxes: bbox_dim = bbox.to_dim(im_md.shape[:2][::-1]) - log.debug(f'bbox: {bbox_dim}') cv.rectangle(im_md, bbox_dim.pt_tl, bbox_dim.pt_br, (0,255,0), 3) cv.imshow('', im_md) while True: @@ -168,4 +170,7 @@ def cli(ctx, opt_fp_in, opt_dir_media, opt_fp_out, opt_data_store, opt_dataset, file_utils.mkdirs(fp_out) df = pd.DataFrame.from_dict(data) df.index.name = 'index' - df.to_csv(fp_out)
\ No newline at end of file + df.to_csv(fp_out) + # save script + cmd_line = ' '.join(sys.argv) + file_utils.write_text(cmd_line, '{}.sh'.format(fp_out))
\ No newline at end of file diff --git a/megapixels/commands/cv/face_vector.py b/megapixels/commands/cv/face_vector.py index 7200d73b..7c03205c 100644 --- a/megapixels/commands/cv/face_vector.py +++ b/megapixels/commands/cv/face_vector.py @@ -15,7 +15,7 @@ from app.settings import app_cfg as cfg help='Override enum media directory') @click.option('--data_store', 'opt_data_store', type=cfg.DataStoreVar, - default=click_utils.get_default(types.DataStore.SSD), + default=click_utils.get_default(types.DataStore.HDD), show_default=True, help=click_utils.show_help(types.Dataset)) @click.option('--dataset', 'opt_dataset', @@ -90,36 +90,28 @@ def cli(ctx, opt_fp_out, opt_dir_media, opt_data_store, opt_dataset, opt_size, log.debug('processing {:,} groups'.format(len(df_img_groups))) vecs = [] - - for image_index, df_img_group in tqdm(df_img_groups): + for record_index, df_img_group in tqdm(df_img_groups): # make fp - roi_index = df_img_group.index.values[0] - # log.debug(f'roi_index: {roi_index}, image_index: {image_index}') - ds_file = df_record.loc[roi_index] # locate image meta - #ds_file = df_record.loc['index', image_index] # locate image meta - - fp_im = data_store.face(str(ds_file.subdir), str(ds_file.fn), str(ds_file.ext)) + ds_record = df_record.iloc[record_index] + fp_im = data_store.face(ds_record.subdir, ds_record.fn, ds_record.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, 'record_index': image_index, 'vec': vec_str}) + for roi_index, df_img in df_img_group.iterrows(): + # get bbox + x, y, w, h = df_img.x, df_img.y, df_img.w, df_img.h + dim = (ds_record.width, ds_record.height) + #dim = im.shape[:2][::-1] + # get face vector + 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, 'record_index': record_index, 'vec': vec_str}) - # save date + # create dataframe df = pd.DataFrame.from_dict(vecs) df.index.name = 'index' + # save CSV file_utils.mkdirs(fp_out) df.to_csv(fp_out)
\ No newline at end of file diff --git a/megapixels/commands/cv/face_vector_mt.py b/megapixels/commands/cv/face_vector_mt.py new file mode 100644 index 00000000..412f9806 --- /dev/null +++ b/megapixels/commands/cv/face_vector_mt.py @@ -0,0 +1,118 @@ +""" +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('-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('--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_out, opt_dir_media, opt_data_store, opt_dataset, 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.models.data_store import DataStore + from app.utils import logger_utils, file_utils, im_utils + from app.processors import face_recognition + + + # ------------------------------------------------- + # init here + + log = logger_utils.Logger.getLogger() + # set data_store + data_store = DataStore(opt_data_store, opt_dataset) + + # get filepath out + fp_out = data_store.metadata(types.Metadata.FACE_VECTOR) if opt_fp_out is None else opt_fp_out + if not opt_force and Path(fp_out).exists(): + log.error('File exists. Use "-f / --force" to overwite') + return + + # init face processors + facerec = face_recognition.RecognitionDLIB() + + # load data + fp_record = data_store.metadata(types.Metadata.FILE_RECORD) + df_record = pd.read_csv(fp_record).set_index('index') + fp_roi = data_store.metadata(types.Metadata.FACE_ROI) + df_roi = pd.read_csv(fp_roi).set_index('index') + + if opt_slice: + df_roi = df_roi[opt_slice[0]:opt_slice[1]] + + # ------------------------------------------------- + # process here + df_img_groups = df_roi.groupby('record_index') + log.debug('processing {:,} groups'.format(len(df_img_groups))) + + vecs = [] + for record_index, df_img_group in tqdm(df_img_groups): + # make fp + ds_record = df_record.iloc[record_index] + fp_im = data_store.face(ds_record.subdir, ds_record.fn, ds_record.ext) + im = cv.imread(fp_im) + for roi_index, df_img in df_img_group.iterrows(): + # get bbox + x, y, w, h = df_img.x, df_img.y, df_img.w, df_img.h + imw = df_img.image_width + imh = df_img.image_height + 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, 'record_index': record_index, 'vec': vec_str}) + + + # save date + df = pd.DataFrame.from_dict(vecs) + df.index.name = 'index' + file_utils.mkdirs(fp_out) + df.to_csv(fp_out)
\ No newline at end of file diff --git a/megapixels/commands/datasets/records.py b/megapixels/commands/datasets/file_record.py index b6ef618b..355b22f2 100644 --- a/megapixels/commands/datasets/records.py +++ b/megapixels/commands/datasets/file_record.py @@ -10,7 +10,12 @@ from app.utils.logger_utils import Logger log = Logger.getLogger() -identity_sources = ['subdir', 'subdir_head', 'subdir_tail'] +# Choose part of the filepath that will be used for the person identity +# eg subdirectory "lfw/media/original/batch_1/train/barack_obama/001.jpg" --> [subdir_tail] --> "barack_obama" +# eg subdirectory "lfw/media/original/batch_1/train/barack_obama/001.jpg" --> [subdir_head] --> "batch_1" +# eg subdirectory "lfw/media/original/batch_1/train/barack_obama/001.jpg" --> [subdir] --> "barack_obama" + +identity_sources = ['subdir', 'numeric'] @click.command() @click.option('-i', '--input', 'opt_fp_in', default=None, @@ -21,7 +26,7 @@ identity_sources = ['subdir', 'subdir_head', 'subdir_tail'] help='Override enum media directory') @click.option('--data_store', 'opt_data_store', type=cfg.DataStoreVar, - default=click_utils.get_default(types.DataStore.SSD), + default=click_utils.get_default(types.DataStore.HDD), show_default=True, help=click_utils.show_help(types.Dataset)) @click.option('--dataset', 'opt_dataset', @@ -35,7 +40,8 @@ identity_sources = ['subdir', 'subdir_head', 'subdir_tail'] 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), +@click.option('--identity', 'opt_identity', type=click.Choice(identity_sources), + default='numeric', help='Identity source, blank for no identity') @click.option('--recursive/--no-recursive', 'opt_recursive', is_flag=True, default=False, help='Use glob recursion (slower)') @@ -44,7 +50,7 @@ def cli(ctx, opt_fp_in, opt_fp_out, opt_dataset, opt_data_store, opt_dir_media, opt_identity, opt_force, opt_recursive): """Generates sha256, uuid, and identity index CSV file""" - import sys + import sys, os from glob import glob from os.path import join from pathlib import Path @@ -53,9 +59,11 @@ def cli(ctx, opt_fp_in, opt_fp_out, opt_dataset, opt_data_store, opt_dir_media, import random import uuid + import cv2 as cv import pandas as pd from tqdm import tqdm from glob import glob + from operator import itemgetter from app.models.data_store import DataStore from app.utils import file_utils, im_utils @@ -91,15 +99,26 @@ def cli(ctx, opt_fp_in, opt_fp_out, opt_dataset, opt_data_store, opt_dir_media, pbar = tqdm(total=len(fp_ims)) - def as_sha256(fp_im): + def pool_mapper(fp_im): pbar.update(1) - return file_utils.sha256(fp_im) + sha256 = file_utils.sha256(fp_im) + im = cv.imread(fp_im) + w, h = im.shape[:2][::-1] + file_size_kb = os.stat(fp_im).st_size // 1000 + num_channels = im_utils.num_channels(im) + return { + 'width': w, + 'height': h, + 'sha256': sha256, + 'file_size_kb': file_size_kb, + 'num_channels': num_channels + } # convert to thread pool - sha256s = [] # ? + pool_maps = [] # ? pool = ThreadPool(opt_threads) with tqdm(total=len(fp_ims)) as pbar: - sha256s = pool.map(as_sha256, fp_ims) + pool_maps = pool.map(pool_mapper, fp_ims) pbar.close() @@ -108,11 +127,12 @@ def cli(ctx, opt_fp_in, opt_fp_out, opt_dataset, opt_data_store, opt_dir_media, data = [] indentity_count = 0 - for sha256, fp_im in zip(sha256s, fp_ims): + for pool_map, fp_im in zip(pool_maps, fp_ims): fpp_im = Path(fp_im) subdir = str(fpp_im.parent.relative_to(fp_in)) + #subdir = '' if subdir is '.' else subdir + log.debug(subdir) - if opt_identity: subdirs = subdir.split('/') if not len(subdirs) > 0: @@ -120,27 +140,46 @@ def cli(ctx, opt_fp_in, opt_fp_out, opt_dataset, opt_data_store, opt_dir_media, 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 + elif opt_identity == 'numeric': + identity = indentity_count # use incrementing number + indentity_count += 1 else: - identity = indentity_count # use incrementing number - indentity_count += 1 + identity = '' data.append({ 'subdir': subdir, + 'num_channels': pool_map['num_channels'], 'fn': fpp_im.stem, 'ext': fpp_im.suffix.replace('.',''), - 'sha256': sha256, + 'sha256': pool_map['sha256'], 'uuid': uuid.uuid4(), - 'identity_key': identity + 'identity_key': identity, + 'width': pool_map['width'], + 'height': pool_map['height'] }) + # create dataframe df_records = pd.DataFrame.from_dict(data) + + 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 {len(df_records)} rows to "{fp_out}"') + # save script + cmd_line = ' '.join(sys.argv) + file_utils.write_text(cmd_line, '{}.sh'.format(fp_out)) + + +''' +# create dataframe + df_records = pd.DataFrame.from_dict(data) + + # add identity key (used for associating identity) if opt_identity: - log.info(f'adding identity index using: "{opt_identity}". This may take a while...') + log.info(f'adding identity index using: "{opt_identity}" subdirectory') # convert dict to DataFrame # sort based on identity_key df_records = df_records.sort_values(by=['identity_key'], ascending=True) @@ -149,19 +188,17 @@ def cli(ctx, opt_fp_in, opt_fp_out, opt_dataset, opt_data_store, opt_dir_media, # populate the identity_index df_records_identity_groups = df_records.groupby('identity_key') # enumerate groups to create identity indices + log.info(f'updating records with identity_key. This may take a while...') + st = time.time() 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) + log.debug('update time: {:.2f}s'.format(time.time() - st)) else: # name everyone person 1, 2, 3... + df_records = df_records.sort_values(by=['subdir'], ascending=True) pass - - 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}')
\ No newline at end of file +'''
\ No newline at end of file diff --git a/megapixels/commands/datasets/msceleb.py b/megapixels/commands/datasets/msceleb.py new file mode 100644 index 00000000..969a1df2 --- /dev/null +++ b/megapixels/commands/datasets/msceleb.py @@ -0,0 +1,66 @@ +''' +Converts MsCelebV1-ImageThumbnails.part.00.tsv to names and images +''' +import click + +from app.settings import types +from app.models.dataset import Dataset +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='Path to input TSV file') +@click.option('-o', '--output', 'opt_fp_out', required=True, + help='Output path for images') +@click.option('--slice', 'opt_slice', type=(int, int), default=(None, None), + help='Slice list of files') +@click.pass_context +def cli(ctx, opt_fp_in, opt_fp_out, opt_slice): + """Converts MSCeleb TSV to images""" + + import sys + import os + from glob import glob + from os.path import join + from pathlib import Path + import time + import base64 + from io import BytesIO + + import pandas as pd + import cv2 as cv + from PIL import Image + from tqdm import tqdm + + from app.utils import file_utils, im_utils + from app.models.data_store import DataStore + + + log = Logger.getLogger() + log.debug(f'opening "{opt_fp_in}" ...') + try: + n_lines = sum(1 for line in open(opt_fp_in)) + except: + n_lines = 1 + + log.debug('{:,}'.format(n_lines)) + + with open(opt_fp_in, 'rb') as fp: + for data_line in tqdm(fp, total=n_lines): + try: + freebase_mid, query_name, search_rank, url_image, url_page, b64_bytes = data_line.decode().split('\t') + # decode image + im64 = base64.b64decode(b64_bytes) + im = Image.open(BytesIO(im64)) + # save image + dir_out = join(opt_fp_out, freebase_mid) + Path(dir_out).mkdir(parents=True, exist_ok=True) + idx = len(os.listdir(dir_out)) + fp_out = join(dir_out, '{}.jpg'.format(file_utils.zpad(idx))) + im.save(fp_out, quality=100) + except Exception as e: + log.error('Could not process: {}, {}. Error: {}'.format(query_name, url_image, e)) diff --git a/megapixels/commands/datasets/msceleb_names.py b/megapixels/commands/datasets/msceleb_names.py new file mode 100644 index 00000000..6ee2ad9a --- /dev/null +++ b/megapixels/commands/datasets/msceleb_names.py @@ -0,0 +1,57 @@ +''' +Converts MsCelebV1-ImageThumbnails.part.00.tsv to names and images +''' +import click + +from app.settings import types +from app.models.dataset import Dataset +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='Path to input TSV file') +@click.option('-o', '--output', 'opt_fp_out', required=True, + help='Output path for images') +@click.pass_context +def cli(ctx, opt_fp_in, opt_fp_out): + """Converts MSCeleb TSV to names file with image count""" + + import sys + import os + from glob import glob + from os.path import join + from pathlib import Path + import time + import base64 + from io import BytesIO + + import pandas as pd + import cv2 as cv + from PIL import Image + from tqdm import tqdm + + from app.utils import file_utils, im_utils + from app.models.data_store import DataStore + + + log = Logger.getLogger() + log.debug(f'opening "{opt_fp_in}" ...') + n_lines = sum(1 for line in open(opt_fp_in)) + log.debug('{:,}'.format(n_lines)) + + with open(opt_fp_in, 'rb') as fp: + for data_line in tqdm(fp, total=n_lines): + freebase_mid, query_name, search_rank, url_image, url_page, b64_bytes = data_line.decode().split('\t') + # decode image + im64 = base64.b64decode(b64_bytes) + im = Image.open(BytesIO(im64)) + # save image + dir_out = join(opt_fp_out, freebase_mid) + Path(dir_out).mkdir(parents=True, exist_ok=True) + idx = len(os.listdir(dir_out)) + fp_out = join(dir_out, '{}.jpg'.format(file_utils.zpad(idx))) + im.save(fp_out, quality=100) diff --git a/megapixels/commands/demo/face_search.py b/megapixels/commands/demo/face_search.py index 6e4bcdad..ca0b8016 100644 --- a/megapixels/commands/demo/face_search.py +++ b/megapixels/commands/demo/face_search.py @@ -15,7 +15,7 @@ log = Logger.getLogger() type=cfg.DataStoreVar, default=click_utils.get_default(types.DataStore.SSD), show_default=True, - help=click_utils.show_help(types.Dataset)) + help=click_utils.show_help(types.DataStore)) @click.option('--dataset', 'opt_dataset', type=cfg.DatasetVar, required=True, @@ -26,7 +26,7 @@ log = Logger.getLogger() @click.option('--gpu', 'opt_gpu', default=0, help='GPU index (use -1 for CPU') @click.pass_context -def cli(ctx, opt_fp_in, opt_data_store, opt_dataset, opt_gpu): +def cli(ctx, opt_fp_in, opt_data_store, opt_dataset, opt_results, opt_gpu): """Display image info""" import sys |
