chrisvlds/MLProject

1import math2import pickle3import tensorflow as tf4import numpy as np5import gradio as gr6import cv27 8print("opencv v is:" + cv2.__version__)9 10interpreter = tf.lite.Interpreter(model_path='lite-model_movenet_singlepose_lightning_3.tflite')11interpreter.allocate_tensors()12loaded_model = pickle.load(open('model1.pkl', 'rb'))13 14 15def draw_keypoints(frame, keypoints, confidence_threshold):16    y, x, c = frame.shape17    shaped = np.squeeze(np.multiply(keypoints, [y, x, 1]))18 19    for kp in shaped:20        ky, kx, kp_conf = kp21        if kp_conf > confidence_threshold:22            cv2.circle(frame, (int(kx), int(ky)), 4, (0, 255, 0), -1)23 24 25EDGES = {26    (11, 12): 'y',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 44def prediction(video):45    repCounterUp = 046    repCounterDown = 047    repCounter = 048    state = 149    cap = cv2.VideoCapture(video)50    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))51    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))52    result = cv2.VideoWriter('outputt.webm', cv2.VideoWriter_fourcc(*'VP90'), 20, (width, height))53    while 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 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        else:100            print("Up")101 102        print("Predicted=%s" % result)103        if prediction == 0:104            print("Down")105            if repCounterDown < 10:106                repCounterDown += 1107            if repCounterUp > 0:108                repCounterUp -= 1109        else:110            print("Up")111            if repCounterUp < 10:112                repCounterUp += 1113            if repCounterDown > 0:114                repCounterDown -= 1115 116        if repCounterDown == 10 and repCounterUp == 0:117            if state == 1:118                state = 0119                repCounter += 1120        elif repCounterUp == 10 and repCounterDown == 0:121            if state == 0:122                state = 1123                repCounter += 1124 125        reps = math.floor(repCounter / 2)126        frame = cv2.rectangle(frame, (150, 0), (900, 200), (0, 0, 0), -1)127        frame = cv2.putText(frame, 'Reps: ' + str(reps), (150, 150), cv2.FONT_HERSHEY_SIMPLEX, 5, (255, 0, 0), 5,128                            cv2.LINE_AA)129        result.write(frame)130 131        if cv2.waitKey(10) & 0xFF == ord('q'):132            break133 134    cap.release()135    result.release()136    cv2.destroyAllWindows()137 138    # left_hip = keypoints_with_scores[0][0][11]139    # left_knee = keypoints_with_scores[0][0][13]140    #141    # shaped = np.squeeze(142    #     np.multiply(interpreter.get_tensor(interpreter.get_output_details()[0]['index']), [480, 640, 1]))143 144    # for kp in shaped:145    #     ky, kx, kp_conf = kp146    #     #print(int(ky), int(kx), kp_conf)147    #148    # shaped[0], shaped[1]149    #150    # for edge, color in EDGES.items():151    #     p1, p2 = edge152    #     y1, x1, c1 = shaped[p1]153    #     y2, x2, c2 = shaped[p2]154    #     #print((int(x2), int(y2)))155    #     input = np.array([[x1, y1, x2, y2]])156    reps = math.floor(repCounter / 2)157    print("Number of reps: " + str(reps))158    return 'outputt.webm', reps159 160 161app = gr.Interface(fn=prediction,162                   inputs=gr.Video(format='webm', label="Video"),163                   outputs=[gr.Video(label="Detection Video", format='webm'), 164                   gr.Text(label="Number of Reps:")],165                   examples=["sample_video.webm"])166 167app.queue().launch()168