110 lines
4.2 KiB
Python
110 lines
4.2 KiB
Python
from typing import List
|
|
|
|
coco_limbSeq = [
|
|
[1, 2], [1, 5], [2, 3], [3, 4],
|
|
[5, 6], [6, 7], [1, 8], [8, 9],
|
|
[9, 10], [1, 11], [11, 12], [12, 13],
|
|
[1, 0], [0, 14], [14, 16], [0, 15],
|
|
[15, 17],
|
|
]
|
|
|
|
coco_colors = [
|
|
[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0],
|
|
[0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255],
|
|
[170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85]
|
|
]
|
|
|
|
body_25_limbSeq = [
|
|
[1, 8], [1, 2], [1, 5], [2, 3],
|
|
[3, 4], [5, 6], [6, 7], [8, 9],
|
|
[9, 10], [10, 11], [8, 12], [12, 13],
|
|
[13, 14], [1, 0], [0, 15], [15, 17],
|
|
[0, 16], [16, 18], [14, 19], [19, 20],
|
|
[14, 21], [11, 22], [22, 23], [11, 24]
|
|
]
|
|
body_25_colors = [
|
|
[255, 0, 85], [255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0],
|
|
[0, 255, 0], [255, 0, 0], [0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255],
|
|
[0, 0, 255], [255, 0, 170], [170, 0, 255], [255, 0, 255], [85, 0, 255], [0, 0, 255], [0, 0, 255],
|
|
[0, 0, 255], [0, 0, 255], [0, 255, 255], [0, 255, 255], [0, 255, 255]
|
|
]
|
|
|
|
body_25B_limbSeq = [
|
|
[0, 1], [0, 2], [1, 3], [2, 4], [5, 7], [6, 8], [7, 9], [8, 10],
|
|
[5, 11], [6, 12], [11, 13], [12, 14], [13, 15], [14, 16], [15, 19],
|
|
[19, 20], [15, 21], [16, 22], [22, 23], [16, 24], [5, 17], [6, 17],
|
|
[17, 18], [11, 12]
|
|
]
|
|
|
|
body_25B_colors = [
|
|
[255, 0, 85], [170, 0, 255], [255, 0, 170], [85, 0, 255], [255, 0, 255],
|
|
[170, 255, 0], [255, 85, 0], [85, 255, 0], [255, 170, 0], [0, 255, 0],
|
|
[255, 255, 0], [0, 170, 255], [0, 255, 85], [0, 85, 255], [0, 255, 170],
|
|
[0, 0, 255], [0, 255, 255], [255, 0, 0], [255, 0, 0], [0, 0, 255],
|
|
[0, 0, 255], [0, 0, 255], [0, 255, 255], [0, 255, 255], [0, 255, 255]
|
|
]
|
|
|
|
class Keypoint:
|
|
def __init__(self, x: float, y: float, confidence: float = 1.0):
|
|
"""
|
|
Initialize a Keypoint object.
|
|
|
|
Args:
|
|
x (float): The x-coordinate of the keypoint.
|
|
y (float): The y-coordinate of the keypoint.
|
|
confidence (float): The confidence score of the keypoint. Default is 1.0.
|
|
"""
|
|
self.x = x
|
|
self.y = y
|
|
self.confidence = confidence
|
|
|
|
def __repr__(self):
|
|
return f"Keypoint(x={self.x}, y={self.y}, confidence={self.confidence})"
|
|
|
|
class Skeleton:
|
|
def __init__(self, keypoints: List[Keypoint]):
|
|
self.keypoints = keypoints
|
|
|
|
def __repr__(self):
|
|
return f"Skeleton(keypoints={self.keypoints})"
|
|
|
|
def is_healthy_skeleton(self):
|
|
for keypoint in self.keypoints:
|
|
if keypoint.confidence == 0.0:
|
|
return False
|
|
return True
|
|
|
|
class Skeleton_Seqence:
|
|
def __init__(self, skeletons: List[Skeleton]):
|
|
self.skeletons_frame = skeletons
|
|
|
|
def __repr__(self):
|
|
return f"Skeleton_Seqence(Skeleton_frames={self.skeletons_frame})"
|
|
|
|
def get_frame(self, frame_index: int) -> Skeleton:
|
|
return self.skeletons_frame[frame_index]
|
|
|
|
def add_frame(self, skeleton: Skeleton):
|
|
self.skeletons_frame.append(skeleton)
|
|
|
|
def is_healthy_seqence(self):
|
|
for skeleton in self.skeletons_frame:
|
|
if not skeleton.is_healthy_skeleton():
|
|
return False
|
|
return True
|
|
|
|
def fix_keypoints(keypoints, last_keypoints, next_keypoints):
|
|
if not keypoints or not last_keypoints or not next_keypoints:
|
|
return keypoints
|
|
for i, keypoint in enumerate(keypoints):
|
|
if keypoint.confidence == 0.0 and last_keypoints and next_keypoints:
|
|
last_keypoint = last_keypoints[i]
|
|
next_keypoint = next_keypoints[i]
|
|
if last_keypoint.confidence > 0 and next_keypoint.confidence > 0:
|
|
keypoint.x = (last_keypoint.x + next_keypoint.x) / 2
|
|
keypoint.y = (last_keypoint.y + next_keypoint.y) / 2
|
|
keypoint.confidence = (last_keypoint.confidence + next_keypoint.confidence) / 2
|
|
return keypoints
|
|
|
|
def get_time_slice_for_Skeleton_Seqences(skeleton_seqences: List[Skeleton_Seqence], frame_index: int) -> List[Skeleton]:
|
|
return [skeleton_seq.get_frame(frame_index) for skeleton_seq in skeleton_seqences] |