import os
import time

import cv2
import sys
sys.path.append('..')
import numpy as np
from math import cos, sin

from lib.FSANET_model import *
import numpy as np
from keras.layers import Average


def draw_axis(img, yaw, pitch, roll, tdx=None, tdy=None, size = 50):
    print(yaw,roll,pitch)
    pitch = pitch * np.pi / 180
    yaw = -(yaw * np.pi / 180)
    roll = roll * np.pi / 180

    if tdx != None and tdy != None:
        tdx = tdx
        tdy = tdy
    else:
        height, width = img.shape[:2]
        tdx = width / 2
        tdy = height / 2

    # X-Axis pointing to right. drawn in red
    x1 = size * (cos(yaw) * cos(roll)) + tdx
    y1 = size * (cos(pitch) * sin(roll) + cos(roll) * sin(pitch) * sin(yaw)) + tdy

    # Y-Axis | drawn in green
    #        v
    x2 = size * (-cos(yaw) * sin(roll)) + tdx
    y2 = size * (cos(pitch) * cos(roll) - sin(pitch) * sin(yaw) * sin(roll)) + tdy

    # Z-Axis (out of the screen) drawn in blue
    x3 = size * (sin(yaw)) + tdx
    y3 = size * (-cos(yaw) * sin(pitch)) + tdy

    cv2.line(img, (int(tdx), int(tdy)), (int(x1),int(y1)),(0,0,255),3)
    cv2.line(img, (int(tdx), int(tdy)), (int(x2),int(y2)),(0,255,0),3)
    cv2.line(img, (int(tdx), int(tdy)), (int(x3),int(y3)),(255,0,0),2)

    return img
    
def draw_results_ssd(detected,input_img,faces,ad,img_size,img_w,img_h,model):
    
    # loop over the detections
    if detected.shape[2]>0:
        for i in range(0, detected.shape[2]):
            # extract the confidence (i.e., probability) associated with the
            # prediction
            confidence = detected[0, 0, i, 2]

            # filter out weak detections
            if confidence > 0.5:
                # compute the (x, y)-coordinates of the bounding box for
                # the face and extract the face ROI
                (h0, w0) = input_img.shape[:2]
                box = detected[0, 0, i, 3:7] * np.array([w0, h0, w0, h0])
                (startX, startY, endX, endY) = box.astype("int")
                # print((startX, startY, endX, endY))
                x1 = startX
                y1 = startY
                w = endX - startX
                h = endY - startY
                
                x2 = x1+w
                y2 = y1+h

                xw1 = max(int(x1 - ad * w), 0)
                yw1 = max(int(y1 - ad * h), 0)
                xw2 = min(int(x2 + ad * w), img_w - 1)
                yw2 = min(int(y2 + ad * h), img_h - 1)

                cv2.rectangle(input_img, (xw1,yw1), (xw2,yw2), (0, 0, 255), 2)
                start=time.time()
                faces[i,:,:,:] = cv2.resize(input_img[yw1:yw2 + 1, xw1:xw2 + 1, :], (img_size, img_size))
                faces[i,:,:,:] = cv2.normalize(faces[i,:,:,:], None, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX)        
                
                face = np.expand_dims(faces[i,:,:,:], axis=0)
                p_result = model.predict(face)
                print('fangxiang',time.time()-start)
                face = face.squeeze()
                img = draw_axis(input_img[yw1:yw2 + 1, xw1:xw2 + 1, :], p_result[0][0], p_result[0][1], p_result[0][2])
                
                input_img[yw1:yw2 + 1, xw1:xw2 + 1, :] = img
    
    return input_img

def main():
    os.makedirs('./img',exist_ok=True)

    img_size = 64
    img_idx = 0
    ad = 0.6
    #Parameters
    num_capsule = 3
    dim_capsule = 16
    routings = 2
    stage_num = [3,3,3]
    lambda_d = 1
    num_classes = 3
    image_size = 64
    num_primcaps = 7*3
    m_dim = 5
    S_set = [num_capsule, dim_capsule, routings, num_primcaps, m_dim]

    model1 = FSA_net_Capsule(image_size, num_classes, stage_num, lambda_d, S_set)()
    model2 = FSA_net_Var_Capsule(image_size, num_classes, stage_num, lambda_d, S_set)()
    
    num_primcaps = 8*8*3
    S_set = [num_capsule, dim_capsule, routings, num_primcaps, m_dim]

    model3 = FSA_net_noS_Capsule(image_size, num_classes, stage_num, lambda_d, S_set)()
    
    weight_file1 = '../pre-trained/300W_LP_models/fsanet_capsule_3_16_2_21_5/fsanet_capsule_3_16_2_21_5.h5'
    model1.load_weights(weight_file1)
    print('Finished loading model 1.')
    
    weight_file2 = '../pre-trained/300W_LP_models/fsanet_var_capsule_3_16_2_21_5/fsanet_var_capsule_3_16_2_21_5.h5'
    weight_file3 = '../pre-trained/300W_LP_models/fsanet_noS_capsule_3_16_2_192_5/fsanet_noS_capsule_3_16_2_192_5.h5'
    model2.load_weights(weight_file2)
    print('Finished loading model 2.')


    model3.load_weights(weight_file3)
    print('Finished loading model 3.')

    inputs = Input(shape=(64,64,3))
    x1 = model1(inputs) #1x1
    x2 = model2(inputs) #var
    x3 = model3(inputs) #w/o
    avg_model = Average()([x1,x2,x3])
    model = Model(inputs=inputs, outputs=avg_model)

    # load our serialized face detector from disk
    print("[INFO] loading face detector...")
    protoPath = os.path.sep.join(["face_detector", "deploy.prototxt"])
    modelPath = os.path.sep.join(["face_detector",
        "res10_300x300_ssd_iter_140000.caffemodel"])
    net = cv2.dnn.readNetFromCaffe(protoPath, modelPath)

    # capture video
    cap = cv2.VideoCapture(0)
    # cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1024*1)
    # cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 768*1)

    while True:
        # get video frame
        ret, input_img = cap.read()

        img_idx = img_idx + 1
        img_h, img_w, _ = np.shape(input_img)

        blob = cv2.dnn.blobFromImage(cv2.resize(input_img, (300, 300)), 1.0,
            (300, 300), (104.0, 177.0, 123.0))
        net.setInput(blob)
        detected = net.forward()

        faces = np.empty((detected.shape[2], img_size, img_size, 3))

        input_img = draw_results_ssd(detected,input_img,faces,ad,img_size,img_w,img_h,model)
            # cv2.imwrite('img/'+str(img_idx)+'.png',input_img)

        cv2.imshow("result", input_img)
        key = cv2.waitKey(1)
        
if __name__ == '__main__':
    main()

