from sklearn.neighbors import KNeighborsClassifier
from sklearn.model_selection import train_test_split
import numpy as np
import pickle
import cv2 as cv
import matplotlib.pyplot as plt
import time
from numpy.random import randint
from os.path import exists
from imutils import paths 
import os
from PIL import ImageDraw

if not os.path.exists('X'):
    XX = []
    yy = []
    for f in os.listdir():
        if ".jpg" not in f:  # 如果不是jpg图片，则跳过
            continue
        img = cv.imread(f, 0)
        img_gray = cv.equalizeHist(img)  # 对图像进行直方图增强
        # 1. 创建级联分类器
        face_cascade = cv.CascadeClassifier()
        # 2. 引入训练好的可用于人脸识别的级联分类器模型
        face_cascade.load("haarcascade_frontalface_alt.xml")
        # 3. 用此级联分类器识别图像中的所有人脸信息，返回一个包含有所有识别的人联系系的列表
        # 列表中每一个元素包含四个值：面部左上角的坐标(x,y) 以及面部的宽和高(w,h)
        faces = face_cascade.detectMultiScale(img_gray)
        # 4. 为图像中的所有面部画框
        for (x, y, w, h) in faces:
            img1=cv.rectangle(img, (x, y), (x + w, y + h), (0, 255, 0), 2)
            img0=img[y:y + h, x:x + w]
            df1=cv.resize(img0, (300, 300), interpolation=cv.INTER_NEAREST)
            if "cxk" in f:
                yy.append(0)
            elif "dlrb" in f:
                yy.append(1)
            elif "hg" in f:
                yy.append(2)
            elif "lh" in f:
                yy.append(3)
            else:
                yy.append(4)
            XX.append(df1.ravel())
        X = np.array(XX,dtype=object)
        y = np.array(yy)
    with open("X", 'wb') as f:
        pickle.dump(X, f)
    with open("y", 'wb') as f:
        pickle.dump(y, f)

acc_data=[]
acc_list=[]
def get_max_K(X_train, X_test, y_train, y_test):
    start = time.time()
    max_k=0
    max_acc=0
    for K in range(1,10):
        knn = KNeighborsClassifier(n_neighbors=K)
        knn.fit(X_train, y_train)
        acc=knn.score(X_test,y_test)
        acc_list.append(acc)
        acc_data.append(K)
        if acc > max_acc:
            max_acc=acc
            max_k=K
    end = time.time()
    duration=(end-start)/60
    return max_acc, max_k,duration

# 1. 获取X和y
with open("X", 'rb') as f:
   X = pickle.load(f)
with open("y", 'rb') as f:
   y = pickle.load(f)

# 2. train_test_split
X_train, X_test, y_train, y_test = train_test_split(X,y,test_size=0.25)

# 3. 创建KNN模型
max_acc, K, duration = get_max_K(X_train, X_test, y_train, y_test)
h1 = np.round(max_acc * 100, 2)
knn = KNeighborsClassifier(n_neighbors=K)
knn.fit(X_train, y_train)

labels = ['cxk','dlrb','hg','lh','zly']
d1 = randint(0,133)
d2 = X[d1]
predict = knn.predict([d2])

print("准确率最高为:", h1,"用时: {0:.2f}分钟".format(duration))
# 随机抽取一张图片
img_file = np.random.choice(list(paths.list_images(".")),1)[0]

# 把图片转化成特征向量
img=cv.imread(img_file,1)
img_gray= cv.cvtColor(img, cv.COLOR_BGR2GRAY)
img_gray = cv.equalizeHist(img_gray) #对图像进行直方图增强

# 1. 创建级联分类器
face_cascade = cv.CascadeClassifier()
# 2. 引入训练好的可用于人脸识别的级联分类器模型
face_cascade.load("haarcascade_frontalface_alt.xml")
# 3. 用此级联分类器识别图像中的所有人脸信息，返回一个包含有所有识别的人联系系的列表
# 列表中每一个元素包含四个值：面部左上角的坐标(x,y) 以及面部的宽和高(w,h)
faces = face_cascade.detectMultiScale(img_gray)
# 4. 为图像中的所有面部画框
for (x,y,w,h) in faces:
    cv.rectangle(img, (x, y), (x+w, y+h), (0, 255, 0), 2)

    cv.putText(img,  # 要显示字体的图片
               labels[predict[0]] ,# 要显示的内容
               (x, y - 10),  # 要显示的位置
               cv.FONT_HERSHEY_SIMPLEX,  # 要使用的字体 -> 一般英文字体
               1,  # 字体放大倍数
               (0, 255, 0),  # 字体颜色
               2)  # 字体线条粗细


#plt.rcParams['font.sans-serif'] = ['SimHei']  # 黑体
#plt.rcParams['axes.unicode_minus'] = False    # 解决无法显示符号的问题
#plt.title('{0:.2f}, 准确率最高为: {1:.2f}, 用时: {2:.2f}分钟'.format(labels[predict[0]], h1, duration))
plt.imshow(cv.cvtColor(img,cv.COLOR_BGR2RGB))
plt.tight_layout()
plt.show()