import time
from centerface import CenterFace
from skimage import transform as trans
import numpy as np
import torch
import cv2
from backbones import iresnet100, iresnet18
from create_database import findOne, load_npy,findAll
from PIL import Image, ImageDraw,ImageFont

def show():
    cap = cv2.VideoCapture("test.mp4")
    ret, frame = cap.read()
    h, w = frame.shape[:2]
    centerface = CenterFace()
    size = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
            int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
    out = cv2.VideoWriter('ccvt6.mp4', cv2.VideoWriter_fourcc('m', 'p', '4', 'v'), 30, size)
    while ret:
        start_time = time.time()
        dets, lms = centerface(frame, h, w, threshold=0.35)
        end_time = time.time()
        print("findOne time: " + str(end_time - start_time))
        for det in dets:
            boxes, score = det[:4], det[4]
            cv2.rectangle(frame, (int(boxes[0]), int(boxes[1])), (int(boxes[2]), int(boxes[3])), (2, 255, 0), 1)
        for lm in lms:
            for i in range(0, 5):
                cv2.circle(frame, (int(lm[i * 2]), int(lm[i * 2 + 1])), 2, (0, 0, 255), -1)
        cv2.imshow('out', frame)
        out.write(frame)
        # Press Q on keyboard to stop recording
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
        ret, frame = cap.read()
    cap.release()
    out.release()
    cv2.destroyAllWindows()

def video():
    model = iresnet100()
    model.load_state_dict(torch.load("./model/backbone100.pth", map_location="cpu"))
    model.eval()
    k_v = load_npy("student.npy")
    count = 0
    #cap = cv2.VideoCapture("http://ivi.bupt.edu.cn/hls/cctv6hd.m3u8")
    cap = cv2.VideoCapture("software.mp4")
    ret, frame = cap.read()
    h, w = frame.shape[:2]
    size = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
            int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
    fps = cap.get(cv2.CAP_PROP_FPS)
    out = cv2.VideoWriter('ttt.mp4', cv2.VideoWriter_fourcc('m', 'p', '4', 'v'), fps, size)
    centerface = CenterFace()
    while ret:
        start_time = time.time()
        dets, lms = centerface(frame, h, w, threshold=0.35)
        end_time = time.time()
        print("detectOneframe time: " + str(end_time - start_time))
        face_list = []
        name_list = []
        for i,det in enumerate(dets):
            boxes, score = det[:4], det[4]
            img_w = int(boxes[2] - boxes[0])
            img_h = int(boxes[3] - boxes[1])
            distace = int(abs(img_w - img_h) / 2)
            img_w1 = int(boxes[0]) - distace
            img_w2 = int(boxes[2]) + distace
            # print(img_w,img_h,distace,max_hw)
            if img_w <= img_h and img_w1 >= 0 and img_w2 <= frame.shape[1]:
                img112 = frame[int(boxes[1]):int(boxes[3]), img_w1:img_w2, :]
                img112 = cv2.resize(img112, (112, 112))
                # cv2.imwrite("./img/man"+str(count)+".jpg", img112)
                # count += 1
                face_list.append(img112)
            else:
                img112 = frame[int(boxes[1]):int(boxes[3]), int(boxes[0]):int(boxes[2]), :]
                img112 = cv2.resize(img112, (112, 112))
                face_list.append(img112)
        if len(face_list) != 0:
            face_list = np.array(face_list)
            face_list = face_list.transpose((0,3,1,2))
            face_list = np.array(face_list, dtype=np.float32)
            face_list -= 127.5
            face_list /= 127.5
            print(face_list.shape)
            face_list = torch.from_numpy(face_list)
            start_time = time.time()

            for face in face_list:
                face = face[np.newaxis, :, :, :]

                name_list.append(findOne(face,model,k_v))
            end_time = time.time()
            print("findOneframe time: "+str(end_time-start_time))
        img_PIL = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
        draw = ImageDraw.Draw(img_PIL)
        font = ImageFont.truetype("font.ttf",12)
        for i,det in enumerate(dets):
            boxes, score = det[:4], det[4]
            # cv2.rectangle(frame, (int(boxes[0]), int(boxes[1])), (int(boxes[2]), int(boxes[3])), (2, 255, 0), 1)
            # cv2.putText(frame, name_list[i], (int(boxes[0]), int(boxes[1])), cv2.FONT_HERSHEY_COMPLEX, 0.4,
            #             (0, 225, 255), 1)
            name = name_list[i][:3]
            if not isinstance(name, np.unicode):
                name = name.decode('utf8')
            draw.text((int(boxes[0]), int(boxes[1])),name,fill=(0, 225, 255),font=font)
            draw.rectangle((int(boxes[0]), int(boxes[1]),int(boxes[2]), int(boxes[3])),outline="green",width=1)
        frame = cv2.cvtColor(np.asarray(img_PIL),cv2.COLOR_RGB2BGR)
        cv2.imshow('out', frame)
        out.write(frame)
        # Press Q on keyboard to stop recording
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
        ret, frame = cap.read()
    cap.release()
    out.release()
    cv2.destroyAllWindows()
