# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/03_tables.ipynb.

# %% auto 0
__all__ = ['Image', 'SCFaceMixin', 'EnfsiMixin', 'VideoMixin', 'EnfsiImage', 'SCImage', 'VideoFrame', 'CPFrame',
           'EnfsiVideoFrame', 'Detector', 'CroppedImage', 'EmbeddingModel', 'FaceImage', 'QualityModel', 'QualityImage']

# %% ../nbs/03_tables.ipynb 3
from typing import Tuple

from sqlalchemy import Column, Integer, String, Enum, Boolean, PickleType, ForeignKey, Table, Float
from sqlalchemy.orm import relationship, declarative_mixin, declared_attr
from sqlalchemy.ext.declarative import declarative_base

from deepface import DeepFace
import cv2
import enum
import numpy as np


# %% ../nbs/03_tables.ipynb 4
#To start table creation.
Base = declarative_base()

# %% ../nbs/03_tables.ipynb 6
## All the attribute Enum for the Image.


class Gender(enum.Enum):
    MALE = "Man"
    FEMALE = "Woman"

class Age(enum.Enum):
    # todo: fill enum with age number.
    CHILD = "0-12"
    ADOLESCENT = '13-17'
    YOUNG_ADULT = '18-30'
    ADULT = '31-45'
    MIDDLE_AGED_ADULT = '46-64'
    SENIOR = '65-100'

    @staticmethod
    def age2enum(age:int)->Enum:
        if age > 65:
            age_enum = Age.SENIOR
        elif age > 45:
            age_enum = Age.MIDDLE_AGED_ADULT
        elif age > 30:
            age_enum = Age.ADULT
        elif age > 18:
            age_enum = Age.YOUNG_ADULT
        elif age > 12:
            age_enum = Age.ADOLESCENT
        elif age > 0:
            age_enum = Age.CHILD
        else:
            age_enum = None
            print(f'Age {age} not in range, None returned')

        return age_enum


class Yaw(enum.Enum):
    FRONTAL = "straight"
    HALF_TURNED = "slightly_turned"
    PROFILE = "sideways"


class Pitch(enum.Enum):
    UP = "upwards"
    HALF_UP = "slightly_upwards"
    FRONTAL = "straight"
    HALF_DOWN = "slightly_downwards"
    DOWN = "downwards"
    
class Roll(enum.Enum):
    FRONTAL = "straight"
    HALF_LEANING = "slightly_inclined"
    HORIZONTAL = "completely_inclined"

class Emotion(enum.Enum):
    # TODO: implementChange to expression
    '''angry, fear, neutral, sad, disgust, happy and surprise'''
    ANGRY = 'angry'
    FEAR = 'fear'
    NEUTRAL = 'neutral'
    SAD = 'sad'
    DISGUST = 'disgust'
    HAPPY = 'happy'
    SURPRISE = 'surprise'
    

class Race(enum.Enum):
    '''asian, white, middle eastern, indian, latino and black'''
    ASIAN='asian'
    WHITE = 'white'
    MIDDLE_EASTERN = 'middle eastern'
    INDIAN = 'indian'
    LATINO = 'latino hispanic'
    BLACK = 'black'


class Distance(enum.Enum):
    '''asian, white, middle eastern, indian, latino and black'''
    FRONTAL = 50
    SHORT= 100
    MEDIUM= 260
    FAR = 420


# %% ../nbs/03_tables.ipynb 8
class Image(Base):
    "Image SQL class"
    __tablename__ = "image"
    image_id = Column(Integer, primary_key=True) # Image primary key
    path = Column(String) # Absolute or relative path
    identity = Column(String) # Person identity of the image
    source = Column(String) # Database the image belongs to
    gender = Column(Enum(Gender))
    age = Column(Enum(Age))
    age_number = Column(Float)
    emotion = Column(Enum(Emotion))
    race = Column(Enum(Race))
    yaw = Column(Enum(Yaw))
    pitch = Column(Enum(Pitch))
    roll = Column(Enum(Roll))
    headgear = Column(Boolean)
    glasses = Column(Boolean)
    beard = Column(Boolean)
    other_occlusions = Column(Boolean)
    low_quality = Column(Boolean)
    

    type = Column(String)
    __mapper_args__ = {
        'polymorphic_identity': 'image',
        'polymorphic_on': type
    }

    croppedImages = relationship("CroppedImage", back_populates="images", lazy='subquery')

    def get_image(self):
        return cv2.imread(self.path)

# %% ../nbs/03_tables.ipynb 9
# class SCImage(Image):
#     "SCFace database"
    
#     __tablename__ = 'scImage'
#     scImage_id = Column(Integer, ForeignKey('image.image_id'), primary_key=True)

#     sc_type = Column(String)
#     distance = Column(Enum(Distance))
#     infrared = Column(Boolean)

#     __mapper_args__ = {
#         'polymorphic_identity': 'scImage',
#     }

# %% ../nbs/03_tables.ipynb 11
@declarative_mixin
class SCFaceMixin:
    "SC Face database mixin"
    sc_type = Column(String)
    distance = Column(Enum(Distance))
    infrared = Column(Boolean)

    @declared_attr
    def image_id(cls):
        return Column(Integer, ForeignKey('image.image_id'), primary_key=True)