[((6, 0, 6, 21), 'sys.path.append', 'sys.path.append', ({(6, 16, 6, 20): '""".."""'}, {}), "('..')", False, 'import sys\n'), ((94, 4, 94, 38), 'os.makedirs', 'os.makedirs', (), '', False, 'import os\n'), ((141, 16, 141, 70), 'os.path.sep.join', 'os.path.sep.join', ({(141, 33, 141, 69): "['face_detector', 'deploy.prototxt']"}, {}), "(['face_detector', 'deploy.prototxt'])", False, 'import os\n'), ((142, 16, 143, 52), 'os.path.sep.join', 'os.path.sep.join', ({(142, 33, 143, 51): "['face_detector', 'res10_300x300_ssd_iter_140000.caffemodel']"}, {}), "(['face_detector', 'res10_300x300_ssd_iter_140000.caffemodel'])", False, 'import os\n'), ((144, 10, 144, 56), 'cv2.dnn.readNetFromCaffe', 'cv2.dnn.readNetFromCaffe', ({(144, 35, 144, 44): 'protoPath', (144, 46, 144, 55): 'modelPath'}, {}), '(protoPath, modelPath)', False, 'import cv2\n'), ((147, 10, 147, 29), 'cv2.VideoCapture', 'cv2.VideoCapture', ({(147, 27, 147, 28): '0'}, {}), '(0)', False, 'import cv2\n'), ((136, 16, 136, 25), 'keras.layers.Average', 'Average', ({}, {}), '()', False, 'from keras.layers import Average\n'), ((156, 26, 156, 45), 'numpy.shape', 'np.shape', ({(156, 35, 156, 44): 'input_img'}, {}), '(input_img)', True, 'import numpy as np\n'), ((163, 16, 163, 68), 'numpy.empty', 'np.empty', ({(163, 25, 163, 67): '(detected.shape[2], img_size, img_size, 3)'}, {}), '((detected.shape[2], img_size, img_size, 3))', True, 'import numpy as np\n'), ((168, 8, 168, 39), 'cv2.imshow', 'cv2.imshow', ({(168, 19, 168, 27): '"""result"""', (168, 29, 168, 38): 'input_img'}, {}), "('result', input_img)", False, 'import cv2\n'), ((169, 14, 169, 28), 'cv2.waitKey', 'cv2.waitKey', ({(169, 26, 169, 27): '1'}, {}), '(1)', False, 'import cv2\n'), ((39, 17, 39, 25), 'math.sin', 'sin', ({(39, 21, 39, 24): 'yaw'}, {}), '(yaw)', False, 'from math import cos, sin\n'), ((158, 37, 158, 70), 'cv2.resize', 'cv2.resize', ({(158, 48, 158, 57): 'input_img', (158, 59, 158, 69): '(300, 300)'}, {}), '(input_img, (300, 300))', False, 'import cv2\n'), ((30, 17, 30, 25), 'math.cos', 'cos', ({(30, 21, 30, 24): 'yaw'}, {}), '(yaw)', False, 'from math import cos, sin\n'), ((30, 28, 30, 37), 'math.cos', 'cos', ({(30, 32, 30, 36): 'roll'}, {}), '(roll)', False, 'from math import cos, sin\n'), ((35, 29, 35, 38), 'math.sin', 'sin', ({(35, 33, 35, 37): 'roll'}, {}), '(roll)', False, 'from math import cos, sin\n'), ((40, 29, 40, 39), 'math.sin', 'sin', ({(40, 33, 40, 38): 'pitch'}, {}), '(pitch)', False, 'from math import cos, sin\n'), ((78, 16, 78, 78), 'cv2.rectangle', 'cv2.rectangle', ({(78, 30, 78, 39): 'input_img', (78, 41, 78, 50): '(xw1, yw1)', (78, 52, 78, 61): '(xw2, yw2)', (78, 63, 78, 74): '(0, 0, 255)', (78, 76, 78, 77): '(2)'}, {}), '(input_img, (xw1, yw1), (xw2, yw2), (0, 0, 255), 2)', False, 'import cv2\n'), ((79, 22, 79, 33), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((80, 33, 80, 105), 'cv2.resize', 'cv2.resize', ({(80, 44, 80, 82): 'input_img[yw1:yw2 + 1, xw1:xw2 + 1, :]', (80, 84, 80, 104): '(img_size, img_size)'}, {}), '(input_img[yw1:yw2 + 1, xw1:xw2 + 1, :], (img_size, img_size))', False, 'import cv2\n'), ((81, 33, 81, 114), 'cv2.normalize', 'cv2.normalize', (), '', False, 'import cv2\n'), ((83, 23, 83, 61), 'numpy.expand_dims', 'np.expand_dims', (), '', True, 'import numpy as np\n'), ((31, 17, 31, 27), 'math.cos', 'cos', ({(31, 21, 31, 26): 'pitch'}, {}), '(pitch)', False, 'from math import cos, sin\n'), ((31, 30, 31, 39), 'math.sin', 'sin', ({(31, 34, 31, 38): 'roll'}, {}), '(roll)', False, 'from math import cos, sin\n'), ((31, 67, 31, 75), 'math.sin', 'sin', ({(31, 71, 31, 74): 'yaw'}, {}), '(yaw)', False, 'from math import cos, sin\n'), ((35, 18, 35, 26), 'math.cos', 'cos', ({(35, 22, 35, 25): 'yaw'}, {}), '(yaw)', False, 'from math import cos, sin\n'), ((36, 17, 36, 27), 'math.cos', 'cos', ({(36, 21, 36, 26): 'pitch'}, {}), '(pitch)', False, 'from math import cos, sin\n'), ((36, 30, 36, 39), 'math.cos', 'cos', ({(36, 34, 36, 38): 'roll'}, {}), '(roll)', False, 'from math import cos, sin\n'), ((36, 66, 36, 75), 'math.sin', 'sin', ({(36, 70, 36, 74): 'roll'}, {}), '(roll)', False, 'from math import cos, sin\n'), ((40, 18, 40, 26), 'math.cos', 'cos', ({(40, 22, 40, 25): 'yaw'}, {}), '(yaw)', False, 'from math import cos, sin\n'), ((62, 47, 62, 73), 'numpy.array', 'np.array', ({(62, 56, 62, 72): '[w0, h0, w0, h0]'}, {}), '([w0, h0, w0, h0])', True, 'import numpy as np\n'), ((31, 42, 31, 51), 'math.cos', 'cos', ({(31, 46, 31, 50): 'roll'}, {}), '(roll)', False, 'from math import cos, sin\n'), ((31, 54, 31, 64), 'math.sin', 'sin', ({(31, 58, 31, 63): 'pitch'}, {}), '(pitch)', False, 'from math import cos, sin\n'), ((36, 42, 36, 52), 'math.sin', 'sin', ({(36, 46, 36, 51): 'pitch'}, {}), '(pitch)', False, 'from math import cos, sin\n'), ((36, 55, 36, 63), 'math.sin', 'sin', ({(36, 59, 36, 62): 'yaw'}, {}), '(yaw)', False, 'from math import cos, sin\n'), ((85, 34, 85, 45), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n')]