Skip to content

Instantly share code, notes, and snippets.

@suprnrdy
Forked from smeschke/swap_body_parts.py
Created January 18, 2020 06:25
Show Gist options
  • Save suprnrdy/99a63af891525f27692292c6a1928822 to your computer and use it in GitHub Desktop.
Save suprnrdy/99a63af891525f27692292c6a1928822 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