chrisvlds/MLProjectV3

1import pickle2import tensorflow as tf3import numpy as np4from matplotlib import pyplot as plt5import cv26import wget7import math8 9print("opencv v is:" + cv2.__version__)10 11interpreter = tf.lite.Interpreter(model_path='lite-model_movenet_singlepose_lightning_3.tflite')12interpreter.allocate_tensors()13loaded_model = pickle.load(open('model1.pkl', 'rb'))14 15 16def draw_keypoints(frame, keypoints, confidence_threshold):17    y, x, c = frame.shape18    shaped = np.squeeze(np.multiply(keypoints, [y, x, 1]))19 20    for kp in shaped:21        ky, kx, kp_conf = kp22        if kp_conf > confidence_threshold:23            cv2.circle(frame, (int(kx), int(ky)), 4, (0, 255, 0), -1)24 25 26EDGES = {27    (11, 12): 'y',28    (11, 13): 'm',29}30 31 32def draw_connections(frame, keypoints, edges, confidence_threshold):33    y, x, c = frame.shape34    shaped = np.squeeze(np.multiply(keypoints, [y, x, 1]))35 36    for edge, color in edges.items():37        p1, p2 = edge38        y1, x1, c1 = shaped[p1]39        y2, x2, c2 = shaped[p2]40 41        if (c1 > confidence_threshold) & (c2 > confidence_threshold):42            cv2.line(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 0, 255), 2)43 44 45repCounterUp = 046repCounterDown = 047repCounter = 048state = 149cap = cv2.VideoCapture('projectV2trim.webm')50width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))51height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))52result = cv2.VideoWriter('iddkkk.webm', cv2.VideoWriter_fourcc(*'VP90'), 20, (width, height))53while cap.isOpened():54    ret, frame = cap.read()55    # Reshape image56    if not ret:57        break58    img = frame.copy()59    img = tf.image.resize_with_pad(np.expand_dims(img, axis=0), 192, 192)60    input_image = tf.cast(img, dtype=tf.float32)61 62    # Setup input and output63    input_details = interpreter.get_input_details()64    output_details = interpreter.get_output_details()65 66    # Make predictions67    interpreter.set_tensor(input_details[0]['index'], np.array(input_image))68    interpreter.invoke()69    keypoints_with_scores = interpreter.get_tensor(output_details[0]['index'])70    #print(keypoints_with_scores)71 72    # Rendering73    draw_connections(frame, keypoints_with_scores, EDGES, 0.4)74    draw_keypoints(frame, keypoints_with_scores, 0.4)75    result.write(frame)76    shaped = np.squeeze(77        np.multiply(interpreter.get_tensor(interpreter.get_output_details()[0]['index']), [480, 640, 1]))78 79    for kp in shaped:80        ky, kx, kp_conf = kp81        # print(int(ky), int(kx), kp_conf)82 83    shaped[0], shaped[1]84 85    for edge, color in EDGES.items():86        p1, p2 = edge87        y1, x1, c1 = shaped[p1]88        y2, x2, c2 = shaped[p2]89        # print((int(x2), int(y2)))90        input = np.array([[x1, y1, x2, y2]])91 92    y = loaded_model.predict(input)93    y = y[0]94    prediction = round(y[0])95 96    print("Predicted=%s" % result)97    if prediction == 0:98        print("Down")99        if repCounterDown < 10:100            repCounterDown += 1101        if repCounterUp > 0:102            repCounterUp -= 1103    else:104        print("Up")105        if repCounterUp < 10:106            repCounterUp += 1107        if repCounterDown > 0:108            repCounterDown -= 1109 110    if repCounterDown == 10 and repCounterUp == 0:111        if state == 1:112            state = 0113            repCounter += 1114    elif repCounterUp == 10 and repCounterDown == 0:115        if state == 0:116            state = 1117            repCounter += 1118 119    reps = math.floor(repCounter / 2)120    frame = cv2.rectangle(frame, (150, 0), (900, 200), (0, 0, 0), -1)121    frame = cv2.putText(frame, 'Reps: '+str(reps), (150, 150), cv2.FONT_HERSHEY_SIMPLEX, 5, (255, 0, 0), 5, cv2.LINE_AA)122    cv2.imshow('MoveNet Lightning', frame)123 124    if cv2.waitKey(10) & 0xFF == ord('q'):125        break126 127cap.release()128result.release()129cv2.destroyAllWindows()130 131left_hip = keypoints_with_scores[0][0][11]132left_knee = keypoints_with_scores[0][0][13]133 134shaped = np.squeeze(np.multiply(interpreter.get_tensor(interpreter.get_output_details()[0]['index']), [480, 640, 1]))135 136# for kp in shaped:137#     ky, kx, kp_conf = kp138#     #print(int(ky), int(kx), kp_conf)139#140# shaped[0], shaped[1]141#142# for edge, color in EDGES.items():143#     p1, p2 = edge144#     y1, x1, c1 = shaped[p1]145#     y2, x2, c2 = shaped[p2]146#     #print((int(x2), int(y2)))147#     input = np.array([[x1, y1, x2, y2]])148reps = math.floor(repCounter/2)149print("Number of reps: "+str(reps))150