chrisvlds/MLProject

1import pandas as pd2import tensorflow as tf3import numpy as np4from matplotlib import pyplot as plt5import cv26import wget7import os8import glob9 10print("opencv v is:" + cv2.__version__)11 12interpreter = tf.lite.Interpreter(model_path='lite-model_movenet_singlepose_lightning_3.tflite')13interpreter.allocate_tensors()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, 13): 'm',28}29 30 31def draw_connections(frame, keypoints, edges, confidence_threshold):32    y, x, c = frame.shape33    shaped = np.squeeze(np.multiply(keypoints, [y, x, 1]))34 35    for edge, color in edges.items():36        p1, p2 = edge37        y1, x1, c1 = shaped[p1]38        y2, x2, c2 = shaped[p2]39 40        if (c1 > confidence_threshold) & (c2 > confidence_threshold):41            cv2.line(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 0, 255), 2)42 43 44data = np.array([[0, 0, 0, 0, 0]])45path = "./squat_down/*.*"46for file in glob.glob(path):47    cap = cv2.imread(file)48    height, width, c = cap.shape49    #result = cv2.imwrite('trialData1.jpg', cap, (width, height))50 51    # Reshape image52    img = cap.copy()53    img = tf.image.resize_with_pad(np.expand_dims(cap, axis=0), 192, 192)54    input_image = tf.cast(img, dtype=tf.float32)55 56    # Setup input and output57    input_details = interpreter.get_input_details()58    output_details = interpreter.get_output_details()59 60    # Make predictions61    interpreter.set_tensor(input_details[0]['index'], np.array(input_image))62    interpreter.invoke()63    keypoints_with_scores = interpreter.get_tensor(output_details[0]['index'])64    #print(keypoints_with_scores)65 66 67    # Rendering68    draw_connections(cap, keypoints_with_scores, EDGES, 0.4)69    draw_keypoints(cap, keypoints_with_scores, 0.4)70    #result.write(cap)71    cv2.imshow('MoveNet Lightning', cap)72 73    #result = cv2.imwrite('trialData1.jpg', cap, (width, height))74 75 76    cv2.destroyAllWindows()77 78    left_hip = keypoints_with_scores[0][0][11]79    left_knee = keypoints_with_scores[0][0][13]80 81    shaped = np.squeeze(np.multiply(interpreter.get_tensor(interpreter.get_output_details()[0]['index']), [480, 640, 1]))82 83    for kp in shaped:84        ky, kx, kp_conf = kp85        #print(int(ky), int(kx), kp_conf)86 87    shaped[0], shaped[1]88 89    for edge, color in EDGES.items():90        p1, p2 = edge91        y1, x1, c1 = shaped[p1]92        y2, x2, c2 = shaped[p2]93        #print((int(x2), int(y2))94        input = np.array([[x1, y1, x2, y2, 0]])95        data = np.concatenate([data, input], axis=0)96 97 98data = np.delete(data, 0, 0)99DF = pd.DataFrame(data)100DF.to_csv("squat_down2.csv")101print(data)102