def get_pose_sparse_img(video_name, index, model): path = '/home/molly/UCF_data/jpegs_256/v_' + video_name img_path = path + '/frame' + str(index).zfill(6) + '.jpg' print(img_path) img = cv2.imread(img_path) shape_dst = np.min(img.shape[0:2]) with torch.no_grad(): paf, heatmap, im_scale = get_outputs(img, model, 'rtpose') humans = paf_to_pose_cpp(heatmap, paf, cfg) image_h, image_w = img.shape[:2] pose_image = np.zeros((image_h, image_w), dtype="uint8") pose_image = cv2.cvtColor(pose_image, cv2.COLOR_GRAY2BGR) centers = {} for human in humans: # draw point for i in range(CocoPart.Background.value): if i not in human.body_parts.keys(): continue body_part = human.body_parts[i] #print(body_part) center = (int(body_part.x * image_w + 0.5), int(body_part.y * image_h + 0.5)) centers[i] = center cv2.circle(pose_image, center, 3, CocoColors[i], thickness=3, lineType=8, shift=0) # draw line for pair_order, pair in enumerate(CocoPairsRender): if pair[0] not in human.body_parts.keys( ) or pair[1] not in human.body_parts.keys(): continue cv2.line(pose_image, centers[pair[0]], centers[pair[1]], CocoColors[pair_order], 3) pose_img_resize = cv2.resize(pose_image, (224, 224)) #print(pose_img_resize.shape) if not os.path.exists('/home/molly/UCF_data/pose_flow/' + video_name + '/'): os.mkdir('/home/molly/UCF_data/pose_flow/' + video_name + '/') cv2.imwrite( '/home/molly/UCF_data/pose_flow/' + video_name + '/frame' + str(index).zfill(6) + '.jpg', pose_img_resize)
def run_eval(image_dir, anno_file, vis_dir, model, preprocess): """Run the evaluation on the test set and report mAP score :param model: the model to test :returns: float, the reported mAP score """ coco = COCO(anno_file) cat_ids = coco.getCatIds(catNms=['person']) img_ids = coco.getImgIds(catIds=cat_ids) # img_ids = img_ids[81:82] # img_paths = img_paths[81:82] print("Total number of validation images {}".format(len(img_ids))) # iterate all val images outputs = [] print("Processing Images in validation set") for i in range(len(img_ids)): if i % 10 == 0 and i != 0: print("Processed {} images".format(i)) img = coco.loadImgs(img_ids[i])[0] file_name = img['file_name'] file_path = os.path.join(image_dir, file_name) oriImg = cv2.imread(file_path) # Get the shortest side of the image (either height or width) shape_dst = np.min(oriImg.shape[0:2]) # Get results of original image paf, heatmap, scale_img = get_outputs(oriImg, model, preprocess) humans = paf_to_pose_cpp(heatmap, paf, cfg) out = draw_humans(oriImg, humans) vis_path = os.path.join(vis_dir, file_name) cv2.imwrite(vis_path, out) # subset indicated how many peoples foun in this image. upsample_keypoints = (heatmap.shape[0]*cfg.MODEL.DOWNSAMPLE/scale_img, heatmap.shape[1]*cfg.MODEL.DOWNSAMPLE/scale_img) append_result(img_ids[i], humans, upsample_keypoints, outputs) # Eval and show the final result! return eval_coco(outputs=outputs, annFile=anno_file, imgIds=img_ids)
def detect_keypoints(oriImg): # Get results of original image with torch.no_grad(): paf, heatmap, im_scale = get_outputs(oriImg, model, 'rtpose') # heatmap2d(paf[:,:,1]) print(paf.shape, heatmap.shape, im_scale) humans = paf_to_pose_cpp(heatmap, paf, cfg) image_h, image_w = oriImg.shape[:2] human_keypoints = [] for human in humans: # draw point centers = {} for i in range(CocoPart.Background.value): if i not in human.body_parts.keys(): continue body_part = human.body_parts[i] center = (int(body_part.x * image_w + 0.5), int(body_part.y * image_h + 0.5)) centers[i] = {'center': center, 'score': body_part.score} human_keypoints.append(centers) return human_keypoints
def main(args): sys.path.append( args.openpose_dir) # In case calling from an external script from lib.network.rtpose_vgg import get_model from lib.network.rtpose_vgg import use_vgg from lib.network import im_transform from evaluate.coco_eval import get_outputs, handle_paf_and_heat from lib.utils.common import Human, BodyPart, CocoPart, CocoColors, CocoPairsRender, draw_humans from lib.utils.paf_to_pose import paf_to_pose_cpp from lib.config import cfg, update_config update_config(cfg, args) model = get_model('vgg19') model = torch.nn.DataParallel(model).cuda() use_vgg(model) # model.load_state_dict(torch.load(args.weight)) checkpoint = torch.load(args.weight) epoch = checkpoint['epoch'] best_loss = checkpoint['best_loss'] state_dict = checkpoint['state_dict'] # state_dict = {key.replace("module.",""):value for key, value in state_dict.items()} # Remove "module." from vgg keys model.load_state_dict(state_dict) # optimizer.load_state_dict(checkpoint['optimizer']) print("=> loaded checkpoint '{}' (epoch {})".format(args.weight, epoch)) model.float() model.eval() image_folders = args.image_folders.split(',') for i, image_folder in enumerate(image_folders): print( f"\nProcessing {i} of {len(image_folders)}: {' '.join(image_folder.split('/')[-4:-2])}" ) if args.all_frames: # Split video and run inference on all frames output_dir = os.path.join(os.path.dirname(image_folder), 'predictions', 'pose2d', 'openpose_pytorch_ft_all') os.makedirs(output_dir, exist_ok=True) video_path = os.path.join( image_folder, 'scan_video.avi') # break up video and run on all frames temp_folder = image_folder.split('/')[-3] + '_openpose' image_folder = os.path.join( '/tmp', f'{temp_folder}') # Overwrite image_folder os.makedirs(image_folder, exist_ok=True) split_video(video_path, image_folder) else: # Just use GT-annotated frames output_dir = os.path.join(os.path.dirname(image_folder), 'predictions', 'pose2d', 'openpose_pytorch_ft') os.makedirs(output_dir, exist_ok=True) img_mask = os.path.join(image_folder, '??????.png') img_names = glob(img_mask) for img_name in img_names: image_file_path = img_name oriImg = cv2.imread(image_file_path) # B,G,R order shape_dst = np.min(oriImg.shape[0:2]) with torch.no_grad(): paf, heatmap, im_scale = get_outputs(oriImg, model, 'rtpose') humans = paf_to_pose_cpp(heatmap, paf, cfg) # Save joints in OpenPose format image_h, image_w = oriImg.shape[:2] people = [] for i, human in enumerate(humans): keypoints = [] for j in range(18): if j == 8: keypoints.extend([ 0, 0, 0 ]) # Add extra joint (midhip) to correspond to body_25 if j not in human.body_parts.keys(): keypoints.extend([0, 0, 0]) else: body_part = human.body_parts[j] keypoints.extend([ body_part.x * image_w, body_part.y * image_h, body_part.score ]) person = {"person_id": [i - 1], "pose_keypoints_2d": keypoints} people.append(person) people_dict = {"people": people} _, filename = os.path.split(image_file_path) name, _ = os.path.splitext(filename) frame_id = int(name) with open( os.path.join(output_dir, f"scan_video_{frame_id:012}_keypoints.json"), 'w') as outfile: json.dump(people_dict, outfile) if args.all_frames: shutil.rmtree(image_folder) # Delete image_folder
help="Modify config options using the command-line", default=None, nargs=argparse.REMAINDER) args = parser.parse_args() # update config file update_config(cfg, args) weight_name = '/home/tensorboy/Downloads/pose_model.pth' model = get_model('vgg19') model.load_state_dict(torch.load(weight_name)) model = torch.nn.DataParallel(model).cuda() model.float() model.eval() test_image = './readme/ski.jpg' oriImg = cv2.imread(test_image) # B,G,R order shape_dst = np.min(oriImg.shape[0:2]) # Get results of original image with torch.no_grad(): paf, heatmap, im_scale = get_outputs(oriImg, model, 'rtpose') print(im_scale) humans = paf_to_pose_cpp(heatmap, paf, cfg) out = draw_humans(oriImg, humans) cv2.imwrite('result.png', out)
def save_results(loader, model, args): # Get cropped image coordinate conversion factors orig_image_w, orig_image_h = loader.dataset.image_size padding = int(args.square_edge / 2.0) x_offset, y_offset = 0, 0 if orig_image_w > args.square_edge: x_offset = torch.tensor(orig_image_w / 2.0 - padding) x_offset = torch.clamp(x_offset, min=0, max=orig_image_w - args.square_edge).item() if orig_image_h > args.square_edge: y_offset = torch.tensor(orig_image_h / 2.0 - padding) y_offset = torch.clamp(y_offset, min=0, max=orig_image_h - args.square_edge).item() cropped_image_w = min(args.square_edge, orig_image_w - x_offset) cropped_image_h = min(args.square_edge, orig_image_h - y_offset) model.eval() with torch.no_grad(): for i, (img, heatmap_target, paf_target, image_ids) in enumerate(loader): img = img.cuda() outputs, _ = model(img) pafs = outputs[-2].cpu().data.numpy().transpose(0, 2, 3, 1) heatmaps = outputs[-1].cpu().data.numpy().transpose(0, 2, 3, 1) for j, (paf, heatmap) in enumerate(zip(pafs, heatmaps)): id = image_ids[j] id_tokens = id.split('_') # eg Lack_TV_Bench_0007_white_floor_08_04_2019_08_28_10_47_dev3_000100 furniture_type = '_'.join(id_tokens[:3]) experiment_id = '_'.join(id_tokens[3:-2]) cam_id = id_tokens[-2] frame_str = id_tokens[-1] json_name = f"scan_video_000000{frame_str}_keypoints.json" humans = paf_to_pose_cpp(heatmap, paf, cfg) # Save joints in OpenPose format out_image_h, out_image_w = 1080, 1920 people = [] for i, human in enumerate(humans): keypoints = [] for j in range(18): if j == 8: keypoints.extend([0, 0, 0]) # Add extra joint (midhip) to correspond to body_25 if j not in human.body_parts.keys(): keypoints.extend([0, 0, 0]) else: body_part = human.body_parts[j] keypoints.extend([(cropped_image_w * body_part.x + x_offset) * out_image_w / orig_image_w, (cropped_image_h * body_part.y + y_offset) * out_image_h / orig_image_h, body_part.score]) person = { "person_id":[i-1], "pose_keypoints_2d":keypoints } people.append(person) # do_plot = True do_plot = False if do_plot: image_path = os.path.join(args.out_data_dir, furniture_type, experiment_id, cam_id, 'images', frame_str + '.png') out_image = Image.open(image_path) keypoints_np = np.array(keypoints).reshape(-1, 3) plot_skeleton(out_image, keypoints_np) people_dict = {"people":people} people_dict["format"] = "body19" output_path = os.path.join(args.out_data_dir, furniture_type, experiment_id, cam_id, 'predictions', 'pose2d', 'openpose_ft') os.makedirs(output_path, exist_ok=True) json_file = os.path.join(output_path, json_name) print(f"Writing: {json_file}") with open(json_file, 'w') as f: json.dump(people_dict, f)