def detect_2d_joints(frame, short=360): """ Args: short: 较短边resize大小 frame: 任意尺寸的RGB图像 Returns: 处理过的图像(ndarray),关节点坐标(NDArray)以及置信度等显示2d姿势相关的要素 """ # 缩放图像和生成目标检测器输入张量 frame = nd.array(frame) x, img = data.transforms.presets.yolo.transform_test(frame, short=short) # print(x.shape, img.shape) # 检测人体 class_ids, scores, bounding_boxes = detector(x) # 生成posenet的输入张量 pose_input, upscale_bbox = detector_to_alpha_pose(img, class_ids, scores, bounding_boxes) # 预测关节点 predict_heatmap = pose_net(pose_input) predict_coords, confidence = heatmap_to_coord_alpha_pose(predict_heatmap, upscale_bbox) # 显示2d姿态 # ax = utils.viz.plot_keypoints(img, predict_coords, confidence, class_ids, bounding_boxes, scores) return { 'img': img, 'coords': predict_coords, 'confidence': confidence, 'class_ids': class_ids, 'bboxes': bounding_boxes, 'scores': scores }
def inference(self, model_input): x, img = model_input class_ids, scores, bounding_boxes = self.detector(x) pose_input, upscale_bbox = detector_to_alpha_pose( img, class_ids, scores, bounding_boxes) predicted_heatmap = self.pose_net(pose_input) pred_coords, confidence = heatmap_to_coord_alpha_pose( predicted_heatmap, upscale_bbox) return pred_coords, confidence
def keypoint_detection(img_path, detector, pose_net): x, img = data.transforms.presets.yolo.load_test(img_path, short=512) class_IDs, scores, bounding_boxs = detector(x) pose_input, upscale_bbox = detector_to_alpha_pose(img, class_IDs, scores, bounding_boxs) predicted_heatmap = pose_net(pose_input) pred_coords, confidence = heatmap_to_coord_alpha_pose(predicted_heatmap, upscale_bbox) ax = plot_keypoints(img, pred_coords, confidence, class_IDs, bounding_boxs, scores, box_thresh=0.5, keypoint_thresh=0.2) plt.show()
def get_pose_estimation(img_object, detector_model="yolo3_mobilenet1.0_coco", pose_model="simple_pose_resnet18_v1b", box_thresh=0.5, keypoint_thresh=0.2): ''' //TODO ''' detector = model_zoo.get_model(detector_model, pretrained=True) pose_net = model_zoo.get_model(pose_model, pretrained=True) # Loading weights for only person class detector.reset_class(["person"], reuse_weights=['person']) try: img_object = utils.download(img_object) except ValueError: pass if "yolo" in detector_model: x, img = data.transforms.presets.yolo.load_test(img_object, short=512) elif "ssd" in detector_model: x, img = data.transforms.presets.ssd.load_test(img_object, short=512) class_IDs, scores, bounding_boxs = detector(x) if "simple_pose" in pose_model: pose_input, upscale_bbox = detector_to_simple_pose( img, class_IDs, scores, bounding_boxs) predicted_heatmap = pose_net(pose_input) pred_coords, confidence = heatmap_to_coord(predicted_heatmap, upscale_bbox) elif "alpha_pose" in pose_model: pose_input, upscale_bbox = detector_to_alpha_pose( img, class_IDs, scores, bounding_boxs) predicted_heatmap = pose_net(pose_input) pred_coords, confidence = heatmap_to_coord_alpha_pose( predicted_heatmap, upscale_bbox) ax = utils.viz.plot_keypoints(img, pred_coords, confidence, class_IDs, bounding_boxs, scores, box_thresh=box_thresh, keypoint_thresh=keypoint_thresh) return ax
def get_poses(self, image): x, image = self.transformer(mx.nd.array(image).astype('uint8'), short=512) x = x.as_in_context(self.ctx) class_IDs, scores, bounding_boxs = self.detector(x) pose_input, upscale_bbox = detector_to_alpha_pose(image, class_IDs, scores, bounding_boxs, ctx=self.ctx) if upscale_bbox is not None and len(upscale_bbox) > 0: pose_input = pose_input.as_in_context(self.ctx) predicted_heatmap = self.pose_net(pose_input).as_in_context( mx.cpu()) return heatmap_to_coord_alpha_pose(predicted_heatmap, upscale_bbox) else: return mx.nd.array([]), mx.nd.array([])
def keypoint_detection(img, detector, pose_net, ctx=mx.cpu(), axes=None): x, img = gcv.data.transforms.presets.yolo.transform_test(img, short=512, max_size=350) x = x.as_in_context(ctx) class_IDs, scores, bounding_boxs = detector(x) plt.cla() pose_input, upscale_bbox = detector_to_alpha_pose(img, class_IDs, scores, bounding_boxs, output_shape=(128, 96), ctx=ctx) if len(upscale_bbox) > 0: predicted_heatmap = pose_net(pose_input) pred_coords, confidence = heatmap_to_coord_alpha_pose(predicted_heatmap, upscale_bbox) axes = plot_keypoints(img, pred_coords, confidence, class_IDs, bounding_boxs, scores, box_thresh=0.5, keypoint_thresh=0.2, ax=axes) plt.draw() plt.pause(0.001) else: axes = plot_image(frame, ax=axes) plt.draw() plt.pause(0.001) return axes
detector.reset_class(["person"], reuse_weights=['person']) detector.collect_params().reset_ctx(ctx) pose_net.collect_params().reset_ctx(ctx) im_fname = utils.download('https://github.com/dmlc/web-data/blob/master/' + 'gluoncv/pose/soccer.png?raw=true', path='soccer.png') x, img = data.transforms.presets.yolo.load_test(im_fname, short=512) print('Shape of pre-processed image:', x.shape) x_ctx = x.as_in_context(ctx) # warming up class_IDs, scores, bounding_boxs = detector(x_ctx) pose_input, upscale_bbox = detector_to_alpha_pose(img, class_IDs, scores, bounding_boxs, ctx=ctx) pose_input = pose_input.as_in_context(ctx) predicted_heatmap = pose_net(pose_input) pred_coords, confidence = heatmap_to_coord_alpha_pose(predicted_heatmap, upscale_bbox) # Benchmark run_cnt = 100 detector_time = [] pose_time = [] for i in range(run_cnt): t1 = time.time() x_ctx = x.as_in_context(ctx)
def predict(img_path): # 1.预处理输入图像和检测人体 x, img = data.transforms.presets.yolo.load_test(img_path, short=256) start = time.time() # detect persons and bbox class_ids, scores, bounding_boxes = detector(x) # 2.预处理检测器的输出张量作为alpha_pose的输入 pose_input, upscale_bbox = detector_to_alpha_pose(img, class_ids, scores, bounding_boxes) global detector_time detector_time += (time.time() - start) print("detector cost time: {:.3f} seconds".format(time.time() - start)) # 3.预测关节点 start_time = time.time() predicted_heatmap = pose_net(pose_input) pred_coords, confidence = heatmap_to_coord_alpha_pose( predicted_heatmap, upscale_bbox) global predictor_2d_time predictor_2d_time += (time.time() - start_time) print("2d pose predictor cost time: {:.3f} seconds".format(time.time() - start_time)) # 4.显示2d姿态 # ax = utils.viz.plot_keypoints(img, pred_coords, confidence, class_IDs, bounding_boxes, scores, box_thresh=0.5, # keypoint_thresh=0.2) # 5.坐标标准化 start_time = time.time() kps = normalize_screen_coordinates(pred_coords.asnumpy(), w=img.shape[1], h=img.shape[0]) receptive_field = pose3d_predictor.receptive_field() pad = (receptive_field - 1) // 2 # Padding on each side causal_shift = 0 # 6.创建生成器作为3d预测器的输入 generator = UnchunkedGenerator(None, None, [kps], pad=pad, causal_shift=causal_shift, augment=False) # 7.3d姿势估计和显示 prediction = predict_3d_pos(generator, pose3d_predictor) global predictor_3d_time, full_time predictor_3d_time += time.time() - start_time full_time += time.time() - start print("3d pose predictor cost time: {:.3f} seconds".format(time.time() - start_time)) rot = np.array([0.14070565, -0.15007018, -0.7552408, 0.62232804], dtype=np.float32) prediction = camera_to_world(prediction, R=rot, t=0) prediction[:, :, 2] -= np.min(prediction[:, :, 2]) elapsed = time.time() - start print("Total elapsed time of predicting image {}: {:.3f} seconds".format( img_path, elapsed)) return prediction, img
def video_to_listPose(vid): cap = cv2.VideoCapture(vid) # load video if (cap.isOpened() == False): # Check if camera opened successfully print("Error opening video stream or file") return frame_count = 0 pose_data_vid = [] dimensions = (0, 0) frame_length = (int)(cap.get(cv2.CAP_PROP_FRAME_COUNT)) pbar = tqdm(total=frame_length, ncols=1, desc='.') pbar.ncols = 100 # Iterate through every frame in video while (cap.isOpened()): ret, frame = cap.read() # read current frame if (frame is None): break # If current frame doesn't exist, finished iterating through frames frame = mx.nd.array(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)).astype( 'uint8') # mxnet readable # Person detection x, frame = gcv.data.transforms.presets.yolo.transform_test( frame) # short=406, max_size=1024 class_IDs, scores, bounding_boxs = detector(x.as_in_context(ctx)) # Pose estimation pose_input, upscale_bbox = detector_to_alpha_pose(frame, class_IDs, scores, bounding_boxs, output_shape=(320, 256)) # Gets current pose keypoints if (upscale_bbox is None): # Caters for no person detection pbar.set_description_str('Skipping ') pose_data_curr = [[-1, -1] for j in range(0, 17)] else: # Caters for person detection pbar.set_description_str('Processing') predicted_heatmap = estimator(pose_input) pred_coords, confidence = heatmap_to_coord_alpha_pose( predicted_heatmap, upscale_bbox) scores = scores.asnumpy() confidence = confidence.asnumpy() pred_coords = pred_coords.asnumpy() # Preparing for json pose_data_curr = curr_pose(frame, pred_coords, confidence, scores, keypoint_thresh=0.2) pose_data_vid.append(pose_data_curr) if (frame_count == 0): dimensions = [frame.shape[1], frame.shape[0]] frame_count += 1 pbar.update(1) cap.release() pbar.close() return dimensions, pose_data_vid