import time

import cv2

from MNPoseDetection.consensus import *
from MNPoseDetection.pose_detection import PoseDetection


if __name__ == "__main__":
    pd = PoseDetection()
    # 读取视频文件
    cap = cv2.VideoCapture(0)

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break    # 显示结果
        start_time = time.time()
        result = pd.detect(frame)
        # print(result)
        for label, info in result.items():
            keypoint = info[KEY_POINTS]
            reliable = info[RELIABLE]
            confidence = info[CONFIDENCE]
            pass_test = info[PASS_TEST]
            if not reliable or not pass_test:
                continue
            cv2.circle(frame, keypoint, 5, (0, 255, 0), -1)
            cv2.putText(
                frame, "{}".format(label), keypoint,
                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1, cv2.LINE_AA
            )
        cv2.imshow('MoveNet Pose Estimation', frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
        print(time.time() - start_time)

    cap.release()
    cv2.destroyAllWindows()