change to cli_proc

1 files changed, 17 insertions, 10 deletions

diff --git a/megapixels/app/models/dataset.py b/megapixels/app/models/dataset.py
index 88986873..a7227a70 100644
--- a/megapixels/app/models/dataset.py
+++ b/megapixels/app/models/dataset.py
@@ -9,6 +9,7 @@ import logging
 
 import pandas as pd
 import numpy as np
+import cv2 as cv
 
 from app.settings import app_cfg as cfg
 from app.settings import types
@@ -68,7 +69,7 @@ class Dataset:
       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()
+      #sys.exit()
 
   def load_metadata(self, metadata_type):
     if metadata_type == types.Metadata.FILE_RECORD:
@@ -107,7 +108,7 @@ class Dataset:
     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)
+    bbox = BBox.from_xywh_dim(ds_roi.x, ds_roi.y, ds_roi.w, ds_roi.h, 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]
@@ -169,17 +170,21 @@ class Dataset:
       roi_index = self._face_vector_roi_idxs[match_idx]
       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}')
+      ds_record = df_record.iloc[record_idx]
       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)
+      dim = (ds_record.width, ds_record.height)
+      im = cv.imread(fp_im)
+      dim = im.shape[:2][::-1]
+      self.log.debug(f'dim: {dim}')
+      s3_url = self.data_store_s3.face(ds_record.uuid)
+      bbox_norm = BBox.from_xywh_norm_dim(ds_roi.x, ds_roi.y, ds_roi.w, ds_roi.h, dim)
+      self.log.debug(f'bbox_norm: {bbox_norm}')
+      score = sim_scores[match_idx]
 
       if types.Metadata.IDENTITY in self._metadata.keys():
         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,
           description=ds_id.description,
@@ -189,7 +194,7 @@ class Dataset:
           num_images=ds_id.num_images)
       else:
         identity = None 
-      image_record = ImageRecord(ds_record, fp_im, s3_url, bbox_norm, identity=identity)
+      image_record = ImageRecord(ds_record, fp_im, s3_url, bbox_norm, score, identity=identity)
       image_records.append(image_record)
     return image_records
 
@@ -222,7 +227,7 @@ class Dataset:
 
 class ImageRecord:
 
-  def __init__(self, ds_record, fp, url, bbox_norm, identity=None):
+  def __init__(self, ds_record, fp, url, bbox_norm, score, identity=None):
     # maybe more other meta will go there
     self.image_index = ds_record.index
     self.sha256 = ds_record.sha256
@@ -232,7 +237,9 @@ class ImageRecord:
     self.height = ds_record.height
     self.url = url
     self.bbox = bbox_norm
+    self.score = score
     self.identity = identity
+
     # image records contain ROIs
     # ROIs are linked to identities