mysql import functions

1 files changed, 19 insertions, 11 deletions

diff --git a/megapixels/app/models/sql_factory.py b/megapixels/app/models/sql_factory.py
index 4adc6f48..ecca0c7f 100644
--- a/megapixels/app/models/sql_factory.py
+++ b/megapixels/app/models/sql_factory.py
@@ -2,8 +2,7 @@ import os
 
 from sqlalchemy import create_engine, Table, Column, String, Integer, DateTime, Float
 from sqlalchemy.orm import sessionmaker
-from sqlalchemy.ext.declarative import declarative_base, declared_attr
-from sqlalchemy.ext.declarative import AbstractConcreteBase, ConcreteBase
+from sqlalchemy.ext.declarative import declarative_base
 
 connection_url = "mysql+mysqldb://{}:{}@{}/{}".format(
   os.getenv("DB_USER"),
@@ -12,15 +11,24 @@ connection_url = "mysql+mysqldb://{}:{}@{}/{}".format(
   os.getenv("DB_NAME")
 )
 
-engine = create_engine(connection_url)
-Session = sessionmaker(bind=engine)
-session = Session()
-Base = declarative_base(engine)
+# Session = sessionmaker(bind=engine)
+# session = Session()
+
 
 class SqlDataset:
-  def __init__(self, name):
+  """
+  Bridge between the facial information CSVs connected to the datasets, and MySQL
+  - each dataset should have files that can be loaded into these database models
+  - names will be fixed to work in SQL (index -> id)
+  - we can then have more generic models for fetching this info after doing a FAISS query
+  """
+  def __init__(self, name, base_model=None):
     self.name = name
     self.tables = {}
+    if base_model is None:
+      engine = create_engine(connection_url)
+      base_model = declarative_base(engine)
+    self.base_model = base_model
 
   def get_table(self, type):
     if type in self.tables:
@@ -41,7 +49,7 @@ class SqlDataset:
   # index,uuid
   # 0,f03fd921-2d56-4e83-8115-f658d6a72287
   def uuid_table(self):
-    class UUID(Base):
+    class UUID(self.base_model):
       __tablename__ = self.name + "_uuid"
       id = Column(Integer, primary_key=True)
       uuid = Column(String(36), nullable=False)
@@ -51,7 +59,7 @@ class SqlDataset:
   # index,h,image_height,image_index,image_width,w,x,y
   # 0,0.33000000000000007,250,0,250,0.32999999999999996,0.33666666666666667,0.35
   def roi_table(self):
-    class ROI(Base):
+    class ROI(self.base_model):
       __tablename__ = self.name + "_roi"
       id = Column(Integer, primary_key=True)
       h = Column(Float, nullable=False)
@@ -67,7 +75,7 @@ class SqlDataset:
   # index,fullname,description,gender,images,image_index
   # 0,A. J. Cook,Canadian actress,f,1,0
   def identity_table(self):
-    class Identity(Base):
+    class Identity(self.base_model):
       __tablename__ = self.name + "_identity"
       id = Column(Integer, primary_key=True)
       fullname = Column(String(36), nullable=False)
@@ -81,7 +89,7 @@ class SqlDataset:
   # index,image_index,pitch,roll,yaw
   # 0,0,11.16264458441435,10.415885631337728,22.99719032415318
   def pose_table(self):
-    class Pose(Base):
+    class Pose(self.base_model):
       __tablename__ = self.name + "_pose"
       id = Column(Integer, primary_key=True)
       image_id = Column(Integer, primary_key=True)