import cv2
from PIL import Image
import torch
import torchvision
from torchvision import datasets, models, transforms
import torch.nn.functional as F
import torch.nn as nn
import argparse

tsfrm = transforms.Compose([
    transforms.Grayscale(3),
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])

# 第一个人、第二个人
classes = ('pc', 'pc1')

# 判断GPU是否可用
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)


# 模型推理
def modelpre(frame):
    # 图像灰度化
    grey_img = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    faces = face_cascade.detectMultiScale(grey_img,1.5,6)
    for (x,y,w,h) in faces:
        center = (x + w//2, y + h//2)
        frame = cv2.rectangle(frame, (x,y), (x+w, y+h),  (255, 0, 255),4)
        faceROI = grey_img[y:y+h,x:x+w]
        model_eval = models.resnet18(pretrained=False)
        num_ftrs = model_eval.fc.in_features
        model_eval.fc = nn.Linear(num_ftrs, 2)
        model_eval.load_state_dict(torch.load('./model/Model-face.pkl', map_location=device))
        # 在推理前，务必调用model.eval()去设置dropout和batch normalization层为评估模式
        model_eval.eval()
        # OpenCV转PIL格式
        image = Image.fromarray(cv2.cvtColor(faceROI, cv2.COLOR_GRAY2RGB))
        # PIL图像数据转换为tensor数据，并归一化
        img = tsfrm(image)
        # 图像增加1维[batch_size,通道,高,宽]
        img = img.unsqueeze(0)
        # 输出推理结果
        output = model_eval(img)
        value, predicted = torch.max(output.data, 1)
        label = predicted.numpy()[0]
        ps = torch.exp(output)
        print(ps)

        print(label)
        if label==0:
            frame=cv2.putText(frame, 'lwh',
                (x, y),                                #坐标
                cv2.FONT_HERSHEY_SIMPLEX,              #字体
                1,                                     #字号
                (255,0,255),                           #颜色
                2)                                     #字的线宽)
        cv2.imshow('Face', frame)


if __name__ == "__main__":
    # -- 读取视频流
    cap = cv2.VideoCapture(0)
    parser = argparse.ArgumentParser(description='Code for Cascade Classifier tutorial.')
    parser.add_argument('--face_cascade', help='Path to face cascade.', default='haarcascade_frontalface_alt.xml')
    args = parser.parse_args()
    face_cascade_name = args.face_cascade
    face_cascade = cv2.CascadeClassifier()
    if not face_cascade.load(cv2.samples.findFile(face_cascade_name)):
        print('--(!)Error loading face cascade')
        exit(0)

    if not cap.isOpened:
        print('--(!)Error opening video capture')
        exit(0)
    while True:
        ret, frame = cap.read()
        modelpre(frame)
        if cv2.waitKey(10) == 27:
            break