Skip to content

Instantly share code, notes, and snippets.

@suprnrdy
Forked from smeschke/smooth_pose_data.py
Created January 18, 2020 06:25
Show Gist options
  • Save suprnrdy/a9580bf0e567f825a4c20cbd1e17993d to your computer and use it in GitHub Desktop.
Save suprnrdy/a9580bf0e567f825a4c20cbd1e17993d to your computer and use it in GitHub Desktop.
Smooth Pose Estimation Data
import pandas as pd
import numpy as np
import cv2, os
import scipy
from scipy import signal
import csv
circle_color, line_color = (255,255,0), (0,0,255)
window_length, polyorder = 13, 2
sd = "workout"
input_source = "/home/stephen/Desktop/" + sd + '.MP4'
# Get pose data - data is generated by OpenPose
df = pd.read_csv('/home/stephen/Desktop/' +sd+ '.csv')
cap = cv2.VideoCapture(input_source)
hw = 720
out = cv2.VideoWriter('/home/stephen/Desktop/smooth_pose.avi',
cv2.VideoWriter_fourcc('M','J','P','G'), 30, (hw,hw))
# There are 15 points in the skeleton
pairs = [[0,1], # head
[1,2],[1,5], # sholders
[2,3],[3,4],[5,6],[6,7], # arms
[1,14],[14,11],[14,8], # hips
[8,9],[9,10],[11,12],[12,13]] # legs
# Smooth it out
for i in range(30): df[str(i)] = signal.savgol_filter(df[str(i)], window_length, polyorder)
frame_number = 0
while True:
print(frame_number)
ret, img = cap.read()
if not ret: break
#img = np.zeros_like(img)
values = np.array(df.values[frame_number], int)
points, lateral_offset = [], 18
points = list(zip(values[:15]+lateral_offset, values[15:]))
cc = 0
for point in points:
cc += 90
xy = tuple(np.array([point[0], point[1]], int))
cv2.circle(img, xy, 5, (cc,cc,cc), 5)
# Draw Skeleton
for pair in pairs:
partA = pair[0]
partB = pair[1]
cv2.line(img, points[partA], points[partB], line_color, 3, lineType=cv2.LINE_AA)
cv2.imshow('Output-Skeleton', img)
k = cv2.waitKey(100)
if k == 27: break
out.write(img)
frame_number+=1
cv2.destroyAllWindows()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment