Spaces:
Sleeping
Sleeping
| import numpy as np | |
| import cv2 | |
| # draw pose | |
| RADIUS = 5 | |
| palette = np.array([[255, 128, 0], [255, 153, 51], [255, 178, 102], | |
| [230, 230, 0], [255, 153, 255], [153, 204, 255], | |
| [255, 102, 255], [255, 51, 255], [102, 178, 255], | |
| [51, 153, 255], [255, 153, 153], [255, 102, 102], | |
| [255, 51, 51], [153, 255, 153], [102, 255, 102], | |
| [51, 255, 51], [0, 255, 0], [0, 0, 255], [255, 0, 0], | |
| [255, 255, 255]]) | |
| SKELETON = [ | |
| [16, 14], [14, 12], [17, 15], [15, 13], [12, 13], [6, 12], | |
| [7, 13], [6, 7], [6, 8], [7, 9], [8, 10], [9, 11], [2, 3], | |
| [1, 2], [1, 3], [2, 4], [3, 5], [4, 6], [5, 7], | |
| ] | |
| POSE_LIMB_COLOR = palette[[9, 9, 9, 9, 7, 7, 7, | |
| 0, 0, 0, 0, 0, 16, 16, 16, 16, 16, 16, 16]] | |
| POSE_KPT_COLOR = palette[[16, 16, 16, 16, | |
| 16, 0, 0, 0, 0, 0, 0, 9, 9, 9, 9, 9, 9]] | |
| def display_frame_pose(image, bboxes, kpts): | |
| assert len(bboxes) == len(kpts), f"""Length of bboxes({len(bboxes)}) and | |
| kpts({len(kpts)}) must be the same.""" | |
| add_pose_display_meta(image, bboxes, kpts, draw_kpts=True, draw_limbs=True) | |
| def add_pose_display_meta(image, | |
| bboxes, | |
| kpts, | |
| draw_kpts=True, | |
| draw_limbs=True, | |
| KPT_THR=0.1): | |
| # Set radus = 5, thickness line = 4 by image size 1280 | |
| ratio = image.shape[1]/1280 | |
| RADIUS = int(5*ratio) | |
| THICKNESS = int(4*ratio) | |
| for bbox, kpt in zip(bboxes, kpts): | |
| # format kpt | |
| kpt = np.array(kpt).reshape(-1, 2) | |
| # keypoints | |
| if draw_kpts: | |
| # display_meta = pyds.nvds_acquire_display_meta_from_pool(batch_meta) | |
| display_meta_num_circles = 1 | |
| for j, pt in enumerate(kpt): | |
| # if sc < KPT_THR: | |
| # continue | |
| r, g, b = POSE_KPT_COLOR[j] | |
| cnt = display_meta_num_circles-1 | |
| # circle_params = display_meta_circle_params[cnt] | |
| circle_params_xc = int(pt[0]) | |
| circle_params_yc = int(pt[1]) | |
| circle_params_radius = RADIUS | |
| # circle_params_circle_color(r, g, b, 1.0) | |
| circle_params_has_bg_color = 1 | |
| # circle_params_bg_color.(r, g, b, 1.0) | |
| image = cv2.circle(image, (circle_params_xc, circle_params_yc), | |
| circle_params_radius, (int(r), int(g), int(b)), -1) | |
| display_meta_num_circles += 1 | |
| if display_meta_num_circles >= 16: | |
| # display_meta = pyds.nvds_acquire_display_meta_from_pool( | |
| # batch_meta) | |
| display_meta_num_circles = 1 | |
| # import ipdb | |
| # ipdb.set_trace() | |
| # pyds.nvds_add_display_meta_to_frame(frame_meta, display_meta) | |
| # limbs | |
| if draw_limbs: | |
| # display_meta = pyds.nvds_acquire_display_meta_from_pool(batch_meta) | |
| display_meta_num_lines = 1 | |
| for i, (s1, s2) in enumerate(SKELETON): | |
| pt1, pt2 = kpt[s1-1], kpt[s2-1] | |
| # sc1, sc2 = score[s1-1], score[s2-1] | |
| # if sc1 < KPT_THR or sc2 < KPT_THR: | |
| # continue | |
| r, g, b = POSE_LIMB_COLOR[i] | |
| cnt = display_meta_num_lines-1 | |
| # line_params = display_meta_line_params[cnt] | |
| line_params_x1 = int(pt1[0]) | |
| line_params_y1 = int(pt1[1]) | |
| line_params_x2 = int(pt2[0]) | |
| line_params_y2 = int(pt2[1]) | |
| line_params_line_width = 2 | |
| # draw line | |
| cv2.line(image, (line_params_x1, line_params_y1), | |
| (line_params_x2, line_params_y2), (int(r), int(g), int(b)), THICKNESS) | |
| # line_params_line_color.set(r, g, b, 1.0) | |
| display_meta_num_lines += 1 | |
| # if display_meta.num_lines >= 16: | |
| # pyds.nvds_add_display_meta_to_frame( | |
| # frame_meta, display_meta) | |
| # display_meta = pyds.nvds_acquire_display_meta_from_pool( | |
| # batch_meta) | |
| # display_meta.num_lines = 1 | |
| # pyds.nvds_add_display_meta_to_frame(frame_meta, display_meta) | |