class SKP:
def __init__(self):
super().__init__()
check_license_and_variables_exist()
self.body = Body()
self.confidence_threshold = 0.5
self.skp_output_path = "skp_output"
self.api = Api(default_license_dir())
self.model_path = os.path.join(
os.environ["CUBEMOS_SKEL_SDK"], "models", "skeleton-tracking", "fp32", "skeleton-tracking.cubemos"
)
self.api.load_model(CM_TargetComputeDevice.CM_CPU, self.model_path)
def get_skp_from_pic(self, pic_path):
try:
img = cv2.imread(pic_path)
skeletons = self.api.estimate_keypoints(img, 192)
file_name = os.path.basename(pic_path)
with open("{}/{}.json".format(self.skp_output_path, os.path.splitext(file_name)[0]), "w") as f:
skps = dict()
for i in skeletons:
self.body.set_body(i)
id = str(skeletons.index(i))
skps[id] = dict()
skps[id]['key_frame'] = int(os.path.splitext(file_name)[0].split('_')[-1])
skps[id]['head'] = self.body.get_head_coordinates()
skps[id]['angles'] = self.body.calculate_angles()
json.dump(skps, f)
# isSaved = cv2.imwrite("{}/{}".format(self.output_path, file_name), img)
# return True if isSaved else False
except Exception as ex:
print("Exception occured: \"{}\"".format(ex))
def main():
try:
check_license_and_variables_exist()
sdk_path = os.environ["CUBEMOS_SKEL_SDK"]
#initialize the api with a valid license key in default_license_dir()
api = Api(default_license_dir())
model_path = os.path.join(sdk_path, "models", "skeleton-tracking",
"fp32", "skeleton-tracking.cubemos")
api.load_model(CM_TargetComputeDevice.CM_CPU, model_path)
except Exception as ex:
print("Exception occured: \"{}\"".format(ex))
pipeline = rs.pipeline()
config = rs.config()
config.enable_stream(rs.stream.depth, 1280, 720, rs.format.z16, 30)
config.enable_stream(rs.stream.color, 1280, 720, rs.format.bgr8, 30)
pipeline.start(config)
c = rs.colorizer()
align_to = rs.stream.color
align = rs.align(align_to)
while True:
frames = pipeline.wait_for_frames()
# Align the depth frame to color frame
aligned_frames = align.process(frames)
color = aligned_frames.get_color_frame()
depth = aligned_frames.get_depth_frame()
if not depth: continue
color_data = color.as_frame().get_data()
color_image = np.asanyarray(color_data)
skeletons = api.estimate_keypoints(color_image, 192)
# perform inference again to demonstrate tracking functionality.
# usually you would estimate the keypoints on another image and then
# update the tracking id
#new_skeletons = api.estimate_keypoints(img, 192)
#new_skeletons = api.update_tracking_id(skeletons, new_skeletons)
render_result(skeletons, color_image, confidence_threshold=0.5)
depth_colormap = c.colorize(depth)
depth_data = depth_colormap.as_frame().get_data()
depth_image = np.asanyarray(depth_data)
depth_image = cv2.cvtColor(depth_image, cv2.COLOR_RGB2BGR)
added_image = cv2.addWeighted(color_image, 0.5, depth_image, 0.5, 0)
cv2.imshow('color', color_image)
cv2.imshow('depth', depth_image)
cv2.imshow('overlay', added_image)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
cv2.destroyAllWindows()
return
def run():
try:
check_license_and_variables_exist()
sdk_path = os.environ["CUBEMOS_SKEL_SDK"]
api = Api(default_license_dir())
model_path = os.path.join(sdk_path, "models", "skeleton-tracking",
"fp32", "skeleton-tracking.cubemos")
api.load_model(CM_TargetComputeDevice.CM_CPU, model_path)
body = Body()
vc = VideoControl()
t1 = threading.Thread(target=vc.start_video, name='control')
t1.start()
while True:
frames = pipe.wait_for_frames()
color_frame = frames.get_color_frame()
color_image = np.asanyarray(color_frame.get_data())
color_image = cv2.flip(color_image, 1)
skeletons = api.estimate_keypoints(color_image, 192)
new_skeletons = api.estimate_keypoints(color_image, 192)
new_skeletons = api.update_tracking_id(skeletons, new_skeletons)
render_result(skeletons, color_image, confidence_threshold)
standard = [52, 105, -1, 34, 156, 26, 176, 29, 62]
correct_score_list = list()
for i in skeletons:
body.set_body(i)
correct_score_list.append(
body.compare_skps_angles(10, standard))
# global correct_score
correct_score = max(
correct_score_list) if correct_score_list else 0
if correct_score + 50 >= correct_rate:
cv2.putText(color_image,
"success: {}".format(correct_score + 50), (20, 25),
cv2.FONT_HERSHEY_COMPLEX, 1, (0, 255, 0), 2)
vc.continue_video()
else:
cv2.putText(color_image,
"correct score: {}".format(correct_score + 50),
(20, 25), cv2.FONT_HERSHEY_COMPLEX, 1, (255, 0, 0),
2)
cv2.namedWindow("preview", cv2.WINDOW_AUTOSIZE)
cv2.imshow("preview", color_image)
if cv2.waitKey(1) & 0xFF == ord('q'):
cv2.destroyAllWindows()
break
except Exception as ex:
print("Exception occured: \"{}\"".format(ex))
finally:
pipe.stop()
quit()
class Cubemos_Tacker():
"""Cubemos skeletal tracker class."""
def __init__(self, intrinsics, verbose=False):
"""Constructor.
Args:
intrinsics (np.array): Camera instrinsics matrix
verbose (bool, optional): Verbose log. Defaults to False.
"""
self.verbose = verbose
self.intrinsics = intrinsics
self.confidence_threshold = 0.3
self.network_height = 16 * 20
self.skeletons = []
self.new_skeletons = []
self.skel3d_np = []
self.skel_ids = []
self.roots_2d = []
self.backs_2d = []
self.keypoint_map = np.asarray(
[18, 19, 1, 0, 5, 6, 7, 2, 3, 4, 11, 12, 13, 8, 9, 10])
self.init_skel_track()
self.init_cubemos_api()
def init_skel_track(self):
"""Skeleton tracking initializer."""
check_license_and_variables_exist()
# Get the path of the native libraries and ressource files
self.sdk_path = os.environ["CUBEMOS_SKEL_SDK"]
if self.verbose:
initialise_logging(self.sdk_path, CM_LogLevel.CM_LL_DEBUG, True,
default_log_dir())
def init_cubemos_api(self):
"""Setup cubemos API."""
# initialize the api with a valid license key in default_license_dir()
self.api = Api(default_license_dir())
model_path = os.path.join(self.sdk_path, "models", "skeleton-tracking",
"fp32", "skeleton-tracking.cubemos")
self.api.load_model(CM_TargetComputeDevice.CM_CPU, model_path)
def track_skeletons(self, image, depth_image):
"""Track skeleton from image frame to deptg frame.
Args:
image (np.array): RGB image
depth_image (np.array): depth image
"""
# perform inference
self.skeletons = self.api.estimate_keypoints(image,
self.network_height)
# perform inference again to demonstrate tracking functionality.
# usually you would estimate the keypoints on another image and then
# update the tracking id
self.new_skeletons = self.api.estimate_keypoints(
image, self.network_height)
self.new_skeletons = self.api.update_tracking_id(
self.skeletons, self.new_skeletons)
self.skel2D_to_skel3D(depth_image)
def render_skeletons(self, image):
"""Overlay detected skeleton on image.
Args:
image (np.array): RGB image
"""
render_result(self.skeletons, image, self.confidence_threshold)
for skel_num, joints in enumerate(self.skel3d_np):
for joint_ndx, pt_3D in enumerate(joints[[0, -1]]):
x_pos = int(self.skeletons[skel_num][0][joint_ndx][0])
y_pos = int(self.skeletons[skel_num][0][joint_ndx][1])
text = f"{pt_3D[0]:.2f}, {pt_3D[1]:.2f}, {pt_3D[2]:.2f}"
text = f"{pt_3D[2]:.2f}"
# text = f"{self.skel_ids[skel_num]}"
cv2.putText(image, text, (x_pos, y_pos),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1,
cv2.LINE_AA)
# add root and back
cv2.circle(image, (int(self.roots_2d[skel_num][0]),
int(self.roots_2d[skel_num][1])), 2,
(255, 0, 255), -1)
cv2.circle(image, (int(self.backs_2d[skel_num][0]),
int(self.backs_2d[skel_num][1])), 2,
(255, 0, 255), -1)
def map_2D_3D(self, pixel, depth):
"""Match camera frame to depth frame
Args:
pixel (np.array): pixel co-ordinates
depth (np.array): depth co-ordinates
Returns:
[np.array]: 3D point co-ordinates
"""
x = (pixel[0] - self.intrinsics.ppx) / self.intrinsics.fx
y = (pixel[1] - self.intrinsics.ppy) / self.intrinsics.fy
r2 = x * x + y * y
f = 1 + self.intrinsics.coeffs[0] * r2 + \
self.intrinsics.coeffs[1] * r2 * r2 + \
self.intrinsics.coeffs[4] * r2 * r2 * r2
ux = x * f + 2 * self.intrinsics.coeffs[2] * x * y + \
self.intrinsics.coeffs[3] * (r2 + 2 * x * x)
uy = y * f + 2 * self.intrinsics.coeffs[3] * x * y + \
self.intrinsics.coeffs[2] * (r2 + 2 * y * y)
x = ux
y = uy
X = depth * x
Y = depth * y
Z = depth
return (X, Y, Z)
def skel2D_to_skel3D(self, depth_image):
"""Convert a 2D skeleton in Camera frame to 3D co-ordinates.
Args:
depth_image (np.array): depth image
"""
self.skel_ids = []
self.skel3d_np = []
self.roots_2d = []
self.backs_2d = []
for skeleton in self.skeletons:
joints = skeleton.joints
# confidences = skeleton[1]
self.skel_ids.append(skeleton.id)
joints_3d = []
for joint in joints:
x_ndx = min(int(joint.x), depth_image.shape[0] - 1)
y_ndx = min(int(joint.y), depth_image.shape[1] - 1)
depth = depth_image[x_ndx, y_ndx]
pt_3D = self.map_2D_3D((x_ndx, y_ndx), depth)
joints_3d.append(pt_3D)
# add root
self.roots_2d.append(((joints[8].x + joints[11].x) / 2,
(joints[8].y + joints[11].y) / 2))
l_hip = joints_3d[8]
r_hip = joints_3d[11]
root_x = (l_hip[0] + r_hip[0]) / 2
root_y = (l_hip[1] + r_hip[1]) / 2
root_z = (l_hip[2] + r_hip[2]) / 2
root = (root_x, root_y, root_z)
joints_3d.append(root)
# add back
self.backs_2d.append(((joints[1].x + self.roots_2d[-1][0]) / 2,
(joints[1].y + self.roots_2d[-1][1]) / 2))
neck = joints_3d[1]
back_x = (neck[0] + root[0]) / 2
back_y = (neck[1] + root[1]) / 2
back_z = (neck[2] + root[2]) / 2
back = (back_x, back_y, back_z)
joints_3d.append(back)
# all all joints
self.skel3d_np.append(joints_3d)
self.skel3d_np = np.asarray(self.skel3d_np)
if len(self.skeletons) > 0:
self.skel3d_np = self.skel3d_np[:, self.keypoint_map, :]
def run(render=False, depth=1850, conts_line_color=(256, 256, 256)):
# try:
check_license_and_variables_exist()
sdk_path = os.environ["CUBEMOS_SKEL_SDK"]
api = Api(default_license_dir())
model_path = os.path.join(sdk_path, "models", "skeleton-tracking", "fp32",
"skeleton-tracking.cubemos")
api.load_model(CM_TargetComputeDevice.CM_CPU, model_path)
res = random.choice(load_res_by_persons(NUM_POSERONS))
base = init_windows(res)
timer = Timer()
while True:
frames = pipe.wait_for_frames()
color_frame = frames.get_color_frame()
color_image = np.asanyarray(color_frame.get_data())
color_image = cv2.flip(color_image, 1)
skeletons = api.estimate_keypoints(color_image, 192)
new_skeletons = api.estimate_keypoints(color_image, 192)
new_skeletons = api.update_tracking_id(skeletons, new_skeletons)
if render: render_result(skeletons, color_image, confidence_threshold)
frames = pipe.wait_for_frames()
depth_frame = frames.get_depth_frame()
if not depth_frame: continue
conts_draw = np.fliplr(np.asanyarray(depth_frame.get_data()))
conts_draw = np.uint8(
np.where((conts_draw > 100) & (conts_draw < depth), 255, 0))
conts_draw = cv2.bilateralFilter(conts_draw, 9, 100, 100)
kernel = np.ones((7, 7), np.uint8)
conts_draw = cv2.morphologyEx(conts_draw, cv2.MORPH_CLOSE, kernel, 3)
contours, hierarchy = cv2.findContours(conts_draw, cv2.RETR_EXTERNAL,
1)
conts_draw = base.copy()
cv2.drawContours(conts_draw, contours, -1, conts_line_color, 5)
correct_score = int(
compare_multi_users(skeletons, get_file_basename(res)) * 100)
if correct_score > correct_rate:
# pass
score_color = (0, 256, 0)
if not timer.timer_started:
timer.start()
else:
if timer.get_time() < correct_maintain_time:
cv2.putText(
conts_draw, "{}".format(
int(correct_maintain_time - timer.get_time())),
(width - 100, 75), cv2.FONT_HERSHEY_COMPLEX, 1,
score_color, 2)
else:
cv2.imwrite('img_saved.jpg', color_image)
cv2.imshow("test", color_image)
timer.stop()
else:
# not pass yet
timer.stop()
score_color = (256, 0, 0)
cv2.putText(conts_draw, "{}".format(correct_score), (width - 100, 50),
cv2.FONT_HERSHEY_COMPLEX, 1, score_color, 2)
cv2.namedWindow("preview", cv2.WINDOW_AUTOSIZE)
cv2.imshow("preview", conts_draw)
key = cv2.waitKey(1)
if key & 0xFF == ord('q') or key == 27:
cv2.destroyAllWindows()
break
while True:
###### Wait for a coherent pair of frames, and align them
frames = pipeline.wait_for_frames()
frames = align.process(frames)
color_frame = frames.get_color_frame()
depth_frame = frames.get_depth_frame()
if not depth_frame or not color_frame: # try until both images are ready
continue
depth_image = np.asanyarray(depth_frame.get_data())
color_image = np.asanyarray(color_frame.get_data())
# perform inference
skeletons = api.estimate_keypoints(color_image, 192)
render_result(skeletons, color_image, CONFIDENCE_THRESHOLD,
depth_frame, depth_scale)
# compute fps
if frame_count == fps_update_rate:
frame_count = 0
fps.update()
curr_fps = fps.fps()
# draw fps on image
cv2.putText(color_image,
str(int(curr_fps)), (0, 15),
cv2.FONT_HERSHEY_SIMPLEX,
fontScale=0.7,
color=(0, 255, 0),