1 files changed, 59 insertions, 6 deletions
diff --git a/megapixels/app/models/sql_factory.py b/megapixels/app/models/sql_factory.py
index e35c3e15..9a44941b 100644
--- a/megapixels/app/models/sql_factory.py
+++ b/megapixels/app/models/sql_factory.py
@@ -19,6 +19,7 @@ connection_url = "mysql+mysqldb://{}:{}@{}/{}".format(
 
 datasets = {}
 loaded = False
+Session = None
 
 def list_datasets():
   return [dataset.describe() for dataset in datasets.values()]
@@ -31,10 +32,11 @@ def get_table(name, table_name):
   return dataset.get_table(table_name) if dataset else None
 
 def load_sql_datasets(replace=False, base_model=None):
-  global datasets, loaded
+  global datasets, loaded, Session
   if loaded:
     return datasets
-  engine = create_engine(connection_url) if replace else None
+  engine = create_engine(connection_url)
+  Session = sessionmaker(bind=engine)
   for path in glob.iglob(os.path.join(cfg.DIR_FAISS_METADATA, "*")):
     dataset = load_sql_dataset(path, replace, engine, base_model)
     datasets[dataset.name] = dataset
@@ -79,6 +81,28 @@ class SqlDataset:
       'tables': list(self.tables.keys()),
     }
 
+  def get_identity(self, id):
+    table = self.get_table('identity_meta')
+    # id += 1
+    identity = table.query.filter(table.image_id <= id).order_by(table.image_id.desc()).first().toJSON()
+    return {
+      'uuid': self.select('uuids', id),
+      'identity': identity,
+      'roi': self.select('roi', id),
+      'pose': self.select('pose', id),
+    }
+
+  def select(self, table, id):
+    table = self.get_table(table)
+    if not table:
+      return None
+    session = Session()
+    # for obj in session.query(table).filter_by(id=id):
+    print(table)
+    obj = session.query(table).filter(table.id == id).first()
+    session.close()
+    return obj.toJSON()
+
   def get_table(self, type):
     if type in self.tables:
       return self.tables[type]
@@ -102,6 +126,11 @@ class SqlDataset:
       __tablename__ = self.name + "_uuid"
       id = Column(Integer, primary_key=True)
       uuid = Column(String(36), nullable=False)
+      def toJSON(self):
+        return {
+          'id': self.id,
+          'uuid': self.uuid,
+        }
     return UUID
 
   # ==> roi.csv <==
@@ -118,6 +147,17 @@ class SqlDataset:
       w = Column(Float, nullable=False)
       x = Column(Float, nullable=False)
       y = Column(Float, nullable=False)
+      def toJSON(self):
+        return {
+          'id': self.id,
+          'image_index': self.image_index,
+          'image_height': self.image_height,
+          'image_width': self.image_width,
+          'w': self.w,
+          'h': self.h,
+          'x': self.x,
+          'y': self.y,
+        }
     return ROI
 
   # ==> identity.csv <==
@@ -132,6 +172,15 @@ class SqlDataset:
       gender = Column(String(1), nullable=False)
       images = Column(Integer, nullable=False)
       image_id = Column(Integer, nullable=False)
+      def toJSON(self):
+        return {
+          'id': self.id,
+          'image_id': self.image_id,
+          'fullname': self.fullname,
+          'images': self.images,
+          'gender': self.gender,
+          'description': self.description,
+        }
     return Identity
 
   # ==> pose.csv <==
@@ -145,8 +194,12 @@ class SqlDataset:
       pitch = Column(Float, nullable=False)
       roll = Column(Float, nullable=False)
       yaw = Column(Float, nullable=False)
+      def toJSON(self):
+        return {
+          'id': self.id,
+          'image_id': self.image_id,
+          'pitch': self.pitch,
+          'roll': self.roll,
+          'yaw': self.yaw,
+        }
     return Pose
-
-
-# Session = sessionmaker(bind=engine)
-# session = Session()