Initial commit
This commit is contained in:
parent
46d36d9bd4
commit
70c5be6697
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
|
|
@ -0,0 +1,56 @@
|
|||
# import pandas as pd
|
||||
# from copy import deepcopy
|
||||
import pandas as pd
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
# import torch.nn.functional as F
|
||||
# import time
|
||||
|
||||
|
||||
class LSTMModel(nn.Module):
|
||||
def __init__(self, input_dim=5, h_RNN_layers=2, h_RNN=256, drop_p=0.2, num_classes=1):
|
||||
super(LSTMModel, self).__init__()
|
||||
self.input_dim = input_dim
|
||||
self.h_RNN_layers = h_RNN_layers # RNN hidden layers
|
||||
self.h_RNN = h_RNN # RNN hidden nodes
|
||||
self.drop_p = drop_p
|
||||
if h_RNN_layers < 2:
|
||||
drop_p = 0
|
||||
self.num_classes = num_classes
|
||||
self.LSTM = nn.LSTM(
|
||||
input_size=self.input_dim,
|
||||
hidden_size=self.h_RNN,
|
||||
num_layers=h_RNN_layers,
|
||||
dropout=drop_p,
|
||||
batch_first=True, # input & output will has batch size as 1s dimension. e.g. (batch, time_step, input_size)
|
||||
)
|
||||
self.fc1 = nn.Linear(self.h_RNN, self.num_classes)
|
||||
|
||||
def forward(self, x, h_s=None):
|
||||
# print('forward started')
|
||||
self.LSTM.flatten_parameters()
|
||||
RNN_out, h_s = self.LSTM(x, h_s)
|
||||
""" h_n shape (n_layers, batch, hidden_size), h_c shape (n_layers, batch, hidden_size) """
|
||||
""" None represents zero initial hidden state. RNN_out has shape=(batch, time_step, output_size) """
|
||||
|
||||
# FC layers
|
||||
out = self.fc1(RNN_out[:, -1, :]) # choose RNN_out at the last time step
|
||||
return out, h_s
|
||||
|
||||
# model = LSTMModel(h_RNN=16, h_RNN_layers=2, drop_p=0.2, num_classes=7)
|
||||
# model.load_state_dict(torch.load('lstm.sav'))
|
||||
# model.eval()
|
||||
# df = pd.read_csv('dataset/2sec_multi_train_data.csv', header=None)
|
||||
# sum = 0
|
||||
# h_s = None
|
||||
# for j in range(0, 80):
|
||||
# xdata = df.iloc[j, :180].values.reshape((36, 5), order='F')
|
||||
# # print(xdata)
|
||||
# #
|
||||
# for i in range(1):
|
||||
# xcurr = torch.Tensor(xdata.reshape(-1, 36, 5))
|
||||
# outputs, h_s = model(xcurr, h_s)
|
||||
# _, predicted = torch.max(outputs.data, 1)
|
||||
# sum += (predicted.cpu().numpy()[0] == 0)
|
||||
#
|
||||
# print(sum)
|
||||
|
|
@ -0,0 +1,347 @@
|
|||
from .visual import CocoPart
|
||||
import numpy as np
|
||||
from helpers import *
|
||||
from default_params import *
|
||||
|
||||
|
||||
def match_ip(ip_set, new_ips, lstm_set, num_matched, consecutive_frames=DEFAULT_CONSEC_FRAMES):
|
||||
len_ip_set = len(ip_set)
|
||||
added = [False for _ in range(len_ip_set)]
|
||||
|
||||
new_len_ip_set = len_ip_set
|
||||
for new_ip in new_ips:
|
||||
if not is_valid(new_ip):
|
||||
continue
|
||||
# assert valid_candidate_hist(new_ip)
|
||||
cmin = [MIN_THRESH, -1]
|
||||
for i in range(len_ip_set):
|
||||
if not added[i] and dist(last_ip(ip_set[i])[0], new_ip) < cmin[0]:
|
||||
# here add dome condition that last_ip(ip_set[0] >-5 or someting)
|
||||
cmin[0] = dist(last_ip(ip_set[i])[0], new_ip)
|
||||
cmin[1] = i
|
||||
|
||||
if cmin[1] == -1:
|
||||
ip_set.append([None for _ in range(consecutive_frames - 1)] + [new_ip])
|
||||
lstm_set.append([None, 0, 0, 0]) # Initial hidden state of lstm is None
|
||||
new_len_ip_set += 1
|
||||
|
||||
else:
|
||||
added[cmin[1]] = True
|
||||
pop_and_add(ip_set[cmin[1]], new_ip, consecutive_frames)
|
||||
|
||||
new_matched = num_matched
|
||||
|
||||
removed_indx = []
|
||||
removed_match = []
|
||||
|
||||
for i in range(len(added)):
|
||||
if not added[i]:
|
||||
pop_and_add(ip_set[i], None, consecutive_frames)
|
||||
if ip_set[i] == [None for _ in range(consecutive_frames)]:
|
||||
if i < num_matched:
|
||||
new_matched -= 1
|
||||
removed_match.append(i)
|
||||
|
||||
new_len_ip_set -= 1
|
||||
removed_indx.append(i)
|
||||
|
||||
for i in sorted(removed_indx, reverse=True):
|
||||
ip_set.pop(i)
|
||||
lstm_set.pop()
|
||||
|
||||
return new_matched, new_len_ip_set, removed_match
|
||||
|
||||
|
||||
def extend_vector(p1, p2, l):
|
||||
p1 += (p1-p2)*l/(2*np.linalg.norm((p1-p2), 2))
|
||||
p2 -= (p1-p2)*l/(2*np.linalg.norm((p1-p2), 2))
|
||||
return p1, p2
|
||||
|
||||
|
||||
def perp(a):
|
||||
b = np.empty_like(a)
|
||||
b[0] = -a[1]
|
||||
b[1] = a[0]
|
||||
return b
|
||||
|
||||
# line segment a given by endpoints a1, a2
|
||||
# line segment b given by endpoints b1, b2
|
||||
# return
|
||||
|
||||
|
||||
def seg_intersect(a1, a2, b1, b2):
|
||||
da = a2-a1
|
||||
db = b2-b1
|
||||
dp = a1-b1
|
||||
dap = perp(da)
|
||||
denom = np.dot(dap, db)
|
||||
num = np.dot(dap, dp)
|
||||
return (num / denom.astype(float))*db + b1
|
||||
|
||||
|
||||
def get_kp(kp):
|
||||
threshold1 = 5e-3
|
||||
|
||||
# dict of np arrays of coordinates
|
||||
inv_pend = {}
|
||||
# print(type(kp[CocoPart.LEar]))
|
||||
numx = (kp[CocoPart.LEar][2]*kp[CocoPart.LEar][0] + kp[CocoPart.LEye][2]*kp[CocoPart.LEye][0] +
|
||||
kp[CocoPart.REye][2]*kp[CocoPart.REye][0] + kp[CocoPart.REar][2]*kp[CocoPart.REar][0])
|
||||
numy = (kp[CocoPart.LEar][2]*kp[CocoPart.LEar][1] + kp[CocoPart.LEye][2]*kp[CocoPart.LEye][1] +
|
||||
kp[CocoPart.REye][2]*kp[CocoPart.REye][1] + kp[CocoPart.REar][2]*kp[CocoPart.REar][1])
|
||||
den = kp[CocoPart.LEar][2] + kp[CocoPart.LEye][2] + kp[CocoPart.REye][2] + kp[CocoPart.REar][2]
|
||||
|
||||
if den < HEAD_THRESHOLD:
|
||||
inv_pend['H'] = None
|
||||
else:
|
||||
inv_pend['H'] = np.array([numx/den, numy/den])
|
||||
|
||||
if all([kp[CocoPart.LShoulder], kp[CocoPart.RShoulder],
|
||||
kp[CocoPart.LShoulder][2] > threshold1, kp[CocoPart.RShoulder][2] > threshold1]):
|
||||
inv_pend['N'] = np.array([(kp[CocoPart.LShoulder][0]+kp[CocoPart.RShoulder][0])/2,
|
||||
(kp[CocoPart.LShoulder][1]+kp[CocoPart.RShoulder][1])/2])
|
||||
else:
|
||||
inv_pend['N'] = None
|
||||
|
||||
if all([kp[CocoPart.LHip], kp[CocoPart.RHip],
|
||||
kp[CocoPart.LHip][2] > threshold1, kp[CocoPart.RHip][2] > threshold1]):
|
||||
inv_pend['B'] = np.array([(kp[CocoPart.LHip][0]+kp[CocoPart.RHip][0])/2,
|
||||
(kp[CocoPart.LHip][1]+kp[CocoPart.RHip][1])/2])
|
||||
else:
|
||||
inv_pend['B'] = None
|
||||
|
||||
if kp[CocoPart.LKnee] is not None and kp[CocoPart.LKnee][2] > threshold1:
|
||||
inv_pend['KL'] = np.array([kp[CocoPart.LKnee][0], kp[CocoPart.LKnee][1]])
|
||||
else:
|
||||
inv_pend['KL'] = None
|
||||
|
||||
if kp[CocoPart.RKnee] is not None and kp[CocoPart.RKnee][2] > threshold1:
|
||||
inv_pend['KR'] = np.array([kp[CocoPart.RKnee][0], kp[CocoPart.RKnee][1]])
|
||||
else:
|
||||
inv_pend['KR'] = None
|
||||
|
||||
if inv_pend['B'] is not None:
|
||||
if inv_pend['N'] is not None:
|
||||
height = np.linalg.norm(inv_pend['N'] - inv_pend['B'], 2)
|
||||
LS, RS = extend_vector(np.asarray(kp[CocoPart.LShoulder][:2]),
|
||||
np.asarray(kp[CocoPart.RShoulder][:2]), height/4)
|
||||
LB, RB = extend_vector(np.asarray(kp[CocoPart.LHip][:2]),
|
||||
np.asarray(kp[CocoPart.RHip][:2]), height/3)
|
||||
ubbox = (LS, RS, RB, LB)
|
||||
|
||||
if inv_pend['KL'] is not None and inv_pend['KR'] is not None:
|
||||
lbbox = (LB, RB, inv_pend['KR'], inv_pend['KL'])
|
||||
else:
|
||||
lbbox = ([0, 0], [0, 0])
|
||||
#lbbox = None
|
||||
else:
|
||||
ubbox = ([0, 0], [0, 0])
|
||||
#ubbox = None
|
||||
if inv_pend['KL'] is not None and inv_pend['KR'] is not None:
|
||||
lbbox = (np.array(kp[CocoPart.LHip][:2]), np.array(kp[CocoPart.RHip][:2]),
|
||||
inv_pend['KR'], inv_pend['KL'])
|
||||
else:
|
||||
lbbox = ([0, 0], [0, 0])
|
||||
#lbbox = None
|
||||
else:
|
||||
ubbox = ([0, 0], [0, 0])
|
||||
lbbox = ([0, 0], [0, 0])
|
||||
#ubbox = None
|
||||
#lbbox = None
|
||||
# condition = (inv_pend["H"] is None) and (inv_pend['N'] is not None and inv_pend['B'] is not None)
|
||||
# if condition:
|
||||
# print("half disp")
|
||||
|
||||
return inv_pend, ubbox, lbbox
|
||||
|
||||
|
||||
def get_angle(v0, v1):
|
||||
return np.math.atan2(np.linalg.det([v0, v1]), np.dot(v0, v1))
|
||||
|
||||
|
||||
def is_valid(ip):
|
||||
|
||||
assert ip is not None
|
||||
|
||||
ip = ip["keypoints"]
|
||||
return (ip['B'] is not None and ip['N'] is not None and ip['H'] is not None)
|
||||
|
||||
|
||||
def get_rot_energy(ip0, ip1):
|
||||
t = ip1["time"] - ip0["time"]
|
||||
ip0 = ip0["keypoints"]
|
||||
ip1 = ip1["keypoints"]
|
||||
m1 = 1
|
||||
m2 = 5
|
||||
m3 = 5
|
||||
energy = 0
|
||||
den = 0
|
||||
N1 = ip1['N'] - ip1['B']
|
||||
N0 = ip0['N'] - ip0['B']
|
||||
d2sq = N1.dot(N1)
|
||||
w2sq = (get_angle(N0, N1)/t)**2
|
||||
energy += m2*d2sq*w2sq
|
||||
|
||||
den += m2*d2sq
|
||||
H1 = ip1['H'] - ip1['B']
|
||||
H0 = ip0['H'] - ip0['B']
|
||||
d1sq = H1.dot(H1)
|
||||
w1sq = (get_angle(H0, H1)/t)**2
|
||||
energy += m1*d1sq*w1sq
|
||||
den += m1*d1sq
|
||||
|
||||
energy = energy/(2*den)
|
||||
# energy = energy/2
|
||||
return energy
|
||||
|
||||
|
||||
def get_angle_vertical(v):
|
||||
return np.math.atan2(-v[0], -v[1])
|
||||
|
||||
|
||||
def get_gf(ip0, ip1, ip2):
|
||||
t1 = ip1["time"] - ip0["time"]
|
||||
t2 = ip2["time"] - ip1["time"]
|
||||
ip0 = ip0["keypoints"]
|
||||
ip1 = ip1["keypoints"]
|
||||
ip2 = ip2["keypoints"]
|
||||
|
||||
m1 = 1
|
||||
m2 = 15
|
||||
g = 10
|
||||
H2 = ip2['H'] - ip2['N']
|
||||
H1 = ip1['H'] - ip1['N']
|
||||
H0 = ip0['H'] - ip0['N']
|
||||
d1 = np.sqrt(H1.dot(H1))
|
||||
theta_1_plus_2_2 = get_angle_vertical(H2)
|
||||
theta_1_plus_2_1 = get_angle_vertical(H1)
|
||||
theta_1_plus_2_0 = get_angle_vertical(H0)
|
||||
# print("H: ",H0,H1,H2)
|
||||
N2 = ip2['N'] - ip2['B']
|
||||
N1 = ip1['N'] - ip1['B']
|
||||
N0 = ip0['N'] - ip0['B']
|
||||
d2 = np.sqrt(N1.dot(N1))
|
||||
# print("N: ",N0,N1,N2)
|
||||
theta_2_2 = get_angle_vertical(N2)
|
||||
theta_2_1 = get_angle_vertical(N1)
|
||||
theta_2_0 = get_angle_vertical(N0)
|
||||
#print("theta_2_2:",theta_2_2,"theta_2_1:",theta_2_1,"theta_2_0:",theta_2_0,sep=", ")
|
||||
theta_1_0 = theta_1_plus_2_0 - theta_2_0
|
||||
theta_1_1 = theta_1_plus_2_1 - theta_2_1
|
||||
theta_1_2 = theta_1_plus_2_2 - theta_2_2
|
||||
|
||||
# print("theta1: ",theta_1_0,theta_1_1,theta_1_2)
|
||||
# print("theta2: ",theta_2_0,theta_2_1,theta_2_2)
|
||||
|
||||
theta2 = theta_2_1
|
||||
theta1 = theta_1_1
|
||||
|
||||
del_theta1_0 = (get_angle(H0, H1))/t1
|
||||
del_theta1_1 = (get_angle(H1, H2))/t2
|
||||
|
||||
del_theta2_0 = (get_angle(N0, N1))/t1
|
||||
del_theta2_1 = (get_angle(N1, N2))/t2
|
||||
# print("del_theta2_1:",del_theta2_1,"del_theta2_0:",del_theta2_0,sep=",")
|
||||
del_theta1 = 0.5 * (del_theta1_1 + del_theta1_0)
|
||||
del_theta2 = 0.5 * (del_theta2_1 + del_theta2_0)
|
||||
|
||||
doubledel_theta1 = (del_theta1_1 - del_theta1_0) / 0.5*(t1 + t2)
|
||||
doubledel_theta2 = (del_theta2_1 - del_theta2_0) / 0.5*(t1 + t2)
|
||||
# print("doubledel_theta2:",doubledel_theta2)
|
||||
|
||||
d1 = d1/d2
|
||||
d2 = 1
|
||||
# print("del_theta",del_theta1,del_theta2)
|
||||
# print("doubledel_theta",doubledel_theta1,doubledel_theta2)
|
||||
|
||||
Q_RD1 = 0
|
||||
Q_RD1 += m1 * d1 * doubledel_theta1 * doubledel_theta1
|
||||
Q_RD1 += (m1*d1*d1 + m1*d1*d2*np.cos(theta1))*doubledel_theta2
|
||||
Q_RD1 += m1*d1*d2*np.sin(theta1)*del_theta2*del_theta2
|
||||
Q_RD1 -= m1*g*d2*np.sin(theta1+theta2)
|
||||
|
||||
Q_RD2 = 0
|
||||
Q_RD2 += (m1*d1*d1 + m1*d1*d2*np.cos(theta1))*doubledel_theta1
|
||||
Q_RD2 += ((m1+m2)*d2*d2 + m1*d1*d1 + 2*m1*d1*d2*np.cos(theta1))*doubledel_theta2
|
||||
Q_RD2 -= 2*m1*d1*d2*np.sin(theta1)*del_theta2*del_theta1 + m1*d1 * \
|
||||
d2*np.sin(theta1)*del_theta1*del_theta1
|
||||
Q_RD2 -= (m1 + m2)*g*d2*np.sin(theta2) + m1*g*d1*np.sin(theta1 + theta2)
|
||||
|
||||
# print("Energy: ", Q_RD1 + Q_RD2)
|
||||
return Q_RD1 + Q_RD2
|
||||
|
||||
|
||||
def get_height_bbox(ip):
|
||||
bbox = ip["box"]
|
||||
assert(type(bbox == np.ndarray))
|
||||
diff_box = bbox[1] - bbox[0]
|
||||
return diff_box[1]
|
||||
|
||||
|
||||
def get_ratio_bbox(ip):
|
||||
bbox = ip["box"]
|
||||
assert(type(bbox == np.ndarray))
|
||||
diff_box = bbox[1] - bbox[0]
|
||||
if diff_box[1] == 0:
|
||||
diff_box[1] += 1e5*diff_box[0]
|
||||
assert(np.any(diff_box > 0))
|
||||
ratio = diff_box[0]/diff_box[1]
|
||||
return ratio
|
||||
|
||||
|
||||
def get_ratio_derivative(ip0, ip1):
|
||||
ratio_der = None
|
||||
time = ip1["time"] - ip0["time"]
|
||||
diff_box = ip1["features"]["ratio_bbox"] - ip0["features"]["ratio_bbox"]
|
||||
assert time != 0
|
||||
ratio_der = diff_box/time
|
||||
|
||||
return ratio_der
|
||||
|
||||
|
||||
def match_ip2(matched_ip_set, unmatched_ip_set, new_ips, re_matrix, gf_matrix, consecutive_frames=DEFAULT_CONSEC_FRAMES):
|
||||
len_matched_ip_set = len(matched_ip_set)
|
||||
added_matched = [False for _ in range(len_matched_ip_set)]
|
||||
len_unmatched_ip_set = len(unmatched_ip_set)
|
||||
added_unmatched = [False for _ in range(len_unmatched_ip_set)]
|
||||
for new_ip in new_ips:
|
||||
if not is_valid(new_ip):
|
||||
continue
|
||||
cmin = [MIN_THRESH, -1]
|
||||
connected_set = None
|
||||
connected_added = None
|
||||
for i in range(len_matched_ip_set):
|
||||
if not added_matched[i] and dist(last_ip(matched_ip_set[i])[0], new_ip) < cmin[0]:
|
||||
# here add dome condition that last_ip(ip_set[0] >-5 or someting)
|
||||
cmin[0] = dist(last_ip(matched_ip_set[i])[0], new_ip)
|
||||
cmin[1] = i
|
||||
connected_set = matched_ip_set
|
||||
connected_added = added_matched
|
||||
for i in range(len_unmatched_ip_set):
|
||||
if not added_unmatched[i] and dist(last_ip(unmatched_ip_set[i])[0], new_ip) < cmin[0]:
|
||||
# here add dome condition that last_ip(ip_set[0] >-5 or someting)
|
||||
cmin[0] = dist(last_ip(unmatched_ip_set[i])[0], new_ip)
|
||||
cmin[1] = i
|
||||
connected_set = unmatched_ip_set
|
||||
connected_added = added_unmatched
|
||||
|
||||
if cmin[1] == -1:
|
||||
unmatched_ip_set.append([None for _ in range(consecutive_frames - 1)] + [new_ip])
|
||||
# re_matrix.append([])
|
||||
# gf_matrix.append([])
|
||||
|
||||
else:
|
||||
connected_added[cmin[1]] = True
|
||||
pop_and_add(connected_set[cmin[1]], new_ip, consecutive_frames)
|
||||
|
||||
i = 0
|
||||
while i < len(added_matched):
|
||||
if not added_matched[i]:
|
||||
pop_and_add(matched_ip_set[i], None, consecutive_frames)
|
||||
if matched_ip_set[i] == [None for _ in range(consecutive_frames)]:
|
||||
matched_ip_set.pop(i)
|
||||
# re_matrix.pop(i)
|
||||
# gf_matrix.pop(i)
|
||||
added_matched.pop(i)
|
||||
continue
|
||||
i += 1
|
||||
|
|
@ -0,0 +1,112 @@
|
|||
import base64
|
||||
import io
|
||||
import time
|
||||
|
||||
import openpifpaf
|
||||
import PIL
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
|
||||
class Processor(object):
|
||||
def __init__(self, width_height, args):
|
||||
self.width_height = width_height
|
||||
|
||||
'''
|
||||
# Load model
|
||||
self.model_cpu, _ = openpifpaf.network.factory()
|
||||
self.model = self.model_cpu.to(args.device)
|
||||
self.processor = openpifpaf.decoder.factory(self.model_cpu.head_metas)
|
||||
# print(self.processor.device)
|
||||
self.device = args.device
|
||||
'''
|
||||
# Load model
|
||||
print("PROCESSOR.PY, Call openpifpaf.NETWORK.factory_from_args()")
|
||||
print("return net_cpu, epoch")
|
||||
self.model, _ = openpifpaf.network.factory_from_args(args)
|
||||
|
||||
print("PROCESSOR.PY, Call model.to(args.device)")
|
||||
self.model = self.model.to(args.device)
|
||||
|
||||
print("\nPROCESSOR.PY, Call openpifpaf.DECODER.factory_from_args()")
|
||||
self.processor = openpifpaf.decoder.factory_from_args(args, self.model)
|
||||
|
||||
# print(self.processor.device)
|
||||
self.device = args.device
|
||||
|
||||
def get_bb(self, kp_set, score=None):
|
||||
bb_list = []
|
||||
for i in range(kp_set.shape[0]):
|
||||
x = kp_set[i, :15, 0]
|
||||
y = kp_set[i, :15, 1]
|
||||
v = kp_set[i, :15, 2]
|
||||
assert np.any(v > 0)
|
||||
if not np.any(v > 0):
|
||||
return None
|
||||
|
||||
# keypoint bounding box
|
||||
x1, x2 = np.min(x[v > 0]), np.max(x[v > 0])
|
||||
y1, y2 = np.min(y[v > 0]), np.max(y[v > 0])
|
||||
if x2 - x1 < 5.0/self.width_height[0]:
|
||||
x1 -= 2.0/self.width_height[0]
|
||||
x2 += 2.0/self.width_height[0]
|
||||
if y2 - y1 < 5.0/self.width_height[1]:
|
||||
y1 -= 2.0/self.width_height[1]
|
||||
y2 += 2.0/self.width_height[1]
|
||||
|
||||
bb_list.append(((x1, y1), (x2, y2)))
|
||||
|
||||
# ax.add_patch(
|
||||
# matplotlib.patches.Rectangle(
|
||||
# (x1, y1), x2s - x1, y2 - y1, fill=False, color=color))
|
||||
#
|
||||
# if score:
|
||||
# ax.text(x1, y1, '{:.4f}'.format(score), fontsize=8, color=color)
|
||||
return bb_list
|
||||
|
||||
@staticmethod
|
||||
def keypoint_sets(annotations):
|
||||
keypoint_sets = [ann.data for ann in annotations]
|
||||
# scores = [ann.score() for ann in annotations]
|
||||
# assert len(scores) == len(keypoint_sets)
|
||||
if not keypoint_sets:
|
||||
return np.zeros((0, 17, 3))
|
||||
keypoint_sets = np.array(keypoint_sets)
|
||||
# scores = np.array(scores)
|
||||
|
||||
return keypoint_sets
|
||||
|
||||
def single_image(self, image):
|
||||
# image_bytes = io.BytesIO(base64.b64decode(b64image))
|
||||
# im = PIL.Image.open(image_bytes).convert('RGB')
|
||||
im = PIL.Image.fromarray(image)
|
||||
|
||||
target_wh = self.width_height
|
||||
if (im.size[0] > im.size[1]) != (target_wh[0] > target_wh[1]):
|
||||
target_wh = (target_wh[1], target_wh[0])
|
||||
if im.size[0] != target_wh[0] or im.size[1] != target_wh[1]:
|
||||
# print(f'!!! have to resize image to {target_wh} from {im.size}')
|
||||
im = im.resize(target_wh, PIL.Image.BICUBIC)
|
||||
width_height = im.size
|
||||
|
||||
start = time.time()
|
||||
preprocess = openpifpaf.transforms.Compose([
|
||||
openpifpaf.transforms.NormalizeAnnotations(),
|
||||
openpifpaf.transforms.CenterPadTight(16),
|
||||
openpifpaf.transforms.EVAL_TRANSFORM,
|
||||
])
|
||||
# processed_image, _, __ = preprocess(im, [], None)
|
||||
processed_image = openpifpaf.datasets.PilImageList([im], preprocess=preprocess)[0][0]
|
||||
# processed_image = processed_image_cpu.contiguous().to(self.device, non_blocking=True)
|
||||
# print(f'preprocessing time {time.time() - start}')
|
||||
|
||||
all_fields = self.processor.batch(self.model, torch.unsqueeze(
|
||||
processed_image.float(), 0), device=self.device)[0]
|
||||
keypoint_sets = self.keypoint_sets(all_fields)
|
||||
|
||||
# Normalize scale
|
||||
keypoint_sets[:, :, 0] /= processed_image.shape[2]
|
||||
keypoint_sets[:, :, 1] /= processed_image.shape[1]
|
||||
|
||||
bboxes = self.get_bb(keypoint_sets)
|
||||
return keypoint_sets, bboxes, width_height
|
||||
|
|
@ -0,0 +1,164 @@
|
|||
from enum import IntEnum, unique
|
||||
from typing import List
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
|
||||
@unique
|
||||
class CocoPart(IntEnum):
|
||||
"""Body part locations in the 'coordinates' list."""
|
||||
Nose = 0
|
||||
LEye = 1
|
||||
REye = 2
|
||||
LEar = 3
|
||||
REar = 4
|
||||
LShoulder = 5
|
||||
RShoulder = 6
|
||||
LElbow = 7
|
||||
RElbow = 8
|
||||
LWrist = 9
|
||||
RWrist = 10
|
||||
LHip = 11
|
||||
RHip = 12
|
||||
LKnee = 13
|
||||
RKnee = 14
|
||||
LAnkle = 15
|
||||
RAnkle = 16
|
||||
|
||||
|
||||
SKELETON_CONNECTIONS_COCO = [(0, 1, (210, 182, 247)), (0, 2, (127, 127, 127)), (1, 2, (194, 119, 227)),
|
||||
(1, 3, (199, 199, 199)), (2, 4, (34, 189, 188)), (3, 5, (141, 219, 219)),
|
||||
(4, 6, (207, 190, 23)), (5, 6, (150, 152, 255)), (5, 7, (189, 103, 148)),
|
||||
(5, 11, (138, 223, 152)), (6, 8, (213, 176, 197)), (6, 12, (40, 39, 214)),
|
||||
(7, 9, (75, 86, 140)), (8, 10, (148, 156, 196)), (11, 12, (44, 160, 44)),
|
||||
(11, 13, (232, 199, 174)), (12, 14,
|
||||
(120, 187, 255)), (13, 15, (180, 119, 31)),
|
||||
(14, 16, (14, 127, 255))]
|
||||
|
||||
|
||||
SKELETON_CONNECTIONS_5P = [('H', 'N', (210, 182, 247)), ('N', 'B', (210, 182, 247)), ('B', 'KL', (210, 182, 247)),
|
||||
('B', 'KR', (210, 182, 247)), ('KL', 'KR', (210, 182, 247))]
|
||||
|
||||
COLOR_ARRAY = [(210, 182, 247), (127, 127, 127), (194, 119, 227), (199, 199, 199), (34, 189, 188),
|
||||
(141, 219, 219), (207, 190, 23), (150, 152, 255), (189, 103, 148), (138, 223, 152)]
|
||||
|
||||
UNMATCHED_COLOR = (180, 119, 31)
|
||||
# activity_dict = {
|
||||
# 1.0: "Falling forward using hands",
|
||||
# 2.0: "Falling forward using knees",
|
||||
# 3: "Falling backwards",
|
||||
# 4: "Falling sideward",
|
||||
# 5: "Falling",
|
||||
# 6: "Walking",
|
||||
# 7: "Standing",
|
||||
# 8: "Sitting",
|
||||
# 9: "Picking up an object",
|
||||
# 10: "Jumping",
|
||||
# 11: "Laying",
|
||||
# 12: "False Fall",
|
||||
# 20: "None"
|
||||
# }
|
||||
activity_dict = {
|
||||
1.0: "Falling forward using hands",
|
||||
2.0: "Falling forward using knees",
|
||||
3: "Falling backwards",
|
||||
4: "Falling sideward",
|
||||
5: "FALL",
|
||||
6: "Normal",
|
||||
7: "Normal",
|
||||
8: "Normal",
|
||||
9: "Normal",
|
||||
10: "Normal",
|
||||
11: "Normal",
|
||||
12: "FALL Warning",
|
||||
20: "None"
|
||||
}
|
||||
|
||||
|
||||
def write_on_image(img: np.ndarray, text: str, color: List) -> np.ndarray:
|
||||
"""Write text at the top of the image."""
|
||||
# Add a white border to top of image for writing text
|
||||
img = cv2.copyMakeBorder(src=img,
|
||||
top=int(0.1 * img.shape[0]),
|
||||
bottom=0,
|
||||
left=0,
|
||||
right=0,
|
||||
borderType=cv2.BORDER_CONSTANT,
|
||||
dst=None,
|
||||
value=[255, 255, 255])
|
||||
for i, line in enumerate(text.split('\n')):
|
||||
y = 30 + i * 30
|
||||
cv2.putText(img=img,
|
||||
text=line,
|
||||
org=(0, y),
|
||||
fontFace=cv2.FONT_HERSHEY_SIMPLEX,
|
||||
fontScale=0.7,
|
||||
color=color,
|
||||
thickness=2)
|
||||
|
||||
return img
|
||||
|
||||
|
||||
def visualise(img: np.ndarray, keypoint_sets: List, width: int, height: int, vis_keypoints: bool = False,
|
||||
vis_skeleton: bool = False, CocoPointsOn: bool = False) -> np.ndarray:
|
||||
"""Draw keypoints/skeleton on the output video frame."""
|
||||
|
||||
if CocoPointsOn:
|
||||
SKELETON_CONNECTIONS = SKELETON_CONNECTIONS_COCO
|
||||
else:
|
||||
SKELETON_CONNECTIONS = SKELETON_CONNECTIONS_5P
|
||||
|
||||
if vis_keypoints or vis_skeleton:
|
||||
for keypoints in keypoint_sets:
|
||||
if not CocoPointsOn:
|
||||
keypoints = keypoints["keypoints"]
|
||||
|
||||
if vis_skeleton:
|
||||
for p1i, p2i, color in SKELETON_CONNECTIONS:
|
||||
if keypoints[p1i] is None or keypoints[p2i] is None:
|
||||
continue
|
||||
|
||||
p1 = (int(keypoints[p1i][0] * width), int(keypoints[p1i][1] * height))
|
||||
p2 = (int(keypoints[p2i][0] * width), int(keypoints[p2i][1] * height))
|
||||
|
||||
if p1 == (0, 0) or p2 == (0, 0):
|
||||
continue
|
||||
|
||||
cv2.line(img=img, pt1=p1, pt2=p2, color=color, thickness=3)
|
||||
|
||||
return img
|
||||
|
||||
|
||||
def visualise_tracking(img: np.ndarray, keypoint_sets: List, width: int, height: int, num_matched: int, vis_keypoints: bool = False,
|
||||
vis_skeleton: bool = False, CocoPointsOn: bool = False) -> np.ndarray:
|
||||
"""Draw keypoints/skeleton on the output video frame."""
|
||||
|
||||
if CocoPointsOn:
|
||||
SKELETON_CONNECTIONS = SKELETON_CONNECTIONS_COCO
|
||||
else:
|
||||
SKELETON_CONNECTIONS = SKELETON_CONNECTIONS_5P
|
||||
|
||||
if vis_keypoints or vis_skeleton:
|
||||
for i, keypoints in enumerate(keypoint_sets):
|
||||
if keypoints is None:
|
||||
continue
|
||||
if not CocoPointsOn:
|
||||
keypoints = keypoints["keypoints"]
|
||||
if vis_skeleton:
|
||||
for p1i, p2i, color in SKELETON_CONNECTIONS:
|
||||
if keypoints[p1i] is None or keypoints[p2i] is None:
|
||||
continue
|
||||
|
||||
p1 = (int(keypoints[p1i][0] * width), int(keypoints[p1i][1] * height))
|
||||
p2 = (int(keypoints[p2i][0] * width), int(keypoints[p2i][1] * height))
|
||||
|
||||
if p1 == (0, 0) or p2 == (0, 0):
|
||||
continue
|
||||
if i < num_matched:
|
||||
color = COLOR_ARRAY[i % 10]
|
||||
else:
|
||||
color = UNMATCHED_COLOR
|
||||
|
||||
cv2.line(img=img, pt1=p1, pt2=p2, color=color, thickness=3)
|
||||
|
||||
return img
|
||||
Loading…
Reference in New Issue