""" Dataset model: container for all CSVs about a dataset """ import os import sys from os.path import join from pathlib import Path import logging import pandas as pd import numpy as np from app.settings import app_cfg as cfg from app.settings import types from app.models.bbox import BBox from app.utils import file_utils, im_utils from app.models.data_store import DataStore, DataStoreS3 from app.utils.logger_utils import Logger # ------------------------------------------------------------------------- # Dataset # ------------------------------------------------------------------------- class Dataset: def __init__(self, opt_data_store, opt_dataset_type): self._dataset_type = opt_dataset_type # enum type self.log = Logger.getLogger() self._metadata = {} self._face_vectors = [] self._nullframe = pd.DataFrame() # empty placeholder self.data_store = DataStore(opt_data_store, self._dataset_type) self.data_store_s3 = DataStoreS3(self._dataset_type) def load_face_vectors(self): metadata_type = types.Metadata.FACE_VECTOR fp_csv = self.data_store.metadata(metadata_type) self.log.info(f'loading: {fp_csv}') if Path(fp_csv).is_file(): self._metadata[metadata_type] = pd.read_csv(fp_csv).set_index('index') # convert DataFrame to list of floats self._face_vectors = self.df_vecs_to_dict(self._metadata[metadata_type]) self._face_vector_roi_idxs = self.df_vec_roi_idxs_to_dict(self._metadata[metadata_type]) self.log.info(f'build face vector dict: {len(self._face_vectors)}') # remove the face vector column, it can be several GB of memory self._metadata[metadata_type].drop('vec', axis=1, inplace=True) #n_dims = len(self._metadata[metadata_type].keys()) - 2 #drop_keys = [f'd{i}' for i in range(1,n_dims+1)] #self._metadata[metadata_type].drop(drop_keys, axis=1, inplace=True) else: self.log.error(f'File not found: {fp_csv}. Exiting.') sys.exit() def load_records(self): metadata_type = types.Metadata.FILE_RECORD fp_csv = self.data_store.metadata(metadata_type) self.log.info(f'loading: {fp_csv}') if Path(fp_csv).is_file(): self._metadata[metadata_type] = pd.read_csv(fp_csv, dtype={'fn':str}).set_index('index') else: self.log.error(f'File not found: {fp_csv}. Exiting.') sys.exit() def load_identities(self): metadata_type = types.Metadata.IDENTITY fp_csv = self.data_store.metadata(metadata_type) self.log.info(f'loading: {fp_csv}') if Path(fp_csv).is_file(): self._metadata[metadata_type] = pd.read_csv(fp_csv).set_index('index') else: self.log.error(f'File not found: {fp_csv}. Exiting.') sys.exit() def metadata(self, opt_metadata_type): return self._metadata.get(opt_metadata_type, None) def index_to_record(self, index): # get record meta df_record = self._metadata[types.Metadata.FILE_RECORD] ds_record = df_record.iloc[index] identity_index = ds_record.identity_index # get identity meta df_identity = self._metadata[types.Metadata.IDENTITY] # future datasets can have multiple identities per images ds_identities = df_identity.iloc[identity_index] # get filepath and S3 url fp_im = self.data_store.face(ds_record.subdir, ds_record.fn, ds_record.ext) s3_url = self.data_store_s3.face(ds_record.uuid) image_record = ImageRecord(ds_record, fp_im, s3_url, ds_identities=ds_identities) return image_record def vector_to_record(self, record_index): '''Accumulates image and its metadata''' df_face_vector = self._metadata[types.Metadata.FACE_VECTOR] ds_face_vector = df_face_vector.iloc[vector_index] # get the match's ROI index image_index = ds_face_vector.image_index # get the roi dataframe df_face_roi = self._metadata[types.Metadata.FACE_ROI] ds_roi = df_face_roi.iloc[image_index] # create BBox dim = (ds_roi.image_width, ds_roi.image_height) bbox = BBox.from_xywh_dim(ds_roi.x, ds_roi.y, ds_roi.w, ds_roi.y, dim) # use the ROI index to get identity index from the identity DataFrame df_sha256 = self._metadata[types.Metadata.SHA256] ds_sha256 = df_sha256.iloc[image_index] sha256 = ds_sha256.sha256 # get the local filepath df_filepath = self._metadata[types.Metadata.FILEPATH] ds_file = df_filepath.iloc[image_index] fp_im = self.data_store.face_image(ds_file.subdir, ds_file.fn, ds_file.ext)\ # get remote path df_uuid = self._metadata[types.Metadata.UUID] ds_uuid = df_uuid.iloc[image_index] uuid = ds_uuid.uuid fp_url = self.data_store_s3.face_image(uuid) fp_url_crop = self.data_store_s3.face_image_crop(uuid) image_record = ImageRecord(image_index, sha256, uuid, bbox, fp_im, fp_url) # now get the identity index (if available) identity_index = ds_sha256.identity_index if identity_index > -1: # then use the identity index to get the identity meta df_identity = df_filepath = self._metadata[types.Metadata.IDENTITY] ds_identity = df_identity.iloc[identity_index] # get the name and description name = ds_identity.fullname desc = ds_identity.description gender = ds_identity.gender n_images = ds_identity.images url = '(url)' # TODO age = '(age)' # TODO nationality = '(nationality)' identity = Identity(identity_index, name=name, desc=desc, gender=gender, n_images=n_images, url=url, age=age, nationality=nationality) image_record.identity = identity else: self.log.info(f'no identity index: {ds_sha256}') return image_record def find_matches(self, query_vec, n_results=5, threshold=0.6): image_records = [] # list of image matches w/identity if available # find most similar feature vectors indexes #match_idxs = self.similar(query_vec, n_results, threshold) sim_scores = np.linalg.norm(np.array([query_vec]) - np.array(self._face_vectors), axis=1) match_idxs = np.argpartition(sim_scores, range(n_results))[:n_results] df_vector = self._metadata[types.Metadata.FACE_VECTOR] df_record = self._metadata[types.Metadata.FILE_RECORD] for match_idx in match_idxs: # get the corresponding face vector row roi_index = self._face_vector_roi_idxs[match_idx] record_idx = df_vector.iloc[roi_index].record_index ds_record = df_record.iloc[record_idx] self.log.debug(f'find match index: {match_idx}, --> roi_index: {roi_index}') fp_im = self.data_store.face(ds_record.subdir, ds_record.fn, ds_record.ext) s3_url = self.data_store_s3.face(ds_record.uuid) image_record = ImageRecord(ds_record, fp_im, s3_url) image_records.append(image_record) return image_records # ---------------------------------------------------------------------- # utilities def df_vecs_to_dict(self, df_vec): # convert the DataFrame CSV to float list of vecs # n_dims = len(df_vec.keys()) - 2 # number of columns with 'd1, d2,...d256' #return [[df[f'd{i}'] for i in range(1,n_dims+1)] for df_idx, df in df_vec.iterrows()] # return [[df[f'd{i}'] for i in range(1,n_dims+1)] for df_idx, df in df_vec.iterrows()] return [list(map(float, x.vec.split(','))) for x in df_vec.itertuples()] def df_vec_roi_idxs_to_dict(self, df): # convert the DataFrame CSV to float list of vecs #return [x.roi_index for x in df.itertuples()] return [int(x.roi_index) for i,x in df.iterrows()] def similar(self, query_vec, n_results): '''Finds most similar N indices of query face vector :query_vec: (list) of 128 floating point numbers of face encoding :n_results: (int) number of most similar indices to return :returns (list) of (int) indices ''' # uses np.linalg based on the ageitgey/face_recognition code return top_idxs class ImageRecord: def __init__(self, ds_record, fp, url, ds_rois=None, ds_identities=None): # maybe more other meta will go there self.image_index = ds_record.index self.sha256 = ds_record.sha256 self.uuid = ds_record.uuid self.filepath = fp self.url = url self._identities = [] # image records contain ROIs # ROIs are linked to identities #self._identities = [Identity(x) for x in ds_identities] @property def identity(self, index): return self._identity def summarize(self): '''Summarizes data for debugging''' log = Logger.getLogger() log.info(f'filepath: {self.filepath}') log.info(f'sha256: {self.sha256}') log.info(f'UUID: {self.uuid}') log.info(f'S3 url: {self.url}') for identity in self._identities: log.info(f'fullname: {identity.fullname}') log.info(f'description: {identity.description}') log.info(f'gender: {identity.gender}') log.info(f'images: {identity.n_images}') class Identity: def __init__(self, idx, name='NA', desc='NA', gender='NA', n_images=1, url='NA', age='NA', nationality='NA'): self.index = idx self.name = name self.description = desc self.gender = gender self.n_images = n_images self.url = url self.age = age self.nationality = nationality