Merge branch 'master' of github.com:adamhrv/megapixels_dev

1 files changed, 96 insertions, 52 deletions

diff --git a/megapixels/app/models/dataset.py b/megapixels/app/models/dataset.py
index 11d568a5..35e10465 100644
--- a/megapixels/app/models/dataset.py
+++ b/megapixels/app/models/dataset.py
@@ -2,6 +2,7 @@
 Dataset model: container for all CSVs about a dataset
 """
 import os
+import sys
 from os.path import join
 from pathlib import Path
 import logging
@@ -12,7 +13,8 @@ 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 import file_utils, im_utils
+from app.models.data_store import DataStore, DataStoreS3
 from app.utils.logger_utils import Logger
 
 # -------------------------------------------------------------------------
@@ -21,36 +23,70 @@ from app.utils.logger_utils import Logger
 
 class Dataset:
   
-  def __init__(self, opt_dataset_type, opt_data_store=types.DataStore.NAS):
+  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 = path_utils.DataStore(opt_data_store, self._dataset_type)
-    self.data_store_s3 = path_utils.DataStoreS3(self._dataset_type)
+    self.data_store = DataStore(opt_data_store, self._dataset_type)
+    self.data_store_s3 = 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 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_idxs = self.df_vec_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)
+    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).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, self._nullframe)
+    return self._metadata.get(opt_metadata_type, None)
 
-  def roi_idx_to_record(self, vector_index):
+  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_image(ds_record.subdir, ds_record.fn, ds_record.ext)
+    s3_url = self.data_store_s3.face_image(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]
@@ -80,7 +116,7 @@ class Dataset:
     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:
+    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]
@@ -95,27 +131,38 @@ class Dataset:
       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 matches(self, query_vec, n_results=5, threshold=0.5):
+  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)
+    #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, n_results)[:n_results]
+
     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)
+      roi_index = self._face_vector_roi_idxs[match_idx]
+      self.log.debug(f'find match index: {match_idx}, --> roi_index: {roi_index}')
+      image_record = self.roi_idx_to_record(roi_index)
+      image_records.append(image_record)
     return image_records
 
   # ----------------------------------------------------------------------
   # utilities
 
-  def df_to_vec_list(self, df):
+  def df_vecs_to_dict(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
+    return [list(map(float,x.vec.split(','))) for x in df.itertuples()]
+
+  def df_vec_idxs_to_dict(self, df):
+    # convert the DataFrame CSV to float list of vecs
+    #return [x.roi_index for x in df.itertuples()]
+    return [x.image_index for x in df.itertuples()]
 
   def similar(self, query_vec, n_results):
     '''Finds most similar N indices of query face vector
@@ -124,45 +171,42 @@ class Dataset:
     :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
+  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._identity = None
+    self._identities = []
+    # image records contain ROIs
+    # ROIs are linked to identities
+
+    #self._identities = [Identity(x) for x in ds_identities]
 
   @property
-  def identity(self):
+  def identity(self, index):
     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}')
+    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: