""" Dataset model: container for all CSVs about a dataset """ import os 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, path_utils from app.utils.logger_utils import Logger # ------------------------------------------------------------------------- # Dataset # ------------------------------------------------------------------------- class Dataset: def __init__(self, opt_dataset_type, opt_data_store=types.DataStore.NAS): 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 = path_utils.DataStore(opt_data_store, self._dataset_type) self.data_store_s3 = path_utils.DataStoreS3(self._dataset_type) def load(self, opt_data_store): '''Loads all CSV files into (dict) of DataFrames''' for metadata_type in types.Metadata: self.log.info(f'load metadata: {metadata_type}') 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') if metadata_type == types.Metadata.FACE_VECTOR: # convert DataFrame to list of floats self._face_vecs = self.df_to_vec_list(self._metadata[metadata_type]) self._metadata[metadata_type].drop('vec', axis=1, inplace=True) else: self.log.error('File not found: {fp_csv}. Replaced with empty DataFrame') self._metadata[metadata_type] = self._nullframe self.log.info('finished loading') def metadata(self, opt_metadata_type): return self._metadata.get(opt_metadata_type, self._nullframe) def roi_idx_to_record(self, vector_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: # 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 return image_record def matches(self, query_vec, n_results=5, threshold=0.5): 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) for match_idx in match_idxs: # get the corresponding face vector row image_record = self.roi_idx_to_record(match_idx) results.append(image_record) return image_records # ---------------------------------------------------------------------- # utilities def df_to_vec_list(self, df): # convert the DataFrame CSV to float list of vecs vecs = [list(map(float,x.vec.split(','))) for x in df.itertuples()] return vecs 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 vecs_sim_scores = np.linalg.norm(np.array([query_vec]) - np.array(self._face_vectors), axis=1) top_idxs = np.argpartition(vecs_sim_scores, n_results)[:n_results] return top_idxs class ImageRecord: def __init__(self, image_index, sha256, uuid, bbox, filepath, url): self.image_index = image_index self.sha256 = sha256 self.uuid = uuid self.bbox = bbox self.filepath = filepath self.url = url self._identity = None @property def identity(self): return self._identity @identity.setter def identity(self, value): self._identity = value 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'BBox: {self.bbox}') log.info(f's3 url: {self.url}') if self._identity: log.info(f'name: {self._identity.name}') log.info(f'age: {self._identity.age}') log.info(f'gender: {self._identity.gender}') log.info(f'nationality: {self._identity.nationality}') log.info(f'images: {self._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