169 lines
4.0 KiB
Python
169 lines
4.0 KiB
Python
#!/usr/bin/env python
|
|
# coding: utf-8
|
|
|
|
# # Overview
|
|
# The kernel shows how to use the [tf_pose_estimation](https://github.com/ildoonet/tf-pose-estimation) package in Python on a series of running videos.
|
|
|
|
# ## Libraries we need
|
|
# Install tf_pose and pycocotools
|
|
|
|
# In[1]:
|
|
|
|
|
|
import os
|
|
def get_ipython():
|
|
return os
|
|
|
|
get_ipython().system('pip install -qq https://www.github.com/ildoonet/tf-pose-estimation')
|
|
|
|
|
|
# In[2]:
|
|
|
|
|
|
get_ipython().system('pip install -qq pycocotools')
|
|
|
|
|
|
# In[3]:
|
|
|
|
|
|
get_ipython().run_line_magic('load_ext', 'autoreload')
|
|
get_ipython().run_line_magic('autoreload', '2')
|
|
import seaborn as sns
|
|
import matplotlib.pyplot as plt
|
|
plt.rcParams["figure.figsize"] = (8, 8)
|
|
plt.rcParams["figure.dpi"] = 125
|
|
plt.rcParams["font.size"] = 14
|
|
plt.rcParams['font.family'] = ['sans-serif']
|
|
plt.rcParams['font.sans-serif'] = ['DejaVu Sans']
|
|
plt.style.use('ggplot')
|
|
sns.set_style("whitegrid", {'axes.grid': False})
|
|
|
|
|
|
# In[4]:
|
|
|
|
|
|
get_ipython().run_line_magic('matplotlib', 'inline')
|
|
import tf_pose
|
|
import cv2
|
|
from glob import glob
|
|
from tqdm import tqdm_notebook
|
|
from PIL import Image
|
|
import numpy as np
|
|
import os
|
|
def video_gen(in_path):
|
|
c_cap = cv2.VideoCapture(in_path)
|
|
while c_cap.isOpened():
|
|
ret, frame = c_cap.read()
|
|
if not ret:
|
|
break
|
|
yield c_cap.get(cv2.CAP_PROP_POS_MSEC), frame[:, :, ::-1]
|
|
c_cap.release()
|
|
|
|
|
|
# In[5]:
|
|
|
|
|
|
video_paths = glob('../input/*.mp4')
|
|
c_video = video_gen(video_paths[0])
|
|
for _ in range(300):
|
|
c_ts, c_frame = next(c_video)
|
|
plt.imshow(c_frame)
|
|
|
|
|
|
# In[6]:
|
|
|
|
|
|
from tf_pose.estimator import TfPoseEstimator
|
|
from tf_pose.networks import get_graph_path, model_wh
|
|
tfpe = tf_pose.get_estimator()
|
|
|
|
|
|
# In[7]:
|
|
|
|
|
|
humans = tfpe.inference(npimg=c_frame, upsample_size=4.0)
|
|
print(humans)
|
|
|
|
|
|
# In[8]:
|
|
|
|
|
|
new_image = TfPoseEstimator.draw_humans(c_frame[:, :, ::-1], humans, imgcopy=False)
|
|
fig, ax1 = plt.subplots(1, 1, figsize=(10, 10))
|
|
ax1.imshow(new_image[:, :, ::-1])
|
|
|
|
|
|
# In[9]:
|
|
|
|
|
|
body_to_dict = lambda c_fig: {'bp_{}_{}'.format(k, vec_name): vec_val
|
|
for k, part_vec in c_fig.body_parts.items()
|
|
for vec_name, vec_val in zip(['x', 'y', 'score'],
|
|
(part_vec.x, 1-part_vec.y, part_vec.score))}
|
|
c_fig = humans[0]
|
|
body_to_dict(c_fig)
|
|
|
|
|
|
# In[10]:
|
|
|
|
|
|
MAX_FRAMES = 200
|
|
body_pose_list = []
|
|
for vid_path in tqdm_notebook(video_paths, desc='Files'):
|
|
c_video = video_gen(vid_path)
|
|
c_ts, c_frame = next(c_video)
|
|
out_path = '{}_out.avi'.format(os.path.split(vid_path)[1])
|
|
out = cv2.VideoWriter(out_path,
|
|
cv2.VideoWriter_fourcc('M','J','P','G'),
|
|
10,
|
|
(c_frame.shape[1], c_frame.shape[0]))
|
|
for (c_ts, c_frame), _ in zip(c_video,
|
|
tqdm_notebook(range(MAX_FRAMES), desc='Frames')):
|
|
bgr_frame = c_frame[:,:,::-1]
|
|
humans = tfpe.inference(npimg=bgr_frame, upsample_size=4.0)
|
|
for c_body in humans:
|
|
body_pose_list += [dict(video=out_path, time=c_ts, **body_to_dict(c_body))]
|
|
new_image = TfPoseEstimator.draw_humans(bgr_frame, humans, imgcopy=False)
|
|
out.write(new_image)
|
|
out.release()
|
|
|
|
|
|
# In[11]:
|
|
|
|
|
|
import pandas as pd
|
|
body_pose_df = pd.DataFrame(body_pose_list)
|
|
body_pose_df.describe()
|
|
|
|
|
|
# In[12]:
|
|
|
|
|
|
fig, m_axs = plt.subplots(1, 2, figsize=(15, 5))
|
|
for c_ax, (c_name, c_rows) in zip(m_axs, body_pose_df.groupby('video')):
|
|
for i in range(17):
|
|
c_ax.plot(c_rows['time'], c_rows['bp_{}_y'.format(i)], label='x {}'.format(i))
|
|
c_ax.legend()
|
|
c_ax.set_title(c_name)
|
|
|
|
|
|
# In[13]:
|
|
|
|
|
|
fig, m_axs = plt.subplots(1, 2, figsize=(15, 5))
|
|
for c_ax, (c_name, n_rows) in zip(m_axs, body_pose_df.groupby('video')):
|
|
for i in range(17):
|
|
c_rows = n_rows.query('bp_{}_score>0.6'.format(i)) # only keep confident results
|
|
c_ax.plot(c_rows['bp_{}_x'.format(i)], c_rows['bp_{}_y'.format(i)], label='BP {}'.format(i))
|
|
c_ax.legend()
|
|
c_ax.set_title(c_name)
|
|
|
|
|
|
# In[14]:
|
|
|
|
|
|
body_pose_df.to_csv('body_pose.csv', index=False)
|
|
|
|
|
|
# In[15]:
|