diff options
Diffstat (limited to 'megapixels')
| -rw-r--r-- | megapixels/app/models/bbox.py | 7 | ||||
| -rw-r--r-- | megapixels/app/models/dataset.py | 53 |
2 files changed, 35 insertions, 25 deletions
diff --git a/megapixels/app/models/bbox.py b/megapixels/app/models/bbox.py index f65f7373..608aaaf8 100644 --- a/megapixels/app/models/bbox.py +++ b/megapixels/app/models/bbox.py @@ -252,6 +252,13 @@ class BBox: # Create from @classmethod + def from_xywh_norm(cls, x, y, w, h): + """Converts w, y, w, h to normalized BBox + :returns BBox + """ + return cls(x, y, x + w, y + h) + + @classmethod def from_xyxy_dim(cls, x1, y1, x2, y2, dim): """Converts x1, y1, w, h to BBox and normalizes :returns BBox 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 |