def video_GPU():
    model = iresnet100()
    model.load_state_dict(torch.load("./model/backbone100.pth", map_location="cpu"))
    model.eval()
    k_v = load_npy("student.npy")
    count = 0
    #cap = cv2.VideoCapture("http://ivi.bupt.edu.cn/hls/cctv6hd.m3u8")
    cap = cv2.VideoCapture("software.mp4")
    ret, frame = cap.read()
    h, w = frame.shape[:2]
    size = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
            int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
    fps = cap.get(cv2.CAP_PROP_FPS)
    out = cv2.VideoWriter('ttt.mp4', cv2.VideoWriter_fourcc('m', 'p', '4', 'v'), fps, size)
    centerface = CenterFace()
    while ret:
        start_time = time.time()
        dets, lms = centerface(frame, h, w, threshold=0.35)
        end_time = time.time()
        print("detectOneframe time: " + str(end_time - start_time))
        face_list = []
        name_list = []
        for i,det in enumerate(dets):
            boxes, score = det[:4], det[4]
            img_w = int(boxes[2] - boxes[0])
            img_h = int(boxes[3] - boxes[1])
            distace = int(abs(img_w - img_h) / 2)
            img_w1 = int(boxes[0]) - distace
            img_w2 = int(boxes[2]) + distace
            # print(img_w,img_h,distace,max_hw)
            if img_w <= img_h and img_w1 >= 0 and img_w2 <= frame.shape[1]:
                img112 = frame[int(boxes[1]):int(boxes[3]), img_w1:img_w2, :]
                img112 = cv2.resize(img112, (112, 112))
                # cv2.imwrite("./img/man"+str(count)+".jpg", img112)
                # count += 1
                face_list.append(img112)
            else:
                img112 = frame[int(boxes[1]):int(boxes[3]), int(boxes[0]):int(boxes[2]), :]
                img112 = cv2.resize(img112, (112, 112))
                face_list.append(img112)
        if len(face_list) != 0:
            face_list = np.array(face_list)
            face_list = face_list.transpose((0,3,1,2))
            face_list = np.array(face_list, dtype=np.float32)
            face_list -= 127.5
            face_list /= 127.5
            print(face_list.shape)
            face_list = torch.from_numpy(face_list)
            start_time = time.time()
            name_list = findAll(face_list, model, k_v)
            # for face in face_list:
            #     face = face[np.newaxis, :, :, :]
            #
            #     name_list.append(findOne(face,model,k_v))
            end_time = time.time()
            print("findOneframe time: "+str(end_time-start_time))
        img_PIL = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
        draw = ImageDraw.Draw(img_PIL)
        font = ImageFont.truetype("font.ttf",18)
        for i,det in enumerate(dets):
            boxes, score = det[:4], det[4]
            # cv2.rectangle(frame, (int(boxes[0]), int(boxes[1])), (int(boxes[2]), int(boxes[3])), (2, 255, 0), 1)
            # cv2.putText(frame, name_list[i], (int(boxes[0]), int(boxes[1])), cv2.FONT_HERSHEY_COMPLEX, 0.4,
            #             (0, 225, 255), 1)
            name = name_list[i][:3]
            if not isinstance(name, np.unicode):
                name = name.decode('utf8')
            draw.text((int(boxes[0]), int(boxes[1])),name,fill=(255, 0, 0),font=font)
            draw.rectangle((int(boxes[0]), int(boxes[1]),int(boxes[2]), int(boxes[3])),outline="green",width=2)
        frame = cv2.cvtColor(np.asarray(img_PIL),cv2.COLOR_RGB2BGR)
        cv2.imshow('out', frame)
        out.write(frame)
        # Press Q on keyboard to stop recording
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
        ret, frame = cap.read()
    cap.release()
    out.release()
    cv2.destroyAllWindows()

