Skip to content

Instantly share code, notes, and snippets.

@smeschke
Created September 21, 2018 19:49
Show Gist options
  • Save smeschke/12936dc6cd58024eb424f380461e8baa to your computer and use it in GitHub Desktop.
Save smeschke/12936dc6cd58024eb424f380461e8baa to your computer and use it in GitHub Desktop.
Ensure that body parts are always in the right places
import pandas as pd
import numpy as np
import cv2, os
import csv
input_source = "/home/stephen/Desktop/me.MP4"
cap = cv2.VideoCapture(input_source)
frame_number = 0
font, scale, colorText, thick = cv2.FONT_HERSHEY_SIMPLEX, .5, (234,234,234), 1
size, color, thickness = 5, (255,255,255), 5
#get pose data - data is generated by open pose video
df = pd.read_csv('/home/stephen/Desktop/open_pose_data.csv')
# there are 15 points in the skeleton
# 0 head
# 1 neck
# 2, 5 shoulders
# 3, 6 elbows
# 4, 7 hands
# 8, 11 hips
# 9, 12 knees
# 10, 13 ankles
# 14 torso
data = []
while cv2.waitKey(10) < 0 and frame_number<len(df.values)-2:
ret, img = cap.read()
if not ret: break
try: values = df.values[frame_number]
except: break
values = np.array(values, int)
points = []
points.append((values[0], values[1]))
points.append((values[2], values[3]))
points.append((values[4], values[5]))
points.append((values[6], values[7]))
points.append((values[8], values[9]))
points.append((values[10], values[11]))
points.append((values[12], values[13]))
points.append((values[14], values[15]))
points.append((values[16], values[17]))
points.append((values[18], values[19]))
points.append((values[20], values[21]))
points.append((values[22], values[23]))
points.append((values[24], values[25]))
points.append((values[26], values[27]))
points.append((values[28], values[29]))
#create a blank list to store the non-swapped poitns
non_swap_points = []
for i in range(15): non_swap_points.append((0,0))
#add the head, that point never changes
non_swap_points[0] = points[0]
#add the neck, that point never changes
non_swap_points[1] = points[1]
#add the torso, that never changes
non_swap_points[14] = points[14]
#swap the left and right shoulders (2 and 5)
if points[2][0]<points[5][0]:
non_swap_points[2] = points[2]
non_swap_points[5] = points[5]
else:
non_swap_points[2] = points[5]
non_swap_points[5] = points[2]
# swap the elbows
if points[3][0]<points[6][0]:
non_swap_points[3] = points[3]
non_swap_points[6] = points[6]
else:
non_swap_points[6] = points[3]
non_swap_points[3] = points[6]
#swap the hands
if points[4][0]<points[7][0]:
non_swap_points[4] = points[4]
non_swap_points[7] = points[7]
else:
non_swap_points[7] = points[4]
non_swap_points[4] = points[7]
#swap the hips
if points[8][0]<points[11][0]:
non_swap_points[11] = points[11]
non_swap_points[8] = points[8]
else:
non_swap_points[8] = points[11]
non_swap_points[11] = points[8]
#swap the knees
if points[9][0]<points[12][0]:
non_swap_points[9] = points[9]
non_swap_points[12] = points[12]
else:
non_swap_points[12] = points[9]
non_swap_points[9] = points[12]
#swap the feet
if points[10][0]<points[13][0]:
non_swap_points[10] = points[10]
non_swap_points[13] = points[13]
else:
non_swap_points[13] = points[10]
non_swap_points[10] = points[13]
for point in non_swap_points:
cv2.circle(img, point, 3, (0, 0, 255), 3)
cv2.putText(img, str(non_swap_points.index(point)), point, font, scale, colorText, thick, cv2.LINE_AA)
cv2.imshow('Output-Skeleton', img)
frame_number+=1
data.append(non_swap_points)
cv2.destroyAllWindows()
with open('/home/stephen/Desktop/swapped_body_parts.csv', 'w') as csvfile:
fieldnames = []
for i in range(30): fieldnames.append(str(i))
writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
writer.writeheader()
for trick in data:
writer.writerow({'0': trick[0][0],
'1':trick[0][1],
'2':trick[1][0],
'3':trick[1][1],
'4':trick[2][0],
'5':trick[2][1],
'6':trick[3][0],
'7':trick[3][1],
'8':trick[4][0],
'9':trick[4][1],
'10':trick[5][0],
'11':trick[5][1],
'12':trick[6][0],
'13':trick[6][1],
'14':trick[7][0],
'15':trick[7][1],
'16':trick[8][0],
'17':trick[8][1],
'18':trick[9][0],
'19':trick[9][1],
'20':trick[10][0],
'21':trick[10][1],
'22':trick[11][0],
'23':trick[11][1],
'24':trick[12][0],
'25':trick[12][1],
'26':trick[13][0],
'27':trick[13][1],
'28':trick[14][0],
'29':trick[14][1]})
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment