Skip to content

Instantly share code, notes, and snippets.

@smeschke
Last active April 7, 2023 07:08
Show Gist options
  • Save smeschke/a3dcf7ce7d7d8bf684b26c8ec08d38f0 to your computer and use it in GitHub Desktop.
Save smeschke/a3dcf7ce7d7d8bf684b26c8ec08d38f0 to your computer and use it in GitHub Desktop.
Saving pose data from video using OpenPose
import cv2, numpy as np, csv
#https://github.com/opencv/opencv/blob/master/samples/dnn/openpose.py
outfile_path = '/home/stephen/Desktop/workout.csv'
protoFile = "/home/stephen/pose/mpi/pose_deploy_linevec_faster_4_stages.prototxt"
weightsFile = "/home/stephen/pose/mpi/pose_iter_160000.caffemodel"
net = cv2.dnn.readNetFromCaffe(protoFile, weightsFile)
data, input_width, input_height, threshold, frame_number = [], 368, 386, 0.1, 0
input_source = "/home/stephen/Desktop/workout.MP4"
cap = cv2.VideoCapture(input_source)
# use the previous location of the body part if the model is wrong
previous_x, previous_y = [0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]
while True:
ret, img = cap.read()
if not ret: break
# get the image shape
img_width, img_height = img.shape[1], img.shape[0]
# get a blob from the image
inputBlob = cv2.dnn.blobFromImage(img, 1.0 / 255, (input_width, input_height),(0, 0, 0), swapRB=False, crop=False)
# set the input and perform a forward pass
net.setInput(inputBlob)
output = net.forward()
# get the output shape
output_width, output_height = output.shape[2], output.shape[3]
# Empty list to store the detected keypoints
x_data, y_data = [], []
# Iterate through the body parts
for i in range(15):
# find probability that point is correct
_, prob, _, point = cv2.minMaxLoc(output[0, i, :, :])
# Scale the point to fit on the original image
x, y = (img_width * point[0]) / output_width, (img_height * point[1]) / output_height
# Is the point likely to be correct?
if prob > threshold:
x_data.append(x)
y_data.append(y)
xy = tuple(np.array([x,y], int))
cv2.circle(img, xy, 5, (25,0,255), 5)
# No? us the location in the previous frame
else:
x_data.append(previous_x[i])
y_data.append(previous_y[i])
# add these points to the list of data
data.append(x_data + y_data)
previous_x, previous_y = x_data, y_data
frame_number+=1
# use this break statement to check your data before processing the whole video
#if frame_number == 300: break
print(frame_number)
cv2.imshow('img', img)
k = cv2.waitKey(1)
if k == 27: break
# write the data to a .csv file
import pandas as pd
df = pd.DataFrame(data)
df.to_csv(outfile_path, index = False)
print('save complete')
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment