ADD file via upload
This commit is contained in:
parent
17b68da48f
commit
e6f1ceaef8
|
|
@ -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