Last active
April 7, 2023 07:08
-
-
Save smeschke/a3dcf7ce7d7d8bf684b26c8ec08d38f0 to your computer and use it in GitHub Desktop.
Saving pose data from video using OpenPose
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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