diff options
Diffstat (limited to 'megapixels/app/models')
| -rw-r--r-- | megapixels/app/models/dataset.py | 179 |
1 files changed, 179 insertions, 0 deletions
diff --git a/megapixels/app/models/dataset.py b/megapixels/app/models/dataset.py new file mode 100644 index 00000000..11d568a5 --- /dev/null +++ b/megapixels/app/models/dataset.py @@ -0,0 +1,179 @@ +""" +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 |