text_html_output_1.png application_vnd.jupyter.stderr_output_1.png
from IPython.display import HTML, display

from tensorflow_docs.vis import embed

import cv2

import imageio

import numpy as np

import tensorflow as tf

import tensorflow_hub as hub

KEYPOINT_EDGE_INDS_TO_COLOR = {(0, 1): 'm', (0, 2): 'c', (1, 3): 'm', (2, 4): 'c', (0, 5): 'm', (0, 6): 'c', (5, 7): 'm', (7, 9): 'm', (6, 8): 'c', (8, 10): 'c', (5, 6): 'y', (5, 11): 'm', (6, 12): 'c', (11, 12): 'y', (11, 13): 'm', (13, 15): 'm', (12, 14): 'c', (14, 16): 'c'}

model = hub.load('https://tfhub.dev/google/movenet/multipose/lightning/1')
movenet = model.signatures['serving_default']

def loop(frame, keypoints, threshold=0.11):
    pass

def draw_keypoints(frame, keypoints, threshold=0.11):
    width, height, _ = frame.shape
    shaped = np.squeeze(np.multiply(keypoints, [width, height, 1]))
    for kp in shaped:
        ky, kx, kp_conf = kp
        if kp_conf > threshold:
            cv2.circle(frame, (int(kx), int(ky)), 4, (255, 0, 0), -1)

def draw_edges(frame, keypoints, edges, threshold=0.11):
    y, x, c = frame.shape
    shaped = np.squeeze(np.multiply(keypoints, [y, x, 1]))
    for edge, color in edges.items():
        p1, p2 = edge
        y1, x1, c1 = shaped[p1]
        y2, x2, c2 = shaped[p2]
        if (c1 > threshold) & (c2 > threshold):
            cv2.line(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 4)

def progress(value, max=100):
    return HTML("\n      <progress\n          value='{value}'\n          max='{max}',\n          style='width: 100%'\n      >\n          {value}\n      </progress>\n  ".format(value=value, max=max))

gif = cv2.VideoCapture('./ngannou.gif')
frame_count = int(gif.get(cv2.CAP_PROP_FRAME_COUNT))

output_frames = []

def run_inference():
    
    #Set the progress bar to 0. It ranges from the first to the last frame
    bar = display(progress(0, frame_count-1), display_id=True)
    
    while gif.isOpened():
        
        #Capture the frame
        ret, frame = gif.read()
        
        #Process the frame : resize to the input size
        if frame is None: 
            break
        
        #Retrieve the frame index
        index = gif.get(cv2.CAP_PROP_POS_FRAMES)
        
        image = frame.copy()
        image = tf.cast(tf.image.resize_with_pad(image, 256, 256), dtype=tf.int32)
        input_image = tf.expand_dims(image, axis=0)

        #Perform inference
        results = movenet(input_image)
        keypoints = results['output_0'].numpy()[:,:,:51].reshape((6,17,3))

        #Loop through the results
        loop(frame, keypoints, threshold=0.11)
        
        #Add the drawings to the output frames
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) #OpenCV processes BGR images instead of RGB
        output_frames.append(frame_rgb)
        
        #Update the progress bar
        bar.update(progress(index, frame_count-1))
    
    #Release the object
    gif.release()

output = np.stack(output_frames, axis=0)
imageio.mimsave('./animation.gif', output, fps=10)
embed.embed_file('./animation.gif')