# Copyright (c) Facebook, Inc. and its affiliates. import os import os.path as osp import sys import numpy as np import cv2 import math import torch import torchvision.transforms as transforms # from PIL import Image # Code from https://github.com/Daniil-Osokin/lightweight-human-pose-estimation.pytorch/blob/master/demo.py # 2D body pose estimator sys.path.append('lite_openpose') from pose2d_models.with_mobilenet import PoseEstimationWithMobileNet from modules.load_state import load_state from modules.pose import Pose, track_poses from modules.keypoints import extract_keypoints, group_keypoints def normalize(img, img_mean, img_scale): img = np.array(img, dtype=np.float32) img = (img - img_mean) * img_scale return img def pad_width(img, stride, pad_value, min_dims): h, w, _ = img.shape h = min(min_dims[0], h) min_dims[0] = math.ceil(min_dims[0] / float(stride)) * stride min_dims[1] = max(min_dims[1], w) min_dims[1] = math.ceil(min_dims[1] / float(stride)) * stride pad = [] pad.append(int(math.floor((min_dims[0] - h) / 2.0))) pad.append(int(math.floor((min_dims[1] - w) / 2.0))) pad.append(int(min_dims[0] - h - pad[0])) pad.append(int(min_dims[1] - w - pad[1])) padded_img = cv2.copyMakeBorder(img, pad[0], pad[2], pad[1], pad[3], cv2.BORDER_CONSTANT, value=pad_value) return padded_img, pad class BodyPoseEstimator(object): """ Hand Detector for third-view input. It combines a body pose estimator (https://github.com/jhugestar/lightweight-human-pose-estimation.pytorch.git) """ def __init__(self, device='cpu'): # print("Loading Body Pose Estimator") self.device=device self.__load_body_estimator() def __load_body_estimator(self): net = PoseEstimationWithMobileNet() pose2d_checkpoint = "lite_openpose/checkpoint_iter_370000.pth" checkpoint = torch.load(pose2d_checkpoint, map_location='cpu') load_state(net, checkpoint) net = net.eval() net = net.to(self.device) self.model = net #Code from https://github.com/Daniil-Osokin/lightweight-human-pose-estimation.pytorch/demo.py def __infer_fast(self, img, input_height_size, stride, upsample_ratio, cpu=False, pad_value=(0, 0, 0), img_mean=(128, 128, 128), img_scale=1/256): height, width, _ = img.shape scale = input_height_size / height scaled_img = cv2.resize(img, (0, 0), fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC) scaled_img = normalize(scaled_img, img_mean, img_scale) min_dims = [input_height_size, max(scaled_img.shape[1], input_height_size)] padded_img, pad = pad_width(scaled_img, stride, pad_value, min_dims) tensor_img = torch.from_numpy(padded_img).permute(2, 0, 1).unsqueeze(0).float() if not cpu: tensor_img = tensor_img.to(self.device) with torch.no_grad(): stages_output = self.model(tensor_img) stage2_heatmaps = stages_output[-2] heatmaps = np.transpose(stage2_heatmaps.squeeze().cpu().data.numpy(), (1, 2, 0)) heatmaps = cv2.resize(heatmaps, (0, 0), fx=upsample_ratio, fy=upsample_ratio, interpolation=cv2.INTER_CUBIC) stage2_pafs = stages_output[-1] pafs = np.transpose(stage2_pafs.squeeze().cpu().data.numpy(), (1, 2, 0)) pafs = cv2.resize(pafs, (0, 0), fx=upsample_ratio, fy=upsample_ratio, interpolation=cv2.INTER_CUBIC) return heatmaps, pafs, scale, pad def detect_body_pose(self, img): """ Output: current_bbox: BBOX_XYWH """ stride = 8 upsample_ratio = 4 orig_img = img.copy() # forward heatmaps, pafs, scale, pad = self.__infer_fast(img, input_height_size=256, stride=stride, upsample_ratio=upsample_ratio) total_keypoints_num = 0 all_keypoints_by_type = [] num_keypoints = Pose.num_kpts for kpt_idx in range(num_keypoints): # 19th for bg total_keypoints_num += extract_keypoints(heatmaps[:, :, kpt_idx], all_keypoints_by_type, total_keypoints_num) pose_entries, all_keypoints = group_keypoints(all_keypoints_by_type, pafs, demo=True) for kpt_id in range(all_keypoints.shape[0]): all_keypoints[kpt_id, 0] = (all_keypoints[kpt_id, 0] * stride / upsample_ratio - pad[1]) / scale all_keypoints[kpt_id, 1] = (all_keypoints[kpt_id, 1] * stride / upsample_ratio - pad[0]) / scale ''' # print(len(pose_entries)) if len(pose_entries)>1: pose_entries = pose_entries[:1] print("We only support one person currently") # assert len(pose_entries) == 1, "We only support one person currently" ''' current_poses, current_bbox = list(), list() for n in range(len(pose_entries)): if len(pose_entries[n]) == 0: continue pose_keypoints = np.ones((num_keypoints, 2), dtype=np.int32) * -1 for kpt_id in range(num_keypoints): if pose_entries[n][kpt_id] != -1.0: # keypoint was found pose_keypoints[kpt_id, 0] = int(all_keypoints[int(pose_entries[n][kpt_id]), 0]) pose_keypoints[kpt_id, 1] = int(all_keypoints[int(pose_entries[n][kpt_id]), 1]) pose = Pose(pose_keypoints, pose_entries[n][18]) current_poses.append(pose.keypoints) current_bbox.append(np.array(pose.bbox)) # enlarge the bbox for i, bbox in enumerate(current_bbox): x, y, w, h = bbox margin = 0.2 x_margin = int(w * margin) y_margin = int(h * margin) x0 = max(x-x_margin, 0) y0 = max(y-y_margin, 0) x1 = min(x+w+x_margin, orig_img.shape[1]) y1 = min(y+h+y_margin, orig_img.shape[0]) current_bbox[i] = np.array((x0, y0, x1, y1)).astype(np.int32) # ltrb # 只拿一个人 body_point_list = [] if len(current_poses) > 0: for item in current_poses[0]: if item[0] == item[1] == -1: body_point_list += [0.0, 0.0, 0.0] else: body_point_list += [float(item[0]), float(item[1]), 1.0] else: for i in range(18): body_point_list += [0.0, 0.0, 0.0] pose_dict = dict() pose_dict["people"] = [] pose_dict["people"].append({ "person_id": [-1], "pose_keypoints_2d": body_point_list, "hand_left_keypoints_2d": [], "hand_right_keypoints_2d": [], "face_keypoints_2d": [], "pose_keypoints_3d": [], "face_keypoints_3d": [], "hand_left_keypoints_3d": [], "hand_right_keypoints_3d": [], }) return current_poses, current_bbox