path: root/megapixels/app/models/dataset.py

"""
Dataset model: container for all CSVs about a dataset
"""
import os
import sys
from os.path import join
from pathlib import Path
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
from app.models.bbox import BBox
from app.utils import file_utils, im_utils
from app.models.data_store import DataStore, DataStoreS3
from app.utils.logger_utils import Logger

# -------------------------------------------------------------------------
# Dataset
# -------------------------------------------------------------------------

class Dataset:
  
  def __init__(self, opt_data_store, opt_dataset_type):
    self._dataset_type = opt_dataset_type  # enum type
    self.log = Logger.getLogger()
    self._metadata = {}
    self._face_vectors = []
    self._nullframe = pd.DataFrame()  # empty placeholder
    self.data_store = DataStore(opt_data_store, self._dataset_type)
    self.data_store_s3 = DataStoreS3(self._dataset_type)

  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_roi_idxs = self.df_vec_roi_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)
      #n_dims = len(self._metadata[metadata_type].keys()) - 2
      #drop_keys = [f'd{i}' for i in range(1,n_dims+1)]
      #self._metadata[metadata_type].drop(drop_keys, axis=1, inplace=True)
    else:
      self.log.error(f'File not found: {fp_csv}. Exiting.')
      sys.exit()

  def _load_file_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, dtype=cfg.FILE_RECORD_DTYPES).set_index('index')
    else:
      self.log.error(f'File not found: {fp_csv}. Exiting.')
      sys.exit()

  def _load_metadata(self, 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')
    else:
      self.log.error(f'File not found: {fp_csv}. Exiting.')
      #sys.exit()

  def load_metadata(self, metadata_type):
    if metadata_type == types.Metadata.FILE_RECORD:
      self._load_file_records()
    elif metadata_type == types.Metadata.FACE_VECTOR:
      self._load_face_vectors()
    else:
      self._load_metadata(metadata_type)
  
  def metadata(self, opt_metadata_type):
    return self._metadata.get(opt_metadata_type, None)

  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(ds_record.subdir, ds_record.fn, ds_record.ext)
    s3_url = self.data_store_s3.face(ds_record.uuid) 
    image_record = ImageRecord(ds_record, fp_im, s3_url)
    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]
    # get the match's ROI index
    image_index = ds_face_vector.image_index
    # get the roi dataframe
    df_face_roi = self._metadata[types.Metadata.FACE_ROI]
    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.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]
    sha256 = ds_sha256.sha256
    # get the local filepath
    df_filepath = self._metadata[types.Metadata.FILEPATH]
    ds_file = df_filepath.iloc[image_index]
    fp_im = self.data_store.face_image(ds_file.subdir, ds_file.fn, ds_file.ext)\
    # get remote path
    df_uuid = self._metadata[types.Metadata.UUID]
    ds_uuid = df_uuid.iloc[image_index]
    uuid = ds_uuid.uuid
    fp_url = self.data_store_s3.face_image(uuid)
    fp_url_crop = self.data_store_s3.face_image_crop(uuid)

    image_record = ImageRecord(image_index, sha256, uuid, bbox, fp_im, fp_url)
    # now get the identity index (if available)
    identity_index = ds_sha256.identity_index
    if identity_index > -1:
      # then use the identity index to get the identity meta
      df_identity = df_filepath = self._metadata[types.Metadata.IDENTITY]
      ds_identity = df_identity.iloc[identity_index]
      # get the name and description
      name = ds_identity.fullname
      desc = ds_identity.description
      gender = ds_identity.gender
      n_images = ds_identity.images
      url = '(url)'  # TODO
      age = '(age)' # TODO
      nationality = '(nationality)'
      identity = Identity(identity_index, name=name, desc=desc, gender=gender, n_images=n_images,
        url=url, age=age, nationality=nationality)
      image_record.identity = identity 
    else:
      self.log.info(f'no identity index: {ds_sha256}')

    return image_record


  def find_matches(self, query_vec, n_results=5, threshold=0.6):
    image_records = []  # list of image matches w/identity if available
    # find most similar feature vectors indexes
    #match_idxs = self.similar(query_vec, n_results, threshold)

    # TODO: add cosine similarity
    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_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

    identities = []

    for match_idx in match_idxs:
      # get the corresponding face vector row
      roi_index = self._face_vector_roi_idxs[match_idx]
      ds_roi = df_roi.iloc[roi_index]
      record_idx = int(ds_roi.record_index)
      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)
      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}')
      self.log.debug(f'match_idx: {match_idx}, record_idx: {record_idx}, roi_index: {roi_index}, len sim_scores: {len(sim_scores)}')
      try:
        score = sim_scores[match_idx]
      except Exception as e:
        self.log.error(e)
      try:
        score = sim_scores[record_idx]
      except Exception as e:
        self.log.error(e)



      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,
          gender=ds_id.gender,
          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, score, identity=identity)
      image_records.append(image_record)
    return image_records

  # ----------------------------------------------------------------------
  # utilities

  def df_vecs_to_dict(self, df_vec):
    # convert the DataFrame CSV to float list of vecs
    # n_dims = len(df_vec.keys()) - 2  # number of columns with 'd1, d2,...d256'
    #return [[df[f'd{i}'] for i in range(1,n_dims+1)] for df_idx, df in df_vec.iterrows()]
    # return [[df[f'd{i}'] for i in range(1,n_dims+1)] for df_idx, df in df_vec.iterrows()]
    return [list(map(float, x.vec.split(','))) for x in df_vec.itertuples()]

  def df_vec_roi_idxs_to_dict(self, df):
    # convert the DataFrame CSV to float list of vecs
    #return [x.roi_index for x in df.itertuples()]
    return [int(x.roi_index) for i,x in df.iterrows()]

  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
    :n_results: (int) number of most similar indices to return
    :returns (list) of (int) indices
    '''
    # uses np.linalg based on the ageitgey/face_recognition code
    
    return top_idxs



class ImageRecord:

  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
    self.uuid = ds_record.uuid
    self.filepath = fp
    self.width = ds_record.width
    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

  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'S3 url: {self.url}')
    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, 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.description = description
    self.gender = gender
    self.roi_index = roi_index
    self.num_images = num_images