# This is a sample Python script.

# Press Shift+F10 to execute it or replace it with your code.
# Press Double Shift to search everywhere for classes, files, tool windows, actions, and settings.
import cv2
import numpy as np
from anchor_generator import generate_anchors
from anchor_decode import decode_bbox
from nms import single_class_non_max_suppression

feature_map_sizes = [[33, 33], [17, 17], [9, 9], [5, 5], [3, 3]]
anchor_sizes = [[0.04, 0.056], [0.08, 0.11], [0.16, 0.22], [0.32, 0.45], [0.64, 0.72]]
anchor_ratios = [[1, 0.62, 0.42]] * 5

anchors = generate_anchors(feature_map_sizes,anchor_sizes,anchor_ratios)
anchors_exp = np.expand_dims(anchors,axis=0)
id2class = {0:'Mask',1:'NoMask'}


def face_mask_detection_caffe(conf_thresh=0.5,iou_thresh=0.4):
    output_info = []
    model_txt = "D:/vs2019Proj/ConsoleApplication1/face_mask_detection/face_mask_detection.prototxt"
    model_bin = "D:/vs2019Proj/ConsoleApplication1/face_mask_detection/face_mask_detection.caffemodel"
    net = cv2.dnn.readNetFromCaffe(model_txt,model_bin)
    print(type(net))
    if net.empty():
        return;
    # cap = cv2.VideoCapture(0)
    cap = cv2.VideoCapture("mask.jpg")
    # img = cv2.imread("mask.jpg")
    if not cap.isOpened():
        return ;
    height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT)
    width = cap.get(cv2.CAP_PROP_FRAME_WIDTH)
    net.setPreferableBackend(cv2.dnn.DNN_BACKEND_OPENCV)
    net.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU)
    while True:
        status,img = cap.read()
        if not status:
            break ;
        blob = cv2.dnn.blobFromImage(img, 1 / 255.0, (260, 260), (104.0, 177.0, 123.0))
        net.setInput(blob,"data")
        outname = net.getUnconnectedOutLayersNames()
        detection = net.forward(outname)
        # print(type(detection))
        y_bboxes = decode_bbox(anchors_exp,detection[0])[0]
        # print(y_bboxes)
        y_cls = detection[1][0]
        bbox_max_scores = np.max(y_cls,axis=1)
        bbox_max_score_classes = np.argmax(y_cls,axis=1)
        keep_idxs = single_class_non_max_suppression(y_bboxes,bbox_max_scores)
        for idx in keep_idxs:
            conf = float(bbox_max_scores[idx])
            class_id = bbox_max_score_classes[idx]
            bbox = y_bboxes[idx]
            xmin = max(0,int(bbox[0]*width))
            ymin = max(0,int(bbox[0]*height))
            xmax = min(int(bbox[2]*width),int(width))
            ymax = min(int(bbox[3]*height),int(height))
            if class_id == 0:
                color = (0,255,0)
            else:
                color = (255,0,0)
            cv2.rectangle(img,(xmin,ymin),(xmax,ymax),color,2)
            cv2.putText(img,"%s:%.2f"%(id2class[class_id],conf),(xmin+2,ymin-2),cv2.FONT_HERSHEY_SIMPLEX,0.8,color)
            # output_info.append([class_id, conf, xmin, ymin, xmax, ymax])
        cv2.imshow("img",img)
        c = cv2.waitKey(10)
        if c == 27:
            break
    return output_info


if __name__ == '__main__':
    info = face_mask_detection_caffe()
    print(info)