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)
