diff options
Diffstat (limited to 'megapixels/app/models/dataset.py')
| -rw-r--r-- | megapixels/app/models/dataset.py | 27 |
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 |