diff options
Diffstat (limited to 'megapixels/app/models/dataset.py')
| -rw-r--r-- | megapixels/app/models/dataset.py | 125 |
1 files changed, 80 insertions, 45 deletions
diff --git a/megapixels/app/models/dataset.py b/megapixels/app/models/dataset.py index 8fef8a7e..35e10465 100644 --- a/megapixels/app/models/dataset.py +++ b/megapixels/app/models/dataset.py @@ -23,7 +23,7 @@ from app.utils.logger_utils import Logger class Dataset: - def __init__(self, opt_data_store, opt_dataset_type, load_files=True): + def __init__(self, opt_data_store, opt_dataset_type): self._dataset_type = opt_dataset_type # enum type self.log = Logger.getLogger() self._metadata = {} @@ -31,31 +31,62 @@ class Dataset: self._nullframe = pd.DataFrame() # empty placeholder self.data_store = DataStore(opt_data_store, self._dataset_type) self.data_store_s3 = DataStoreS3(self._dataset_type) - self.load_metadata() - def load_metadata(self): - '''Loads all CSV files into (dict) of DataFrames''' - self.log.info(f'creating dataset: {self._dataset_type}...') - for metadata_type in types.Metadata: - self.log.info(f'load metadata: {metadata_type}') - 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') - if metadata_type == types.Metadata.FACE_VECTOR: - # convert DataFrame to list of floats - self._face_vectors = self.df_to_vec_list(self._metadata[metadata_type]) - self.log.info(f'build face vector dict: {len(self._face_vectors)}') - self._metadata[metadata_type].drop('vec', axis=1, inplace=True) - else: - self.log.error(f'File not found: {fp_csv}. Exiting.') - sys.exit() - self.log.info('finished loading') + def load_face_vectors(self): + metadata_type = types.Metadata.FACE_VECTOR + 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') + # convert DataFrame to list of floats + self._face_vectors = self.df_vecs_to_dict(self._metadata[metadata_type]) + self._face_vector_idxs = self.df_vec_idxs_to_dict(self._metadata[metadata_type]) + 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) + else: + self.log.error(f'File not found: {fp_csv}. Exiting.') + sys.exit() + + def load_records(self): + metadata_type = types.Metadata.FILE_RECORD + 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') + else: + self.log.error(f'File not found: {fp_csv}. Exiting.') + sys.exit() + + def load_identities(self): + metadata_type = types.Metadata.IDENTITY + 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') + else: + self.log.error(f'File not found: {fp_csv}. Exiting.') + sys.exit() def metadata(self, opt_metadata_type): - return self._metadata.get(opt_metadata_type, self._nullframe) + return self._metadata.get(opt_metadata_type, None) - def roi_idx_to_record(self, vector_index): + def index_to_record(self, index): + # get record meta + df_record = self._metadata[types.Metadata.FILE_RECORD] + ds_record = df_record.iloc[index] + identity_index = ds_record.identity_index + # get identity meta + df_identity = self._metadata[types.Metadata.IDENTITY] + # future datasets can have multiple identities per images + ds_identities = df_identity.iloc[identity_index] + # get filepath and S3 url + fp_im = self.data_store.face_image(ds_record.subdir, ds_record.fn, ds_record.ext) + s3_url = self.data_store_s3.face_image(ds_record.uuid) + image_record = ImageRecord(ds_record, fp_im, s3_url, ds_identities=ds_identities) + return image_record + + def vector_to_record(self, record_index): '''Accumulates image and its metadata''' df_face_vector = self._metadata[types.Metadata.FACE_VECTOR] ds_face_vector = df_face_vector.iloc[vector_index] @@ -115,18 +146,24 @@ class Dataset: for match_idx in match_idxs: # get the corresponding face vector row - self.log.debug(f'find match index: {match_idx}') - image_record = self.roi_idx_to_record(match_idx) + roi_index = self._face_vector_roi_idxs[match_idx] + self.log.debug(f'find match index: {match_idx}, --> roi_index: {roi_index}') + image_record = self.roi_idx_to_record(roi_index) image_records.append(image_record) return image_records # ---------------------------------------------------------------------- # utilities - def df_to_vec_list(self, df): + def df_vecs_to_dict(self, df): # convert the DataFrame CSV to float list of vecs return [list(map(float,x.vec.split(','))) for x in df.itertuples()] + def df_vec_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.image_index for x in df.itertuples()] + def similar(self, query_vec, n_results): '''Finds most similar N indices of query face vector :query_vec: (list) of 128 floating point numbers of face encoding @@ -141,37 +178,35 @@ class Dataset: class ImageRecord: - def __init__(self, image_index, sha256, uuid, bbox, filepath, url): - self.image_index = image_index - self.sha256 = sha256 - self.uuid = uuid - self.bbox = bbox - self.filepath = filepath + def __init__(self, ds_record, fp, url, ds_rois=None, ds_identities=None): + # maybe more other meta will go there + self.image_index = ds_record.index + self.sha256 = ds_record.sha256 + self.uuid = ds_record.uuid + self.filepath = fp self.url = url - self._identity = None + self._identities = [] + # image records contain ROIs + # ROIs are linked to identities + + #self._identities = [Identity(x) for x in ds_identities] @property - def identity(self): + def identity(self, index): return self._identity - @identity.setter - def identity(self, value): - self._identity = value - def summarize(self): '''Summarizes data for debugging''' log = Logger.getLogger() log.info(f'filepath: {self.filepath}') log.info(f'sha256: {self.sha256}') log.info(f'UUID: {self.uuid}') - log.info(f'BBox: {self.bbox}') - log.info(f's3 url: {self.url}') - if self._identity: - log.info(f'name: {self._identity.name}') - log.info(f'age: {self._identity.age}') - log.info(f'gender: {self._identity.gender}') - log.info(f'nationality: {self._identity.nationality}') - log.info(f'images: {self._identity.n_images}') + 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}') class Identity: |