def get_pose_bbox(img, human): height, width = img.shape[:2] kp_preds = human['keypoints'] kp_scores = human['kp_score'] kp_preds = torch.cat( (kp_preds, torch.unsqueeze((kp_preds[5, :] + kp_preds[6, :]) / 2, 0))) kp_scores = torch.cat( (kp_scores, torch.unsqueeze((kp_scores[5, :] + kp_scores[6, :]) / 2, 0))) keypoints = [] for n in range(kp_scores.shape[0]): keypoints.append(float(kp_preds[n, 0])) keypoints.append(float(kp_preds[n, 1])) keypoints.append(float(kp_scores[n])) bbox = get_box(keypoints, height, width) return bbox
def vis_frame(frame, im_res, add_bbox=False, format='coco'): ''' frame: frame image im_res: im_res of predictions format: coco or mpii return rendered image ''' if format == 'coco': l_pair = [ (0, 1), (0, 2), (1, 3), (2, 4), # Head (5, 6), (5, 7), (7, 9), (6, 8), (8, 10), (17, 11), (17, 12), # Body (11, 13), (12, 14), (13, 15), (14, 16) ] p_color = [ (0, 255, 255), (0, 191, 255), (0, 255, 102), (0, 77, 255), (0, 255, 0), # Nose, LEye, REye, LEar, REar (77, 255, 255), (77, 255, 204), (77, 204, 255), (191, 255, 77), (77, 191, 255), (191, 255, 77), # LShoulder, RShoulder, LElbow, RElbow, LWrist, RWrist (204, 77, 255), (77, 255, 204), (191, 77, 255), (77, 255, 191), (127, 77, 255), (77, 255, 127), (0, 255, 255) ] # LHip, RHip, LKnee, Rknee, LAnkle, RAnkle, Neck line_color = [(0, 215, 255), (0, 255, 204), (0, 134, 255), (0, 255, 50), (77, 255, 222), (77, 196, 255), (77, 135, 255), (191, 255, 77), (77, 255, 77), (77, 222, 255), (255, 156, 127), (0, 127, 255), (255, 127, 77), (0, 77, 255), (255, 77, 36)] elif format == 'coco_with_foot': l_pair = [ (0, 1), (0, 2), (1, 3), (2, 4), # Head (5, 6), (5, 7), (7, 9), (6, 8), (8, 10), (23, 11), (23, 12), # Body (11, 13), (12, 14), (13, 15), (14, 16), (17, 19), (18, 19), (20, 22), (21, 22), # Foot (15, 19), (16, 22) ] # Foot p_color = [ (0, 255, 255), (0, 191, 255), (0, 255, 102), (0, 77, 255), (0, 255, 0), # Nose, LEye, REye, LEar, REar (77, 255, 255), (77, 255, 204), (77, 204, 255), (191, 255, 77), (77, 191, 255), (191, 255, 77), # LShoulder, RShoulder, LElbow, RElbow, LWrist, RWrist (204, 77, 255), (77, 255, 204), (191, 77, 255), (77, 255, 191), (127, 77, 255), (77, 255, 127), (0, 255, 255), # LHip, RHip, LKnee, Rknee, LAnkle, RAnkle, Neck (0, 255, 255), (0, 191, 255), (0, 255, 102), (0, 77, 255), (0, 255, 0), (77, 255, 255) ] # Foot line_color = [(0, 215, 255), (0, 255, 204), (0, 134, 255), (0, 255, 50), (77, 255, 222), (77, 196, 255), (77, 135, 255), (191, 255, 77), (77, 255, 77), (77, 222, 255), (255, 156, 127), (0, 127, 255), (255, 127, 77), (0, 77, 255), (255, 77, 36), (0, 77, 255), (0, 77, 255), (0, 77, 255), (0, 77, 255), (255, 156, 127), (255, 156, 127)] # Foot elif format == 'mpii': l_pair = [(8, 9), (11, 12), (11, 10), (2, 1), (1, 0), (13, 14), (14, 15), (3, 4), (4, 5), (8, 7), (7, 6), (6, 2), (6, 3), (8, 12), (8, 13)] p_color = [ PURPLE, BLUE, BLUE, RED, RED, BLUE, BLUE, RED, RED, PURPLE, PURPLE, PURPLE, RED, RED, BLUE, BLUE ] line_color = [ PURPLE, BLUE, BLUE, RED, RED, BLUE, BLUE, RED, RED, PURPLE, PURPLE, RED, RED, BLUE, BLUE ] else: raise NotImplementedError # im_name = os.path.basename(im_res['imgname']) img = frame.copy() height, width = img.shape[:2] img = cv2.resize(img, (int(width / 2), int(height / 2))) for human in im_res['result']: part_line = {} kp_preds = human['keypoints'] kp_scores = human['kp_score'] kp_preds = torch.cat( (kp_preds, torch.unsqueeze((kp_preds[5, :] + kp_preds[6, :]) / 2, 0))) kp_scores = torch.cat( (kp_scores, torch.unsqueeze((kp_scores[5, :] + kp_scores[6, :]) / 2, 0))) # Draw bboxes if add_bbox: from PoseFlow.poseflow_infer import get_box keypoints = [] for n in range(kp_scores.shape[0]): keypoints.append(float(kp_preds[n, 0])) keypoints.append(float(kp_preds[n, 1])) keypoints.append(float(kp_scores[n])) bbox = get_box(keypoints, height, width) bg = img.copy() # color = get_color_fast(int(abs(human['idx'][0][0]))) cv2.rectangle(bg, (int(bbox[0] / 2), int(bbox[2] / 2)), (int(bbox[1] / 2), int(bbox[3] / 2)), BLUE, 1) transparency = 0.8 img = cv2.addWeighted(bg, transparency, img, 1 - transparency, 0) # Draw indexes of humans if 'idx' in human.keys(): bg = img.copy() cv2.putText(bg, ''.join(str(human['idx'])), (int(bbox[0] / 2), int((bbox[2] + 26) / 2)), DEFAULT_FONT, 0.5, BLACK, 1) transparency = 0.8 img = cv2.addWeighted(bg, transparency, img, 1 - transparency, 0) # Draw keypoints for n in range(kp_scores.shape[0]): if kp_scores[n] <= 0.35: continue cor_x, cor_y = int(kp_preds[n, 0]), int(kp_preds[n, 1]) part_line[n] = (int(cor_x / 2), int(cor_y / 2)) bg = img.copy() cv2.circle(bg, (int(cor_x / 2), int(cor_y / 2)), 2, p_color[n], -1) # Now create a mask of logo and create its inverse mask also transparency = max(0, min(1, kp_scores[n])) img = cv2.addWeighted(bg, transparency, img, 1 - transparency, 0) # Draw limbs for i, (start_p, end_p) in enumerate(l_pair): if start_p in part_line and end_p in part_line: start_xy = part_line[start_p] end_xy = part_line[end_p] bg = img.copy() X = (start_xy[0], end_xy[0]) Y = (start_xy[1], end_xy[1]) mX = np.mean(X) mY = np.mean(Y) length = ((Y[0] - Y[1])**2 + (X[0] - X[1])**2)**0.5 angle = math.degrees(math.atan2(Y[0] - Y[1], X[0] - X[1])) stickwidth = (kp_scores[start_p] + kp_scores[end_p]) + 1 polygon = cv2.ellipse2Poly((int(mX), int(mY)), (int(length / 2), stickwidth), int(angle), 0, 360, 1) cv2.fillConvexPoly(bg, polygon, line_color[i]) # cv2.line(bg, start_xy, end_xy, line_color[i], (2 * (kp_scores[start_p] + kp_scores[end_p])) + 1) transparency = max( 0, min(1, 0.5 * (kp_scores[start_p] + kp_scores[end_p]))) img = cv2.addWeighted(bg, transparency, img, 1 - transparency, 0) img = cv2.resize(img, (width, height), interpolation=cv2.INTER_CUBIC) return img
def vis_frame_dense(frame, im_res, add_bbox=False, format='coco'): ''' frame: frame image im_res: im_res of predictions format: coco or mpii return rendered image ''' l_pair = [ (0, 1), (0, 2), (1, 3), (2, 4), # Head (5, 19), (19, 21), (21, 17), (17, 22), (22, 20), (20, 6), # Shoulder (6, 28), (28, 30), (30, 8), (8, 32), (32, 34), (34, 10), # right arm (5, 27), (27, 29), (29, 7), (7, 31), (31, 33), (33, 9), #left arm (17, 47), (47, 48), (48, 18), #spine (20, 24), (24, 26), (26, 12), #right spine (19, 23), (23, 25), (25, 11), #left spine (12, 36), (36, 38), (38, 18), (18, 37), (37, 35), (35, 11), #row (12, 40), (40, 42), (42, 14), (14, 44), (44, 46), (46, 16), #right leg (11, 39), (39, 41), (41, 13), (13, 43), (43, 45), (45, 15) #left leg ] p_color = [ (0, 255, 255), (0, 191, 255), (0, 255, 102), (0, 77, 255), (0, 255, 0), # Nose, LEye, REye, LEar, REar (77, 255, 255), (77, 255, 204), (77, 204, 255), (191, 255, 77), (77, 191, 255), (191, 255, 77), # LShoulder, RShoulder, LElbow, RElbow, LWrist, RWrist (204, 77, 255), (77, 255, 204), (191, 77, 255), (77, 255, 191), (127, 77, 255), (77, 255, 127), (0, 255, 255), # LHip, RHip, LKnee, Rknee, LAnkle, RAnkle, Neck (0, 255, 255), (0, 191, 255), (0, 255, 102), (0, 77, 255), (0, 255, 0), # unkonw (77, 255, 255), (77, 255, 204), (77, 204, 255), (191, 255, 77), (77, 191, 255), (191, 255, 77), # unknown (204, 77, 255), (77, 255, 204), (191, 77, 255), (77, 255, 191), (127, 77, 255), (77, 255, 127), (0, 255, 255), #unknown (0, 255, 255), (0, 191, 255), (0, 255, 102), (0, 77, 255), (0, 255, 0), # unkonw (77, 255, 255), (77, 255, 204), (77, 204, 255), (191, 255, 77), (77, 191, 255), (191, 255, 77), # unknown (204, 77, 255), (77, 255, 204), (191, 77, 255), (77, 255, 191), (127, 77, 255), (77, 255, 127), (0, 255, 255) ] #unknown line_color = [ (0, 215, 255), (0, 255, 204), (0, 134, 255), (0, 255, 50), (77, 255, 222), (77, 255, 222), (77, 255, 222), (77, 196, 255), (77, 196, 255), (77, 196, 255), # Shoulder (77, 135, 255), (77, 135, 255), (77, 135, 255), (191, 255, 77), (191, 255, 77), (191, 255, 77), (77, 255, 77), (77, 255, 77), (77, 255, 77), (77, 222, 255), (77, 222, 255), (77, 222, 255), (255, 156, 127), (255, 156, 127), (255, 156, 127), (255, 156, 127), (255, 156, 127), (255, 156, 127), (255, 156, 127), (255, 156, 127), (255, 156, 127), (255, 156, 127), (255, 156, 127), (255, 156, 127), (255, 156, 127), (255, 156, 127), (255, 156, 127), (0, 127, 255), (0, 127, 255), (0, 127, 255), (255, 127, 77), (255, 127, 77), (255, 127, 77), (0, 77, 255), (0, 77, 255), (0, 77, 255), (255, 77, 36), (255, 77, 36), (255, 77, 36) ] # im_name = os.path.basename(im_res['imgname']) img = frame.copy() height, width = img.shape[:2] for human in im_res['result']: part_line = {} kp_preds = human['keypoints'] kp_scores = human['kp_score'] kp_preds = torch.cat( (kp_preds, torch.unsqueeze((kp_preds[5, :] + kp_preds[6, :]) / 2, 0))) kp_scores = torch.cat( (kp_scores, torch.unsqueeze((kp_scores[5, :] + kp_scores[6, :]) / 2, 0))) # Draw bboxes if add_bbox: from PoseFlow.poseflow_infer import get_box keypoints = [] for n in range(kp_scores.shape[0]): keypoints.append(float(kp_preds[n, 0])) keypoints.append(float(kp_preds[n, 1])) keypoints.append(float(kp_scores[n])) bbox = get_box(keypoints, height, width) # color = get_color_fast(int(abs(human['idx'][0]))) cv2.rectangle(img, (int(bbox[0]), int(bbox[2])), (int(bbox[1]), int(bbox[3])), BLUE, 2) # Draw indexes of humans if 'idx' in human.keys(): cv2.putText(img, ''.join(str(human['idx'])), (int(bbox[0]), int( (bbox[2] + 26))), DEFAULT_FONT, 1, BLACK, 2) # Draw keypoints for n in range(kp_scores.shape[0]): if kp_scores[n] <= 0.35: continue cor_x, cor_y = int(kp_preds[n, 0]), int(kp_preds[n, 1]) part_line[n] = (cor_x, cor_y) cv2.circle(img, (cor_x, cor_y), 3, p_color[n], -1) # Draw limbs for i, (start_p, end_p) in enumerate(l_pair): if start_p in part_line and end_p in part_line: start_xy = part_line[start_p] end_xy = part_line[end_p] cv2.line(img, start_xy, end_xy, line_color[i], (kp_scores[start_p] + kp_scores[end_p]) + 1) return img
def vis_frame_fast(frame, im_res, add_bbox=False, format='coco'): ''' frame: frame image im_res: im_res of predictions format: coco or mpii return rendered image ''' if format == 'coco': l_pair = [ (0, 1), (0, 2), (1, 3), (2, 4), # Head (5, 6), (5, 7), (7, 9), (6, 8), (8, 10), (17, 11), (17, 12), # Body (11, 13), (12, 14), (13, 15), (14, 16) ] p_color = [ (0, 255, 255), (0, 191, 255), (0, 255, 102), (0, 77, 255), (0, 255, 0), # Nose, LEye, REye, LEar, REar (77, 255, 255), (77, 255, 204), (77, 204, 255), (191, 255, 77), (77, 191, 255), (191, 255, 77), # LShoulder, RShoulder, LElbow, RElbow, LWrist, RWrist (204, 77, 255), (77, 255, 204), (191, 77, 255), (77, 255, 191), (127, 77, 255), (77, 255, 127), (0, 255, 255) ] # LHip, RHip, LKnee, Rknee, LAnkle, RAnkle, Neck line_color = [(0, 215, 255), (0, 255, 204), (0, 134, 255), (0, 255, 50), (77, 255, 222), (77, 196, 255), (77, 135, 255), (191, 255, 77), (77, 255, 77), (77, 222, 255), (255, 156, 127), (0, 127, 255), (255, 127, 77), (0, 77, 255), (255, 77, 36)] elif format == 'coco_with_foot': l_pair = [ (0, 1), (0, 2), (1, 3), (2, 4), # Head (5, 6), (5, 7), (7, 9), (6, 8), (8, 10), (23, 11), (23, 12), # Body (11, 13), (12, 14), (13, 15), (14, 16), (17, 19), (18, 19), (20, 22), (21, 22), # Foot (15, 19), (16, 22) ] # Foot p_color = [ (0, 255, 255), (0, 191, 255), (0, 255, 102), (0, 77, 255), (0, 255, 0), # Nose, LEye, REye, LEar, REar (77, 255, 255), (77, 255, 204), (77, 204, 255), (191, 255, 77), (77, 191, 255), (191, 255, 77), # LShoulder, RShoulder, LElbow, RElbow, LWrist, RWrist (204, 77, 255), (77, 255, 204), (191, 77, 255), (77, 255, 191), (127, 77, 255), (77, 255, 127), (0, 255, 255), # LHip, RHip, LKnee, Rknee, LAnkle, RAnkle, Neck (0, 255, 255), (0, 191, 255), (0, 255, 102), (0, 77, 255), (0, 255, 0), (77, 255, 255) ] # Foot line_color = [(0, 215, 255), (0, 255, 204), (0, 134, 255), (0, 255, 50), (77, 255, 222), (77, 196, 255), (77, 135, 255), (191, 255, 77), (77, 255, 77), (77, 222, 255), (255, 156, 127), (0, 127, 255), (255, 127, 77), (0, 77, 255), (255, 77, 36), (0, 77, 255), (0, 77, 255), (0, 77, 255), (0, 77, 255), (255, 156, 127), (255, 156, 127)] # Foot elif format == 'mpii': l_pair = [(8, 9), (11, 12), (11, 10), (2, 1), (1, 0), (13, 14), (14, 15), (3, 4), (4, 5), (8, 7), (7, 6), (6, 2), (6, 3), (8, 12), (8, 13)] p_color = [ PURPLE, BLUE, BLUE, RED, RED, BLUE, BLUE, RED, RED, PURPLE, PURPLE, PURPLE, RED, RED, BLUE, BLUE ] else: NotImplementedError # im_name = os.path.basename(im_res['imgname']) img = frame.copy() height, width = img.shape[:2] for human in im_res['result']: part_line = {} kp_preds = human['keypoints'] kp_scores = human['kp_score'] kp_preds = torch.cat( (kp_preds, torch.unsqueeze((kp_preds[5, :] + kp_preds[6, :]) / 2, 0))) kp_scores = torch.cat( (kp_scores, torch.unsqueeze((kp_scores[5, :] + kp_scores[6, :]) / 2, 0))) # Draw bboxes if add_bbox: from PoseFlow.poseflow_infer import get_box keypoints = [] for n in range(kp_scores.shape[0]): keypoints.append(float(kp_preds[n, 0])) keypoints.append(float(kp_preds[n, 1])) keypoints.append(float(kp_scores[n])) bbox = get_box(keypoints, height, width) # color = get_color_fast(int(abs(human['idx'][0]))) cv2.rectangle(img, (int(bbox[0]), int(bbox[2])), (int(bbox[1]), int(bbox[3])), BLUE, 2) # Draw indexes of humans if 'idx' in human.keys(): cv2.putText(img, ''.join(str(human['idx'])), (int(bbox[0]), int( (bbox[2] + 26))), DEFAULT_FONT, 1, BLACK, 2) # Draw keypoints for n in range(kp_scores.shape[0]): if kp_scores[n] <= 0.35: continue cor_x, cor_y = int(kp_preds[n, 0]), int(kp_preds[n, 1]) part_line[n] = (cor_x, cor_y) cv2.circle(img, (cor_x, cor_y), 3, p_color[n], -1) # Draw limbs for i, (start_p, end_p) in enumerate(l_pair): if start_p in part_line and end_p in part_line: start_xy = part_line[start_p] end_xy = part_line[end_p] cv2.line(img, start_xy, end_xy, line_color[i], 2 * (kp_scores[start_p] + kp_scores[end_p]) + 1) return img