import cv2
import mediapipe as mp

class FaceDetection():
    # initialize the face detection class with arguments from https://google.github.io/mediapipe/solutions/face_detection.html
    def __init__(self, model_selection = 0, threshold = 0.5):
        self.model_selection = model_selection
        self.threshold = threshold

        self.mp_draw = mp.solutions.drawing_utils
        self.face_detection = mp.solutions.face_detection.FaceDetection(model_selection = self.model_selection, min_detection_confidence = self.threshold)

    # gets bounding boxes using self.face_detection, returns a list of element, elment = (score, bbox_dict)
    def get_bboxs(self, frame):
        mp_detections = self.face_detection.process(frame)
        score_bboxs = []
        if mp_detections.detections:
            for detection in mp_detections.detections:
                score = detection.score[0]
                mp_bbox = detection.location_data.relative_bounding_box
                bbox_dict = {
                    'x_min': mp_bbox.xmin,
                    'y_min': mp_bbox.ymin,
                    'w': mp_bbox.width,
                    'h': mp_bbox.height
                }

                score_bboxs.append([score, bbox_dict])

        return score_bboxs

    # draws the bbox onto the frame
    def draw_bbox(self, face_probs, bbox_dict, frame, col = (255, 0, 255), gender = None, gender_score = None):
        x_min, y_min, w, h = bbox_dict.values()
        frame_h, frame_w, _ = frame.shape
        bbox = int(x_min * frame_w), int(y_min * frame_h), int(w * frame_w), int(h * frame_h)

        # prepare text, depending on what atributes we predict
        text = str(round(face_probs, 3))
        if gender:
            text = gender + ": " + str(round(gender_score, 2))

        # draw bbox
        cv2.rectangle(frame, bbox, col, 2)
        cv2.putText(frame, text, (bbox[0], bbox[1] - 10),
                    cv2.FONT_HERSHEY_COMPLEX, 0.5, col, 1)


[(11, 'mediapipe.solutions.face_detection.FaceDetection', 'mp.solutions.face_detection.FaceDetection', True, 'import mediapipe as mp\n'), (44, 'cv2.rectangle', 'cv2.rectangle', False, 'import cv2\n'), (45, 'cv2.putText', 'cv2.putText', False, 'import cv2\n')]