# %% ../nbs/03_tables.ipynb 12
@declarative_mixin
class EnfsiMixin:
    "ENFSI database mixin"
    year = Column(Integer)

    @declared_attr
    def image_id(cls):
        return Column(Integer, ForeignKey('image.image_id'), primary_key=True)

# %% ../nbs/03_tables.ipynb 13
@declarative_mixin
class VideoMixin:
    source_video = Column(String)
    n_frame = Column(Integer)

    @declared_attr
    def image_id(cls):
        return Column(Integer, ForeignKey('image.image_id'), primary_key=True)

    def get_image(self):
        video = cv2.VideoCapture(self.path)
        video.set(1, self.n_frame)
        ret, image = video.read()
        if ret:
            return image

# %% ../nbs/03_tables.ipynb 15
class EnfsiImage(EnfsiMixin, Image):
    __tablename__ = 'enfsiImage'
    __mapper_args__ = {
        'polymorphic_identity': 'enfsiImage',
    }

# %% ../nbs/03_tables.ipynb 17
class SCImage(SCFaceMixin, Image):    
    __tablename__ = 'scImage'
    __mapper_args__ = {
        'polymorphic_identity': 'scImage',
    }

# %% ../nbs/03_tables.ipynb 19
class VideoFrame(VideoMixin, Image):
    __tablename__ = 'videoFrame'
    __mapper_args__ = {
        'polymorphic_identity': 'videoFrame',
    }

# %% ../nbs/03_tables.ipynb 20
class CPFrame(VideoFrame):
    "ChokePoint video frame"

    def get_image(self):
        "Especial method for getting the images in ChokePoint."
        pass

# %% ../nbs/03_tables.ipynb 21
class EnfsiVideoFrame(EnfsiMixin, VideoMixin, Image):
    __tablename__ = 'enfsiVideoFrame'
    __mapper_args__ = {
        'polymorphic_identity': 'enfsiVideoFrame',
    }

# %% ../nbs/03_tables.ipynb 22
class Detector(Base):
    "Detector SQL class"
    __tablename__ = "detector"
    detector_id = Column(Integer, primary_key=True)
    name = Column(String)

# %% ../nbs/03_tables.ipynb 23
class CroppedImage(Base):
    __tablename__ = 'croppedImage'
    croppedImage_id = Column(Integer, primary_key=True)

    image_id = Column(Integer, ForeignKey('image.image_id'))
    detector_id = Column(Integer, ForeignKey('detector.detector_id'))

    bounding_box = Column(PickleType)
    landmarks = Column(PickleType)
    face_detected = Column(Boolean)

    images = relationship("Image", foreign_keys=[image_id])
    detectors = relationship("Detector", foreign_keys=[detector_id])
    faceImages = relationship("FaceImage", back_populates="croppedImages")

    def get_cropped_image(self):
        image = self.images.get_image()
        if self.face_detected:
            return image[self.bounding_box[1]:self.bounding_box[1] + self.bounding_box[3],
                   self.bounding_box[0]:self.bounding_box[0] + self.bounding_box[2], :]
        else:
            return image

    def get_aligned_image(self, target_size:Tuple[int,int]=(112,112), ser_fiq = None):
        
        if self.detectors.name == 'mtcnn_serfiq':
            image = self.images.get_image() 
            aligned_image = ser_fiq.apply_mtcnn(image)                     
            return np.transpose(aligned_image, (1,2,0)) 
        
        else:
            
            aligned_img = DeepFace.detectFace(img_path = self.images.path, 
            target_size = target_size, 
            detector_backend = self.detectors.name, 
            align=True,
            enforce_detection=True)
            return aligned_img*255

# %% ../nbs/03_tables.ipynb 24
class EmbeddingModel(Base):
    __tablename__ = "embeddingModel"
    embeddingModel_id = Column(Integer, primary_key=True)
    name = Column(String)

# %% ../nbs/03_tables.ipynb 25
class FaceImage(Base):
    __tablename__ = 'faceImage'
    faceImage_id = Column(Integer, primary_key=True)

    croppedImage_id = Column(Integer, ForeignKey('croppedImage.croppedImage_id'))
    embeddingModel_id = Column(Integer, ForeignKey('embeddingModel.embeddingModel_id'))

    embeddings = Column(PickleType)
    confusion_score = Column(Float)

    croppedImages = relationship("CroppedImage", foreign_keys=[croppedImage_id])
    embeddingModels = relationship("EmbeddingModel", foreign_keys=[embeddingModel_id])

# %% ../nbs/03_tables.ipynb 26
class QualityModel(Base):
    __tablename__ = "qualityModel"
    qualityModel_id = Column(Integer, primary_key=True)
    name = Column(String)

# %% ../nbs/03_tables.ipynb 27
class QualityImage(Base):
    __tablename__ = 'qualityImage'
    qualityImage_id = Column(Integer, primary_key=True)

    faceImage_id = Column(Integer, ForeignKey('faceImage.faceImage_id'))
    qualityModel_id = Column(Integer, ForeignKey('qualityModel.qualityModel_id'))

    quality = Column(Float)
    quality_vec = Column(PickleType)

    faceImages = relationship("FaceImage", foreign_keys=[faceImage_id])
    qualityModels = relationship("QualityModel", foreign_keys=[qualityModel_id])
