from time import time
import cv2
import numpy as np
import tensorflow as tf

def draw(frame, keypoints_with_scores):
    for item in keypoints_with_scores[0][0]:
        if item[2] > 0.3:  # confiance
            x = int(item[1]*720)  # numpy array x est le 2ème
            y = int(item[0]*720)
            cv2.circle(frame, (x, y), 6, color=(0,0,255), thickness=-1)
    return frame

models = [  "lite-model_movenet_singlepose_thunder_3.tflite",
            "lite-model_movenet_singlepose_thunder_tflite_float16_4.tflite",
            "lite-model_movenet_singlepose_thunder_tflite_int8_4.tflite"]

interpreter = tf.lite.Interpreter(model_path=models[0])
interpreter.allocate_tensors()

cv2.namedWindow('color', cv2.WND_PROP_FULLSCREEN)
cam = cv2.VideoCapture(2)
cam.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
cam.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)

t0 = time()
nbr = 0

while 1:

    nbr += 1
    ret, frame = cam.read()

    if not ret:
        continue

    # Extraction d'une image carrée [rows, columns]
    frame = frame[:, 280:1000]

    image = tf.expand_dims(frame, axis=0)
    # Resize and pad the image to keep the aspect ratio and fit the expected size.
    image = tf.image.resize_with_pad(image, 256, 256)

    # TF Lite format expects tensor type of float32.
    input_image = tf.cast(image, dtype=tf.float32)
    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()

    interpreter.set_tensor(input_details[0]['index'], input_image.numpy())
    interpreter.invoke()

    # Output is a [1, 1, 17, 3] numpy array.
    keypoints_with_scores = interpreter.get_tensor(output_details[0]['index'])

    frame = draw(frame, keypoints_with_scores)

    # Affichage de l'image
    cv2.imshow('color', frame)

    # Calcul du FPS, affichage toutes les 10 s
    if time() - t0 > 10:
        print("FPS =", round(nbr/10, 2))
        t0, nbr = time(), 0

    # Pour quitter
    if cv2.waitKey(1) == 27:
        break

cv2.destroyAllWindows()