diff --git a/utils/plots.py b/utils/plots.py index e75bc7b3..fdacc438 100644 --- a/utils/plots.py +++ b/utils/plots.py @@ -431,3 +431,59 @@ def plot_results(start=0, stop=0, bucket='', id=(), labels=(), save_dir=''): ax[1].legend() fig.savefig(Path(save_dir) / 'results.png', dpi=200) + + +def output_to_keypoint(output): + # Convert model output to target format [batch_id, class_id, x, y, w, h, conf] + targets = [] + for i, o in enumerate(output): + kpts = o[:,6:] + o = o[:,:6] + for index, (*box, conf, cls) in enumerate(o.cpu().numpy()): + targets.append([i, cls, *list(*xyxy2xywh(np.array(box)[None])), conf, *list(kpts.cpu().numpy()[index])]) + return np.array(targets) + + +def plot_skeleton_kpts(im, kpts, steps, orig_shape=None): + #Plot the skeleton and keypointsfor coco datatset + 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]] + radius = 5 + num_kpts = len(kpts) // steps + + for kid in range(num_kpts): + r, g, b = pose_kpt_color[kid] + x_coord, y_coord = kpts[steps * kid], kpts[steps * kid + 1] + if not (x_coord % 640 == 0 or y_coord % 640 == 0): + if steps == 3: + conf = kpts[steps * kid + 2] + if conf < 0.5: + continue + cv2.circle(im, (int(x_coord), int(y_coord)), radius, (int(r), int(g), int(b)), -1) + + for sk_id, sk in enumerate(skeleton): + r, g, b = pose_limb_color[sk_id] + pos1 = (int(kpts[(sk[0]-1)*steps]), int(kpts[(sk[0]-1)*steps+1])) + pos2 = (int(kpts[(sk[1]-1)*steps]), int(kpts[(sk[1]-1)*steps+1])) + if steps == 3: + conf1 = kpts[(sk[0]-1)*steps+2] + conf2 = kpts[(sk[1]-1)*steps+2] + if conf1<0.5 or conf2<0.5: + continue + if pos1[0]%640 == 0 or pos1[1]%640==0 or pos1[0]<0 or pos1[1]<0: + continue + if pos2[0] % 640 == 0 or pos2[1] % 640 == 0 or pos2[0]<0 or pos2[1]<0: + continue + cv2.line(im, pos1, pos2, (int(r), int(g), int(b)), thickness=2)