def video_GPU_retinaface():
    model = iresnet100()
    model.load_state_dict(torch.load("./model/backbone100.pth", map_location="cpu"))
    model.eval()
    k_v = load_npy("student.npy")
    count = 0
    #cap = cv2.VideoCapture("http://ivi.bupt.edu.cn/hls/cctv6hd.m3u8")
    cap = cv2.VideoCapture("software.mp4")
    ret, frame = cap.read()
    h, w = frame.shape[:2]
    size = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
            int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
    fps = cap.get(cv2.CAP_PROP_FPS)
    out = cv2.VideoWriter('ttt.mp4', cv2.VideoWriter_fourcc('m', 'p', '4', 'v'), fps, size)
    centerface = CenterFace()
    while ret:
        start_time = time.time()
        dets, lms = centerface(frame, h, w, threshold=0.35)
        end_time = time.time()
        print("detectOneframe time: " + str(end_time - start_time))
        face_list = []
        name_list = []
        print(dets.shape)
        for i,det in enumerate(dets):
            boxes, score = det[:4], det[4]
            img_w = int(boxes[2] - boxes[0])
            img_h = int(boxes[3] - boxes[1])
            distace = int(abs(img_w - img_h) / 2)
            img_w1 = int(boxes[0]) - distace
            img_w2 = int(boxes[2]) + distace
            # print(img_w,img_h,distace,max_hw)
            if img_w <= img_h and img_w1 >= 0 and img_w2 <= frame.shape[1]:
                img112 = frame[int(boxes[1]):int(boxes[3]), img_w1:img_w2, :]
                img112 = cv2.resize(img112, (112, 112))
                # cv2.imwrite("./img/man"+str(count)+".jpg", img112)
                # count += 1
                face_list.append(img112)
            else:
                img112 = frame[int(boxes[1]):int(boxes[3]), int(boxes[0]):int(boxes[2]), :]
                img112 = cv2.resize(img112, (112, 112))
                face_list.append(img112)
        if len(face_list) != 0:
            face_list = np.array(face_list)
            face_list = face_list.transpose((0,3,1,2))
            face_list = np.array(face_list, dtype=np.float32)
            face_list -= 127.5
            face_list /= 127.5
            print(face_list.shape)
            face_list = torch.from_numpy(face_list)
            start_time = time.time()
            name_list = findAll(face_list, model, k_v)
            # for face in face_list:
            #     face = face[np.newaxis, :, :, :]
            #
            #     name_list.append(findOne(face,model,k_v))
            end_time = time.time()
            print("findOneframe time: "+str(end_time-start_time))
        img_PIL = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
        draw = ImageDraw.Draw(img_PIL)
        font = ImageFont.truetype("font.ttf",18)
        for i,det in enumerate(dets):
            boxes, score = det[:4], det[4]
            # cv2.rectangle(frame, (int(boxes[0]), int(boxes[1])), (int(boxes[2]), int(boxes[3])), (2, 255, 0), 1)
            # cv2.putText(frame, name_list[i], (int(boxes[0]), int(boxes[1])), cv2.FONT_HERSHEY_COMPLEX, 0.4,
            #             (0, 225, 255), 1)
            name = name_list[i][:3]
            if not isinstance(name, np.unicode):
                name = name.decode('utf8')
            draw.text((int(boxes[0]), int(boxes[1])),name,fill=(255, 0, 0),font=font)
            draw.rectangle((int(boxes[0]), int(boxes[1]),int(boxes[2]), int(boxes[3])),outline="green",width=2)
        frame = cv2.cvtColor(np.asarray(img_PIL),cv2.COLOR_RGB2BGR)
        cv2.imshow('out', frame)
        out.write(frame)
        # Press Q on keyboard to stop recording
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
        ret, frame = cap.read()
    cap.release()
    out.release()
    cv2.destroyAllWindows()


video_GPU_retinaface()
#video_GPU()
#show()