change bbox to norm, refine face extractor

1 files changed, 16 insertions, 9 deletions

diff --git a/megapixels/app/models/dataset.py b/megapixels/app/models/dataset.py
index eb0109a7..bbef9ff5 100644
--- a/megapixels/app/models/dataset.py
+++ b/megapixels/app/models/dataset.py
@@ -44,6 +44,9 @@ class Dataset:
       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()
@@ -53,7 +56,7 @@ class Dataset:
     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')
+      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()
@@ -142,33 +145,37 @@ class Dataset:
     # 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, n_results)[:n_results]
+    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]
-      df_record = self._metadata[types.Metadata.FILE_RECORD]
-      ds_record = df_record.iloc[roi_index]
+      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)
-      #roi_index = self._face_vector_roi_idxs[match_idx]
-      #image_record = self.roi_idx_to_record(roi_index)
       image_records.append(image_record)
     return image_records
 
   # ----------------------------------------------------------------------
   # utilities
 
-  def df_vecs_to_dict(self, df):
+  def df_vecs_to_dict(self, df_vec):
     # convert the DataFrame CSV to float list of vecs
-    return [list(map(float,x.vec.split(','))) for x in df.itertuples()]
+    # 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 [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