1 files changed, 28 insertions, 25 deletions

diff --git a/megapixels/app/models/dataset.py b/megapixels/app/models/dataset.py
index 1b91467b..88986873 100644
--- a/megapixels/app/models/dataset.py
+++ b/megapixels/app/models/dataset.py
@@ -154,36 +154,42 @@ class Dataset:
     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]
+    df_vector = self._metadata[types.Metadata.FACE_VECTOR]
+    df_roi = self._metadata[types.Metadata.FACE_ROI]
     if types.Metadata.IDENTITY in self._metadata.keys():
       df_identity = self._metadata[types.Metadata.IDENTITY]
     else:
       df_identity = None
-    df_roi = self._metadata[types.Metadata.FACE_ROI]
 
     identities = []
 
     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_roi = df_roi.iloc[roi_index]
+      record_idx = int(ds_roi.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) 
       identities = []
+      
+      bbox_norm = BBox.from_xywh_norm(ds_roi.x, ds_roi.y, ds_roi.w, ds_roi.w)
+
       if types.Metadata.IDENTITY in self._metadata.keys():
-        ds_id = df_identity.loc[df_identity['identity_key'] == ik].iloc[0]
-        identity = Identity(idx, 
+        ds_id = df_identity.loc[df_identity['identity_key'] == ds_record.identity_key].iloc[0]
+        identity = Identity(record_idx, 
           name_display=ds_id.name_display,
-          name_full=ds_id.name_full,
           description=ds_id.description,
           gender=ds_id.gender,
-          image_index=ds_id.image_index,
-          identity_key=ds_id.identity_key)
-        identities.append(identity)
-      image_record = ImageRecord(ds_record, fp_im, s3_url, identities=identities)
+          roi_index=roi_index,
+          identity_key=ds_id.identity_key,
+          num_images=ds_id.num_images)
+      else:
+        identity = None 
+      image_record = ImageRecord(ds_record, fp_im, s3_url, bbox_norm, identity=identity)
       image_records.append(image_record)
     return image_records
 
@@ -216,7 +222,7 @@ class Dataset:
 
 class ImageRecord:
 
-  def __init__(self, ds_record, fp, url, rois=None, identities=None):
+  def __init__(self, ds_record, fp, url, bbox_norm, identity=None):
     # maybe more other meta will go there
     self.image_index = ds_record.index
     self.sha256 = ds_record.sha256
@@ -225,15 +231,11 @@ class ImageRecord:
     self.width = ds_record.width
     self.height = ds_record.height
     self.url = url
-    self.rois = rois
-    self.identities = identities
+    self.bbox = bbox_norm
+    self.identity = identity
     # image records contain ROIs
     # ROIs are linked to identities
 
-  @property
-  def identity(self, index):
-    return self._identity
-
   def summarize(self):
     '''Summarizes data for debugging'''
     log = Logger.getLogger()
@@ -241,19 +243,20 @@ class ImageRecord:
     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}')
+    if self.identity:
+      log.info(f'name: {self.identity.name_display}')
+      log.info(f'description: {self.identity.description}')
+      log.info(f'gender: {self.identity.gender}')
+      log.info(f'images: {self.identity.num_images}')
 
 
 class Identity:
 
-  def __init__(self, idx, name_display=None, name_full=None, description=None, gender=None, roi=None):
+  def __init__(self, idx, identity_key=None, name_display=None, num_images=None,
+      description=None, gender=None, roi_index=None):
     self.index = idx
     self.name_display = name_display
-    self.name_full = name_full
     self.description = description
     self.gender = gender
-    self.roi = roi
+    self.roi_index = roi_index
+    self.num_images = num_images