# import the necessary packages
import numpy as np
import cv2 as cv

# Load Model
print("loading emotion detection model ...")
#credits to https://github.com/HSE-asavchenko/face-emotion-recognition
model_path = 'robot/emotion_detection/models/emotions/affectnet_emotions/'
net = cv.dnn.readNet(model_path+'mobilenet_7.pbtxt',model_path+'mobilenet_7.pb')
labels = None
labelsFile = model_path+"labels.txt"
with open(labelsFile, 'rt') as f:
    labels = f.read().rstrip('\n').split('\n')
print("... face emotion model loaded")

def setTarget(gpu):
    if gpu:
        net.setPreferableBackend(cv.dnn.DNN_BACKEND_CUDA)
        net.setPreferableTarget(cv.dnn.DNN_TARGET_CUDA)
    else:
        net.setPreferableBackend(cv.dnn.DNN_BACKEND_OPENCV)
        net.setPreferableTarget(cv.dnn.DNN_TARGET_CPU)

setTarget(False)

def detectEmotion(face):
    # transform the image to suitable input of the Inception DNN
    blob = cv.dnn.blobFromImage(face, 1.0, (224, 224), (123.68, 116.779, 103.939), True, False)
    # put the input to the network
    net.setInput(blob)
    # launch the network and get the produced output
    scores = net.forward()[0]
    # process the output typical for softmax classifier
    emotion = np.argmax(scores) # 0..6 see labels.txt
    return emotion
      
def displayEmotion(image, emotion):
    # display the detected emotion
    label = labels[emotion]
    print(label)
    cv.putText(image, label, (10,50), cv.FONT_HERSHEY_SIMPLEX, 1.4, (0, 0, 255), 3, cv.LINE_AA)

def main(filename):
    # load image
    image = cv.imread(filename)

    # detect emotion
    emotion = detectEmotion(image)

    # display emotion
    displayEmotion(image, emotion)
    
    # show the output image
    cv.imshow("Emotion", image)
    cv.waitKey(0)

    cv.destroyAllWindows()

if __name__ == '__main__':
    main('happy.png')
    main('fear.png')
    main('angry.png')
