import cv2
import numpy as np
import dlib
import os
import matplotlib.pyplot as plt

def AffineTransform(dst_img, img, src_point, dst_point):
    rows, cols, ch = dst_img.shape
    M = cv2.getAffineTransform(src_point, dst_point)
    dst = cv2.warpAffine(img, M, (cols, rows))
    return dst

def face_cut(img):
    detector = cv2.CascadeClassifier('E:/data/haarcascade_frontalface_default.xml')
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    faces = detector.detectMultiScale(gray, 1.1, 3)
    x = faces[0][0]
    y = faces[0][1]
    w = faces[0][2]
    h = faces[0][3]
    img1 = img[y:y + h, x:x + w]
    return img1

def face_detector(frame):
    detector = dlib.get_frontal_face_detector()
    predictor = dlib.shape_predictor('E:/shape_predictor_68_face_landmarks.dat')
    face_key_point = []
    # cv2读取图像
    # 取灰度
    img_gray = frame
    img = frame
    # 人脸数rects
    rects = detector(img_gray, 0)
    face_key_point = np.empty([0,1,2], dtype=np.float32)
    for i in range(len(rects)):
        landmarks = np.matrix([[p.x, p.y] for p in predictor(img,rects[i]).parts()])
        for idx, point in enumerate(landmarks):
            # 68点的坐标
            pos = (point[0, 0], point[0, 1])
            #print(idx,pos)
            temp_point = np.empty([1,1,2], dtype=np.float32)
            temp_point[0,0] = pos
            face_key_point = np.concatenate([face_key_point, temp_point])

            # 利用cv2.circle给每个特征点画一个圈，共68个
            cv2.circle(img, pos, 5, color=(0, 255, 0))
            # 利用cv2.putText输出1-68
            font = cv2.FONT_HERSHEY_SIMPLEX
            cv2.putText(img, str(idx+1), pos, font, 0.8, (0, 0, 255), 1,cv2.LINE_AA)

    return face_key_point

def calcOpticalFlow(old_gray, frame_gray, p0):
    lk_params = dict(winSize=(15,15),
                 maxLevel=2)
    p1, st, err = cv2.calcOpticalFlowPyrLK(old_gray, frame_gray, p0, None, **lk_params)
    good_old = p0[st==1]
    good_new = p1[st==1]
    return good_old, good_new

def addOpticalFlow(vectors):
    n = vectors.shape[0]
    result = [0,0]
    for i in range(n):
        result += vectors[i]
    return result

def imshow(name, img):
    cv2.imshow(name, img)
    k = cv2.waitKey(0)&0xff
    return k


def pre_process(initial_img, img, first, p0):
    a = 2
    b = 16
    c = 34
    img = face_cut(img)
    if first == True:
        p0 = face_detector(img)
        return p0, img
    else:
        p1 = face_detector(img)
        dst_point = np.float32([p0[a,0], p0[b,0], p0[c,0]])
        src_point = np.float32([p1[a,0], p1[b,0], p1[c,0]])

        temp_frame = AffineTransform(initial_img, img, src_point, dst_point)
        return temp_frame

def viz_flow(flow):
    # 色调H：用角度度量，取值范围为0°～360°，从红色开始按逆时针方向计算，红色为0°，绿色为120°,蓝色为240°
    # 饱和度S：取值范围为0.0～1.0
    # 亮度V：取值范围为0.0(黑色)～1.0(白色)
    h, w = flow.shape[:2]
    hsv = np.zeros((h, w, 3), np.uint8)
    mag, ang = cv2.cartToPolar(flow[...,0], flow[...,1])
    hsv[...,0] = ang*180/np.pi/2
    hsv[...,1] = cv2.normalize(mag,None,0,255,cv2.NORM_MINMAX)
    # flownet是将V赋值为255, 此函数遵循flownet，饱和度S代表像素位移的大小，亮度都为最大，便于观看
    # 也有的光流可视化讲s赋值为255，亮度代表像素位移的大小，整个图片会很暗，很少这样用
    hsv[...,2] = 255
    bgr = cv2.cvtColor(hsv,cv2.COLOR_HSV2BGR)
    return bgr

def show_gray_hist(img):
    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    plt.hist(img.ravel(), 256)
    plt.show()

def show_bgr_hist(img):
    color = ('b', 'g', 'r')
    for i,col in enumerate(color):
        histr = cv2.calcHist([img], [i], None, [256], [0,256])
        plt.plot(histr, color =col)
        plt.xlim([0,256])
    plt.show()

def get_path():
    path1 = 'E:/CAS(ME)^2/rawvideo/s19/'
    path2 = 'E:/ZLW_generate_flow/s19/'
    path_list1 = []
    path_list2 = []
    #counnt=1表示进入到第一个目录下
    count = 1
    for root, dirs, files in os.walk(path1):
        if count == 1:
            #print(root)
            #print(dirs)
            #print(files)
            for i in range(len(files)):
                temp_path = path1 + files[i]
                path_list1.append(temp_path)
                temp_path = path2 + files[i][:-4]
                path_list2.append(temp_path)
            break
        count+=1
    return path_list1, path_list2

a = 2
b = 16
c = 34
file_path_915 = 'E:/s19_debug1/s19/19_0402beatingpregnantwoman/915.jpg'
file_path_1 = 'E:/s19_debug1/s19/19_0402beatingpregnantwoman/1.jpg'
img_915 = cv2.imread(file_path_915)
img_1 = cv2.imread(file_path_1)
face_cut = face_cut(img_915)
#imshow('face_cut', face_cut)
p1 = face_detector(face_cut)
#imshow('face_cut', face_cut)
p0 = face_detector(img_1)
src_point = np.float32([p1[a,0], p1[b,0], p1[c,0]])
dst_point = np.float32([p0[a,0], p0[b,0], p0[c,0]])
dst = AffineTransform(img_1, face_cut, src_point, dst_point)
if not os.path.exists('E:/debug_s19_915'):
    os.mkdir('E:/debug_s19_915')
cv2.imwrite('E:/debug_s19_915/' + '1.jpg', img_1)
cv2.imwrite('E:/debug_s19_915/' + '915.jpg', face_cut)
cv2.imwrite('E:/debug_s19_915/' + 'dst.jpg', dst)
imshow('dst', dst)