# face detector 3.2
# DNN-based face detector

import cv2
from matplotlib import pyplot as plt
import numpy as np


# visualize functions
def show_img_with_matplotlib(color_img, title, pos):
    img_rgb = color_img[:, :, ::-1]
    ax = plt.subplot(1, 1, pos)
    plt.imshow(img_rgb)
    plt.title(title, fontsize=8)
    plt.axis('off')


# load image
image = cv2.imread("picture/005.jpg")

# load the pretrained model
# net = cv2.dnn.readNetFromCaffe("deploy.prototxt", "res10_300x300_ssd_iter_140000_fp16.caffemodel")
net = cv2.dnn.readNetFromTensorflow("opencv_face_detector_uint8.pb", "opencv_face_detector.pbtxt")

# 为了获得最佳精度，必须分别对蓝色、绿色和红色通道执行 (104, 177, 123) 通道均值减法，并将图像调整为 300 x 300 的 BGR 图像，
# 在 OpenCV 中可以通过使用 cv2.dnn.blobFromImage() 函数进行此预处理：
blob = cv2.dnn.blobFromImage(image, 1.0, (300, 300), [104, 117, 123], False, False)

# 下一步是将 blob 设置为输入以获得结果，对整个网络执行前向计算以计算输出
net.setInput(blob)
detections = net.forward()
# print(detections)

detected_faces = 0
w, h = image.shape[1], image.shape[0]
print(w, h)
# 最后一步是迭代检测并绘制结果，仅在相应置信度大于最小阈值时才将其可视化：
for i in range(0, detections.shape[2]):
    # 获取当前检测结果的置信度
    confidence = detections[0, 0, i, 2]
    # 如果置信大于最小置信度，则将其可视化
    if confidence > 0.7:
        detected_faces += 1
        # 获取当前检测结果的坐标
        box = detections[0, 0, i, 3:7] * np.array([w,h,w,h])
        (startX, startY, endX, endY) = box.astype('int')
        # 绘制检测结果和置信度
        text = "{:.3f}%".format(confidence * 100)
        y = startY - 10 if startY - 10 > 10 else startY + 10
        cv2.rectangle(image, (startX, startY), (endX, endY), (255, 0, 0), 3)
        cv2.putText(image, text, (startX, y), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 0, 255), 2)

# 可视化
show_img_with_matplotlib(image, "DNN face detector:" + str(detected_faces), 1)
plt.show()