def __init__(self, checkpoint_path, device, img_mean=np.array([128, 128, 128], dtype=np.float32), img_scale=np.float32(1/255), use_tensorrt=False): from models.with_mobilenet import PoseEstimationWithMobileNet from modules.load_state import load_state self.img_mean = img_mean self.img_scale = img_scale self.device = 'cpu' if device != 'CPU': if torch.cuda.is_available(): self.device = torch.device('cuda:0') else: print('No CUDA device found, inferring on CPU') net = PoseEstimationWithMobileNet() checkpoint = torch.load(checkpoint_path, map_location='cpu') if use_tensorrt: from torch2trt import TRTModule net = TRTModule() net.load_state_dict(checkpoint) else: load_state(net, checkpoint) net = net.to(self.device) net.eval() self.net = net
def init_pose(checkpoint_path): net = PoseEstimationWithMobileNet() checkpoint = torch.load(checkpoint_path, map_location='cpu') load_state(net, checkpoint) net.eval() net = net.cuda() env = [net, []] return None, env
def openpose_to_jit(): x = torch.randn(1,3,256,456) net = PoseEstimationWithMobileNet().cpu() checkpoint = torch.load(r'.\weights\checkpoint_iter_370000.pth', map_location='cpu') load_state(net, checkpoint) net.eval() net(x) script_model = torch.jit.trace(net, x) script_model.save('test.jit')
def __init__(self): self.name = 'OpenPose' net = PoseEstimationWithMobileNet() checkpoint = torch.load('./checkpoint_iter_370000.pth', map_location='cpu') load_state(net, checkpoint) self.net = net.eval() if envars.USE_GPU(): self.net = self.net.cuda() self.stride = 8 self.upsample_ratio = 4 self.height_size = 256 self.kpt_names = [ 'nose', 'neck', 'r_sho', 'r_elb', 'r_wri', 'l_sho', 'l_elb', 'l_wri', 'r_hip', 'r_knee', 'r_ank', 'l_hip', 'l_knee', 'l_ank', 'r_eye', 'l_eye', 'r_ear', 'l_ear' ] self.connections = [('nose', 'r_eye'), ('r_eye', 'r_ear'), ('nose', 'l_eye'), ('l_eye', 'l_ear'), ('nose', 'neck'), ('neck', 'r_sho'), ('r_sho', 'r_elb'), ('r_elb', 'r_wri'), ('neck', 'l_sho'), ('l_sho', 'l_elb'), ('l_elb', 'l_wri'), ('neck', 'r_hip'), ('r_hip', 'r_knee'), ('r_knee', 'r_ank'), ('neck', 'l_hip'), ('l_hip', 'l_knee'), ('l_knee', 'l_ank')]
def init(cpu = False): net = PoseEstimationWithMobileNet() checkpoint_path = "checkpoint_iter_370000.pth" checkpoint = torch.load(checkpoint_path, map_location='cpu') #load the existing model load_state(net, checkpoint) net = net.eval() if not cpu: net = net.cuda() return net
def run_demo(args, image_provider, height_size, cpu, track, smooth): net = PoseEstimationWithMobileNet() checkpoint = torch.load(args.checkpoint_path, map_location='cpu') load_state(net, checkpoint) net = net.eval() if not cpu: net = net.cuda() stride = 8 upsample_ratio = 4 num_keypoints = Pose.num_kpts previous_poses = [] delay = 33 for d in image_provider: img, image_name = d["image"], d["image_name"] orig_img = img.copy() heatmaps, pafs, scale, pad = infer_fast(net, img, height_size, stride, upsample_ratio, cpu) total_keypoints_num = 0 all_keypoints_by_type = [] 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 keypoints_out = [] for n in range(len(pose_entries)): if len(pose_entries[n]) == 0: continue pose_keypoints = np.ones((num_keypoints, 3), dtype=np.float32) * -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_keypoints[kpt_id, 2] = 0.94 if pose_keypoints[kpt_id, 0] != -1 else 0 keypoints_out.append(pose_keypoints) save_json(image_name, keypoints_out, args)
class LightBody(object): def __init__(self, model_path): self.model = PoseEstimationWithMobileNet() checkpoint = torch.load(model_path, map_location='cpu') load_state(self.model, checkpoint) self.model = self.model.eval() self.model = self.model.cuda() def __call__(self, img, height_size=256): stride = 8 upsample_ratio = 4 num_keypoints = Pose.num_kpts orig_img = img.copy() heatmaps, pafs, scale, pad = infer_fast(self.model, img, height_size, stride, upsample_ratio, False) total_keypoints_num = 0 all_keypoints_by_type = [] 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) 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 current_poses = [] 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) for pose in current_poses: pose.draw(img) img = cv2.addWeighted(orig_img, 0.6, img, 0.4, 0) cv2.imwrite("/data1/qilei_chen/DEVELOPMENTS/test1.jpg", img) return current_poses
def _pose_dect_init(self, device): """Initialize the pose detection model. Arguments: device {torch.device}: device to implement the models on. Returns: PoseEstimationWithMobileNet: initialized OpenPose model. """ weight_path = self.__params.pose_weights model = PoseEstimationWithMobileNet() weight = torch.load(weight_path, map_location='cpu') load_state(model, weight) model = model.eval() if device.type != 'cpu': model = model.cuda() return model
def callback(self, data): try: cv_image = self.bridge.imgmsg_to_cv2(data, "bgr8") except CvBridgeError as e: print(e) ## Rescale Image size rescale_factor = 1 width = int(cv_image.shape[1] * rescale_factor) height = int(cv_image.shape[0] * rescale_factor) dim = (width, height) resized_img = cv2.resize(cv_image, dim) net = PoseEstimationWithMobileNet() checkpoint = torch.load( "/home/zheng/lightweight-human-pose-estimation.pytorch/checkpoint_iter_370000.pth", map_location='cpu') load_state(net, checkpoint) height_size = 256 net = net.eval() net = net.cuda() net.eval() stride = 8 upsample_ratio = 4 num_keypoints = Pose.num_kpts previous_poses = [] delay = 33 # img = cv2.imread("/home/zheng/lightweight-human-pose-estimation.pytorch/data/image_1400.jpg") img = asarray(cv_image) orig_img = img heatmaps, pafs, scale, pad = infer_fast(net, img, height_size, stride, upsample_ratio, cpu="store_true") total_keypoints_num = 0 all_keypoints_by_type = [] 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 current_poses = [] ## Collect all keypoint in numpy array to send it to Ros" pose_keypoints_ros_data = np.zeros(16) my_array_for_publishing = Float32MultiArray() #### pose_keypoints = np.ones((num_keypoints, 2), dtype=np.int32) * -1 for kpt_id in range(8): if pose_entries[0][kpt_id] != -1.0: # keypoint was found pose_keypoints[kpt_id, 0] = int( all_keypoints[int(pose_entries[0][kpt_id]), 0]) pose_keypoints[kpt_id, 1] = int( all_keypoints[int(pose_entries[0][kpt_id]), 1]) pose = Pose(pose_keypoints, pose_entries[0][18]) current_poses.append(pose) pose_keypoints_ros_data[2 * kpt_id] = pose.keypoints[kpt_id][0] pose_keypoints_ros_data[2 * kpt_id + 1] = pose.keypoints[kpt_id][1] for pose in current_poses: pose.draw(img) img = cv2.addWeighted(orig_img, 0.6, img, 0.4, 0) my_array_for_publishing.data = [ pose_keypoints_ros_data[0], pose_keypoints_ros_data[1], pose_keypoints_ros_data[2], pose_keypoints_ros_data[3], pose_keypoints_ros_data[4], pose_keypoints_ros_data[5], pose_keypoints_ros_data[6], pose_keypoints_ros_data[7], pose_keypoints_ros_data[8], pose_keypoints_ros_data[9], pose_keypoints_ros_data[10], pose_keypoints_ros_data[11], pose_keypoints_ros_data[12], pose_keypoints_ros_data[13], pose_keypoints_ros_data[14], pose_keypoints_ros_data[15], ] # cv2.imshow('Lightweight Human Pose Estimation Python Demo', img) self.image_pub.publish(self.bridge.cv2_to_imgmsg(img, "bgr8")) self.keypts_pub.publish(my_array_for_publishing) # cv2.imwrite('/home/zheng/Bureau/image_1400_key.jpg',img) cv2.waitKey(2)
setting.text3.SetLabelText("{} : {}".format(mainstudentcap.username, mainstudentcap.score)) setting.text2.SetLabelText("{} : {}".format(secondstudentcap.username, secondstudentcap.score)) setting.text1.SetLabelText("{} : {}".format(thirdstudentcap.username, thirdstudentcap.score)) if __name__ == '__main__': net = PoseEstimationWithMobileNet() checkpoint = torch.load("checkpoint\\checkpoint_iter_370000.pth", map_location='cpu') load_state(net, checkpoint) net = net.cuda() net = net.eval() capture = None capture = cv2.VideoCapture("data\\HsinDance.mp4") capture.set(cv2.CAP_PROP_FRAME_WIDTH, 1280) capture.set(cv2.CAP_PROP_FRAME_HEIGHT, 720) capture4 = cv2.VideoCapture("data\\IronmanDance24.mp4") capture4.set(cv2.CAP_PROP_FRAME_WIDTH, 1280) capture4.set(cv2.CAP_PROP_FRAME_HEIGHT, 720) screen_width = user32.GetSystemMetrics(0) screen_height = user32.GetSystemMetrics(1) capture2 = cv2.VideoCapture(0) capture2.set(cv2.CAP_PROP_FRAME_WIDTH, 1280) capture2.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)
humans.append(human) #socket.send_string(json.dumps(humans)) img = cv2.addWeighted(orig_img, 0.8, img, 0.8, 0) cv2.imshow('Human Pose', img) key = cv2.waitKey(1) if key == 27: # esc return if __name__ == '__main__': net = PoseEstimationWithMobileNet() checkpoint = torch.load("checkpoint.pth.tar", map_location='cpu') load_state(net, checkpoint) net.eval() #example = torch.randn(1, 3, 256, 456) #traced_script_module = torch.jit.script(net, example) #traced_script_module.save('human-pose.pt') #x = torch.rand(1, 3, 256, 456) #sm = torch.jit.trace(net,[Variable(x)]) #sm.save("human-pose.pt") frame_provider = VideoCapture(0) frame_provider.start() run_demo(net, frame_provider, 256, "cuda")
class PoseDetector(object): def __init__(self, model_path: str): self.device = torch.device( "cuda" if torch.cuda.is_available() else "cpu") self.net = PoseEstimationWithMobileNet() self.net.to(self.device) checkpoint = torch.load(model_path) load_state(self.net, checkpoint) self.net.eval() self.image = None self.avg_heatmap = None self.avg_paf = None self.track = True self.stride = 8 self.upsample_ratio = 4 self.height_size = 256 self.smooth = 1 self.num_keypoints = Pose.num_kpts def __inference(self, image, multiscale=False): img = image.copy() base_height = 368 scales = [1] if multiscale: scales = [0.5, 1.0, 1.5, 2.0] stride = 8 normed_img = self.__normalize(img) height, width, _ = normed_img.shape scales_ratios = [ scale * base_height / float(height) for scale in scales ] avg_heatmap = np.zeros((height, width, 19), dtype=np.float32) avg_paf = np.zeros((height, width, 38), dtype=np.float32) for ratio in scales_ratios: scaled_img = cv2.resize(normed_img, (0, 0), fx=ratio, fy=ratio, interpolation=cv2.INTER_CUBIC) min_dims = [base_height, max(scaled_img.shape[1], base_height)] padded_img, pad = self.__pad_width(scaled_img, stride, min_dims) tensor_img = torch.from_numpy(padded_img).permute( 2, 0, 1).unsqueeze(0).float().to(self.device) stages_output = self.net(tensor_img) stage2_heatmap = stages_output[-2] heatmap = np.transpose(stage2_heatmap.squeeze().cpu().data.numpy(), (1, 2, 0)) heatmap = cv2.resize(heatmap, (0, 0), fx=stride, fy=stride, interpolation=cv2.INTER_CUBIC) heatmap = heatmap[pad[0]:heatmap.shape[0] - pad[2], pad[1]:heatmap.shape[1] - pad[3]:, :] heatmap = cv2.resize(heatmap, (width, height), interpolation=cv2.INTER_CUBIC) self.avg_heatmap = avg_heatmap + heatmap / len(scales_ratios) stage2_paf = stages_output[-1] paf = np.transpose(stage2_paf.squeeze().cpu().data.numpy(), (1, 2, 0)) paf = cv2.resize(paf, (0, 0), fx=stride, fy=stride, interpolation=cv2.INTER_CUBIC) paf = paf[pad[0]:paf.shape[0] - pad[2], pad[1]:paf.shape[1] - pad[3], :] paf = cv2.resize(paf, (width, height), interpolation=cv2.INTER_CUBIC) self.avg_paf = avg_paf + paf / len(scales_ratios) def __inference_fast(self, img, net_input_height_size, stride, upsample_ratio): height, width, _ = img.shape self.scale = net_input_height_size / height scaled_img = cv2.resize(img, (0, 0), fx=self.scale, fy=self.scale, interpolation=cv2.INTER_LINEAR) scaled_img = self.__normalize(scaled_img) min_dims = [ net_input_height_size, max(scaled_img.shape[1], net_input_height_size) ] padded_img, self.pad = self.__pad_width(scaled_img, stride, min_dims) tensor_img = torch.from_numpy(padded_img).permute( 2, 0, 1).unsqueeze(0).float() tensor_img = tensor_img.to(self.device) stages_output = self.net(tensor_img) stage2_heatmaps = stages_output[-2] heatmaps = np.transpose(stage2_heatmaps.squeeze().cpu().data.numpy(), (1, 2, 0)) self.avg_heatmap = 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)) self.avg_paf = cv2.resize(pafs, (0, 0), fx=upsample_ratio, fy=upsample_ratio, interpolation=cv2.INTER_CUBIC) total_keypoints_num = 0 all_keypoints_by_type = [] for kpt_idx in range(self.num_keypoints): # 19th for bg total_keypoints_num += extract_keypoints( self.avg_heatmap[:, :, kpt_idx], all_keypoints_by_type, total_keypoints_num) self.pose_entries, self.all_keypoints = group_keypoints( all_keypoints_by_type, self.avg_paf) for kpt_id in range(self.all_keypoints.shape[0]): self.all_keypoints[kpt_id, 0] = \ int((self.all_keypoints[kpt_id, 0] * self.stride / self.upsample_ratio - self.pad[1]) / self.scale) self.all_keypoints[kpt_id, 1] = \ int((self.all_keypoints[kpt_id, 1] * self.stride / self.upsample_ratio - self.pad[0]) / self.scale) def visualize_prediction(self, image): orig_img = image.copy() if not np.array_equal(self.image, image): self.image = image self.__inference_fast(self.image, self.height_size, self.stride, self.upsample_ratio) current_poses = [] for n in range(len(self.pose_entries)): if len(self.pose_entries[n]) == 0: continue pose_keypoints = np.ones( (self.num_keypoints, 2), dtype=np.int32) * -1 for kpt_id in range(self.num_keypoints): if self.pose_entries[n][kpt_id] != -1.0: # keypoint was found pose_keypoints[kpt_id, 0] = int( self.all_keypoints[int(self.pose_entries[n][kpt_id]), 0]) pose_keypoints[kpt_id, 1] = int( self.all_keypoints[int(self.pose_entries[n][kpt_id]), 1]) pose = Pose(pose_keypoints, self.pose_entries[n][18]) current_poses.append(pose) # if self.track: # previous_poses = [] # track_poses(previous_poses, current_poses, smooth=smooth) # previous_poses = current_poses for pose in current_poses: pose.draw(image) image = cv2.addWeighted(orig_img, 0.6, image, 0.4, 0) # plt.imshow(image) # plt.show() return image def __get_auto_grading_outputs(self, image): if not np.array_equal(self.image, image): self.image = image self.__inference_fast(self.image, self.height_size, self.stride, self.upsample_ratio) return self.all_keypoints, self.pose_entries def __visualize_prediction_slow(self, image): if not np.array_equal(self.image, image): self.image = image self.__inference(self.image) total_keypoints_num = 0 all_keypoints_by_type = [] for kpt_idx in range(18): # 19th for bg total_keypoints_num += extract_keypoints( self.avg_heatmap[:, :, kpt_idx], all_keypoints_by_type, total_keypoints_num) pose_entries, all_keypoints = group_keypoints(all_keypoints_by_type, self.avg_paf) coco_keypoints, scores = convert_to_coco_format( pose_entries, all_keypoints) for keypoints in coco_keypoints: for idx in range(len(keypoints) // 3): cv2.circle( image, (int(keypoints[idx * 3]), int(keypoints[idx * 3 + 1])), 3, (255, 0, 255), -1) plt.imshow(image) plt.show() def __get_auto_grading_outputs_slow(self, image): if not np.array_equal(self.image, image): self.image = image self.__inference(self.image) all_peaks = [] peak_counter = 0 thre1 = 0.1 thre2 = 0.05 for part in range(18): map_ori = self.avg_heatmap[:, :, part] one_heatmap = gaussian_filter(map_ori, sigma=3) map_left = np.zeros(one_heatmap.shape) map_left[1:, :] = one_heatmap[:-1, :] map_right = np.zeros(one_heatmap.shape) map_right[:-1, :] = one_heatmap[1:, :] map_up = np.zeros(one_heatmap.shape) map_up[:, 1:] = one_heatmap[:, :-1] map_down = np.zeros(one_heatmap.shape) map_down[:, :-1] = one_heatmap[:, 1:] peaks_binary = np.logical_and.reduce( (one_heatmap >= map_left, one_heatmap >= map_right, one_heatmap >= map_up, one_heatmap >= map_down, one_heatmap > thre1)) peaks = list( zip(np.nonzero(peaks_binary)[1], np.nonzero(peaks_binary)[0])) # note reverse peaks_with_score = [x + (map_ori[x[1], x[0]], ) for x in peaks] peak_id = range(peak_counter, peak_counter + len(peaks)) peaks_with_score_and_id = [ peaks_with_score[i] + (peak_id[i], ) for i in range(len(peak_id)) ] all_peaks.append(peaks_with_score_and_id) peak_counter += len(peaks) # find connection in the specified sequence, center 29 is in the position 15 limbSeq = [[2, 3], [2, 6], [3, 4], [4, 5], [6, 7], [7, 8], [2, 9], [9, 10], [10, 11], [2, 12], [12, 13], [13, 14], [2, 1], [1, 15], [15, 17], [1, 16], [16, 18], [3, 17], [6, 18]] # the middle joints heatmap correspondence mapIdx = [[31, 32], [39, 40], [33, 34], [35, 36], [41, 42], [43, 44], [19, 20], [21, 22], [23, 24], [25, 26], [27, 28], [29, 30], [47, 48], [49, 50], [53, 54], [51, 52], [55, 56], [37, 38], [45, 46]] connection_all = [] special_k = [] mid_num = 10 for k in range(len(mapIdx)): score_mid = self.avg_paf[:, :, [x - 19 for x in mapIdx[k]]] candA = all_peaks[limbSeq[k][0] - 1] candB = all_peaks[limbSeq[k][1] - 1] nA = len(candA) nB = len(candB) if nA != 0 and nB != 0: connection_candidate = [] for i in range(nA): for j in range(nB): vec = np.subtract(candB[j][:2], candA[i][:2]) norm = math.sqrt(vec[0] * vec[0] + vec[1] * vec[1]) vec = np.divide(vec, norm) startend = list( zip( np.linspace(candA[i][0], candB[j][0], num=mid_num), np.linspace(candA[i][1], candB[j][1], num=mid_num))) vec_x = np.array([ score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 0] for I in range(len(startend)) ]) vec_y = np.array([ score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 1] for I in range(len(startend)) ]) score_midpts = np.multiply( vec_x, vec[0]) + np.multiply(vec_y, vec[1]) score_with_dist_prior = sum(score_midpts) / len( score_midpts) + min( 0.5 * image.shape[0] / norm - 1, 0) criterion1 = len(np.nonzero( score_midpts > thre2)[0]) > 0.8 * len(score_midpts) criterion2 = score_with_dist_prior > 0 if criterion1 and criterion2: connection_candidate.append([ i, j, score_with_dist_prior, score_with_dist_prior + candA[i][2] + candB[j][2] ]) connection_candidate = sorted(connection_candidate, key=lambda x: x[2], reverse=True) connection = np.zeros((0, 5)) for c in range(len(connection_candidate)): i, j, s = connection_candidate[c][0:3] if i not in connection[:, 3] and j not in connection[:, 4]: connection = np.vstack( [connection, [candA[i][3], candB[j][3], s, i, j]]) if len(connection) >= min(nA, nB): break connection_all.append(connection) else: special_k.append(k) connection_all.append([]) # last number in each row is the total parts number of that person # the second last number in each row is the score of the overall configuration subset = -1 * np.ones((0, 20)) candidate = np.array( [item for sublist in all_peaks for item in sublist]) for k in range(len(mapIdx)): if k not in special_k: partAs = connection_all[k][:, 0] partBs = connection_all[k][:, 1] indexA, indexB = np.array(limbSeq[k]) - 1 for i in range(len(connection_all[k])): # = 1:size(temp,1) found = 0 subset_idx = [-1, -1] for j in range(len(subset)): # 1:size(subset,1): if subset[j][indexA] == partAs[i] or subset[j][ indexB] == partBs[i]: subset_idx[found] = j found += 1 if found == 1: j = subset_idx[0] if subset[j][indexB] != partBs[i]: subset[j][indexB] = partBs[i] subset[j][-1] += 1 subset[j][-2] += candidate[ partBs[i].astype(int), 2] + connection_all[k][i][2] elif found == 2: # if found 2 and disjoint, merge them j1, j2 = subset_idx membership = ((subset[j1] >= 0).astype(int) + (subset[j2] >= 0).astype(int))[:-2] if len(np.nonzero(membership == 2)[0]) == 0: # merge subset[j1][:-2] += (subset[j2][:-2] + 1) subset[j1][-2:] += subset[j2][-2:] subset[j1][-2] += connection_all[k][i][2] subset = np.delete(subset, j2, 0) else: # as like found == 1 subset[j1][indexB] = partBs[i] subset[j1][-1] += 1 subset[j1][-2] += candidate[ partBs[i].astype(int), 2] + connection_all[k][i][2] # if find no partA in the subset, create a new subset elif not found and k < 17: row = -1 * np.ones(20) row[indexA] = partAs[i] row[indexB] = partBs[i] row[-1] = 2 row[-2] = sum( candidate[connection_all[k][i, :2].astype(int), 2]) + connection_all[k][i][2] subset = np.vstack([subset, row]) # delete some rows of subset which has few parts occur deleteIdx = [] for i in range(len(subset)): if subset[i][-1] < 4 or subset[i][-2] / subset[i][-1] < 0.4: deleteIdx.append(i) subset = np.delete(subset, deleteIdx, axis=0) # subset: n*20 array, 0-17 is the index in candidate, 18 is the total score, 19 is the total parts # candidate: x, y, score, id return candidate, subset @staticmethod def __normalize(img, img_mean=(128, 128, 128), img_scale=1 / 256): img = np.array(img, dtype=np.float32) img = (img - img_mean) * img_scale return img @staticmethod def __pad_width(img, stride, min_dims, pad_value=(0, 0, 0)): 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 def __call__(self, image): return self.__get_auto_grading_outputs(image)
# Optimizes a model # Relies on the platform => Must be compiled for every device # The first command lin parameter must be the path to the old model import sys import torch # import torch2trt from models.with_mobilenet import PoseEstimationWithMobileNet from torch2trt.torch2trt import * MODEL_NAME = sys.argv[1] NEW_MODEL_NAME = MODEL_NAME.replace(".pth", "_opt.pth") WIDTH = 224 HEIGHT = 224 net = PoseEstimationWithMobileNet() net.load_state_dict(torch.load(sys.argv[1]), strict=False) net.eval().cuda() data = torch.ones((1, 3, HEIGHT, WIDTH)).cuda() model_trt = torch2trt(net, [data], fp16_mode=True, max_workspace_size=1 << 25) torch.save(model_trt.state_dict(), NEW_MODEL_NAME)
def convert_to_skelets(in_, out_, cpu=False, height_size=256): # height_size - network input layer height size # cpu - True if we would like to run in CPU print('start convert to skelets') # mask that shows - this is bed mask = cv2.imread(os.path.join('mask', 'mask.jpg'), 0) mask = cv2.normalize(mask, None, alpha=0, beta=1, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_32F) net = PoseEstimationWithMobileNet() load_state(net, checkpoint) net = net.eval() if not cpu: net = net.cuda() stride = 8 upsample_ratio = 4 max_number = 963 num_img = 0 stream = cv2.VideoCapture("rtsp://*****:*****@62.140.233.76:554") # for num in range(0, max_number + 1): while (True): # frame = 'frame' + str(num) + '.jpg' # img = cv2.imread(os.path.join(in_, frame), cv2.IMREAD_COLOR) r, img = stream.read() # cv2.destroyAllWindows() # find the place of the bed - and add border to it, so we can cut the unnecessary part # apply object detection and find bed # output is an image with black pixels of not bed, and white pixels of bed heatmaps, pafs, scale, pad = infer_fast(net, img, height_size, stride, upsample_ratio, cpu) total_keypoints_num = 0 all_keypoints_by_type = [] for kpt_idx in range(18): 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) 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 # how many persons in image num_persons = len(pose_entries) # num_img more than time_period - we delete first second and add the last second bones_detected = np.zeros(len(bones_to_detect)) bones_xa = np.zeros(len(bones_to_detect)) bones_ya = np.zeros(len(bones_to_detect)) bones_xb = np.zeros(len(bones_to_detect)) bones_yb = np.zeros(len(bones_to_detect)) bones_in_bed = np.zeros(len(bones_to_detect)) for n in range(num_persons): count_person_not_in_bed = 1 for id_x in range(len(bones_to_detect)): bones_detected[id_x] = 0 bones_xa[id_x] = 0 bones_ya[id_x] = 0 bones_xb[id_x] = 0 bones_yb[id_x] = 0 bones_in_bed[id_x] = 0 if len(pose_entries[n]) == 0: continue for id_, part_id in enumerate(bones_to_detect): kpt_a_id = BODY_PARTS_KPT_IDS[part_id][0] global_kpt_a_id = pose_entries[n][kpt_a_id] kpt_b_id = BODY_PARTS_KPT_IDS[part_id][1] global_kpt_b_id = pose_entries[n][kpt_b_id] # if both points are detected if global_kpt_a_id != -1 and global_kpt_b_id != -1: bones_xa[id_], bones_ya[id_] = all_keypoints[ int(global_kpt_a_id), 0:2] bones_xb[id_], bones_yb[id_] = all_keypoints[ int(global_kpt_b_id), 0:2] if mask[int(bones_ya[id_])][int( bones_xa[id_])] == 1 and mask[int( bones_yb[id_])][int(bones_xb[id_])] == 1: bones_in_bed[id_] = 1 bones_detected[id_] = 1 sum_bones = 0 for id_, val in enumerate(bones_in_bed): sum_bones += val if sum_bones == len(bones_in_bed): # anomaly # we take mean vector of 2 vectors of bones 6 and 9 bone_xa = (bones_xa[0] + bones_xa[2]) / 2 bone_ya = (bones_ya[0] + bones_ya[2]) / 2 bone_xb = (bones_xb[0] + bones_xb[2]) / 2 bone_yb = (bones_yb[0] + bones_yb[2]) / 2 x1 = bone_xb - bone_xa y1 = bone_yb - bone_ya x2 = 100 y2 = 0 global anomaly_checker alfa = math.acos( (x1 * x2 + y1 * y2) / (math.sqrt(x1**2 + y1**2) * math.sqrt(x2**2 + y2**2))) # if alfa is close to 90 degree - anomaly if min(abs(alfa - rad_90), abs(alfa - rad_270)) <= threshold: print('num_persons', num_persons) if num_persons == 1: anomaly_checker = np.delete(anomaly_checker, 0) anomaly_checker = np.append(anomaly_checker, 1) cv2.imwrite(os.path.join('out_out', frame), img) if np.sum(anomaly_checker) >= SEC_WITHOUT_HELP: print('ALARM!') num_img += 1 if not os.path.exists(out_): os.mkdir(out_) if cv2.waitKey(1) & 0xFF == ord('q'): break print('done convert to skelets')