Initial commit

This commit is contained in:
王洋 2021-10-23 17:32:59 +08:00
parent 46d36d9bd4
commit 70c5be6697
14 changed files with 679 additions and 0 deletions

Binary file not shown.

BIN
examples/outfallingdown.avi Normal file

Binary file not shown.

BIN
input/edlerly_fall.mp4 Normal file

Binary file not shown.

BIN
input/edlerly_fall_5.mp4 Normal file

Binary file not shown.

BIN
input/elderly_fall_4.mp4 Normal file

Binary file not shown.

BIN
input/fallingdown.mp4 Normal file

Binary file not shown.

Binary file not shown.

0
model/__init__.py Normal file
View File

BIN
model/lstm_weights.sav Normal file

Binary file not shown.

56
model/model.py Normal file
View File

@ -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
vis/__init__.py Normal file
View File

347
vis/inv_pendulum.py Normal file
View File

@ -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

112
vis/processor.py Normal file
View File

@ -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

164
vis/visual.py Normal file
View File

@ -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