397 lines
18 KiB
Python
397 lines
18 KiB
Python
from pathlib import Path
|
|
import torch
|
|
import argparse
|
|
import os
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from hamer.configs import CACHE_DIR_HAMER
|
|
from hamer.models import HAMER, download_models, load_hamer#, DEFAULT_CHECKPOINT
|
|
from hamer.utils import recursive_to
|
|
from hamer.datasets.vitdet_dataset import ViTDetDataset, DEFAULT_MEAN, DEFAULT_STD
|
|
from hamer.utils.renderer import Renderer, cam_crop_to_full
|
|
|
|
LIGHT_BLUE=(0.65098039, 0.74117647, 0.85882353)
|
|
openpose_indices = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20]
|
|
gt_indices = openpose_indices
|
|
|
|
from vitpose_model import ViTPoseModel
|
|
|
|
import json
|
|
from typing import Dict, Optional, List, Tuple
|
|
import math
|
|
|
|
def convert_crop_coords_to_orig_img(bbox, keypoints, crop_size):
|
|
# import IPython; IPython.embed(); exit()
|
|
cx, cy, h = bbox[:, 0], bbox[:, 1], bbox[:, 2]
|
|
|
|
# unnormalize to crop coords
|
|
# keypoints = 0.5 * crop_size * (keypoints + 1.0)
|
|
|
|
# rescale to orig img crop
|
|
keypoints *= h[..., None, None] / crop_size
|
|
|
|
# transform into original image coords
|
|
keypoints[:,:,0] = (cx - h/2)[..., None] + keypoints[:,:,0]
|
|
keypoints[:,:,1] = (cy - h/2)[..., None] + keypoints[:,:,1]
|
|
return keypoints
|
|
|
|
def get_keypoints_rectangle(keypoints: np.array, threshold: float) -> Tuple[float, float, float]:
|
|
"""
|
|
Compute rectangle enclosing keypoints above the threshold.
|
|
Args:
|
|
keypoints (np.array): Keypoint array of shape (N, 3).
|
|
threshold (float): Confidence visualization threshold.
|
|
Returns:
|
|
Tuple[float, float, float]: Rectangle width, height and area.
|
|
"""
|
|
print(keypoints.shape)
|
|
valid_ind = keypoints[:, -1] > threshold
|
|
if valid_ind.sum() > 0:
|
|
valid_keypoints = keypoints[valid_ind][:, :-1]
|
|
max_x = valid_keypoints[:,0].max()
|
|
max_y = valid_keypoints[:,1].max()
|
|
min_x = valid_keypoints[:,0].min()
|
|
min_y = valid_keypoints[:,1].min()
|
|
width = max_x - min_x
|
|
height = max_y - min_y
|
|
area = width * height
|
|
return width, height, area
|
|
else:
|
|
return 0,0,0
|
|
|
|
def render_keypoints(img: np.array,
|
|
keypoints: np.array,
|
|
pairs: List,
|
|
colors: List,
|
|
thickness_circle_ratio: float,
|
|
thickness_line_ratio_wrt_circle: float,
|
|
pose_scales: List,
|
|
threshold: float = 0.1,
|
|
alpha: float = 1.0) -> np.array:
|
|
"""
|
|
Render keypoints on input image.
|
|
Args:
|
|
img (np.array): Input image of shape (H, W, 3) with pixel values in the [0,255] range.
|
|
keypoints (np.array): Keypoint array of shape (N, 3).
|
|
pairs (List): List of keypoint pairs per limb.
|
|
colors: (List): List of colors per keypoint.
|
|
thickness_circle_ratio (float): Circle thickness ratio.
|
|
thickness_line_ratio_wrt_circle (float): Line thickness ratio wrt the circle.
|
|
pose_scales (List): List of pose scales.
|
|
threshold (float): Only visualize keypoints with confidence above the threshold.
|
|
Returns:
|
|
(np.array): Image of shape (H, W, 3) with keypoints drawn on top of the original image.
|
|
"""
|
|
img_orig = img.copy()
|
|
width, height = img.shape[1], img.shape[2]
|
|
area = width * height
|
|
|
|
lineType = 8
|
|
shift = 0
|
|
numberColors = len(colors)
|
|
thresholdRectangle = 0.1
|
|
|
|
person_width, person_height, person_area = get_keypoints_rectangle(keypoints, thresholdRectangle)
|
|
if person_area > 0:
|
|
ratioAreas = min(1, max(person_width / width, person_height / height))
|
|
thicknessRatio = np.maximum(np.round(math.sqrt(area) * thickness_circle_ratio * ratioAreas), 2)
|
|
thicknessCircle = np.maximum(1, thicknessRatio if ratioAreas > 0.05 else -np.ones_like(thicknessRatio))
|
|
thicknessLine = np.maximum(1, np.round(thicknessRatio * thickness_line_ratio_wrt_circle))
|
|
radius = thicknessRatio / 2
|
|
|
|
img = np.ascontiguousarray(img.copy())
|
|
for i, pair in enumerate(pairs):
|
|
index1, index2 = pair
|
|
if keypoints[index1, -1] > threshold and keypoints[index2, -1] > threshold:
|
|
thicknessLineScaled = int(round(min(thicknessLine[index1], thicknessLine[index2]) * pose_scales[0]))
|
|
colorIndex = index2
|
|
color = colors[colorIndex % numberColors]
|
|
keypoint1 = keypoints[index1, :-1].astype(np.int32)
|
|
keypoint2 = keypoints[index2, :-1].astype(np.int32)
|
|
cv2.line(img, tuple(keypoint1.tolist()), tuple(keypoint2.tolist()), tuple(color.tolist()), thicknessLineScaled, lineType, shift)
|
|
for part in range(len(keypoints)):
|
|
faceIndex = part
|
|
if keypoints[faceIndex, -1] > threshold:
|
|
radiusScaled = int(round(radius[faceIndex] * pose_scales[0]))
|
|
thicknessCircleScaled = int(round(thicknessCircle[faceIndex] * pose_scales[0]))
|
|
colorIndex = part
|
|
color = colors[colorIndex % numberColors]
|
|
center = keypoints[faceIndex, :-1].astype(np.int32)
|
|
cv2.circle(img, tuple(center.tolist()), radiusScaled, tuple(color.tolist()), thicknessCircleScaled, lineType, shift)
|
|
return img
|
|
|
|
def render_hand_keypoints(img, right_hand_keypoints, threshold=0.1, use_confidence=False, map_fn=lambda x: np.ones_like(x), alpha=1.0):
|
|
if use_confidence and map_fn is not None:
|
|
#thicknessCircleRatioLeft = 1./50 * map_fn(left_hand_keypoints[:, -1])
|
|
thicknessCircleRatioRight = 1./50 * map_fn(right_hand_keypoints[:, -1])
|
|
else:
|
|
#thicknessCircleRatioLeft = 1./50 * np.ones(left_hand_keypoints.shape[0])
|
|
thicknessCircleRatioRight = 1./50 * np.ones(right_hand_keypoints.shape[0])
|
|
thicknessLineRatioWRTCircle = 0.75
|
|
pairs = [0,1, 1,2, 2,3, 3,4, 0,5, 5,6, 6,7, 7,8, 0,9, 9,10, 10,11, 11,12, 0,13, 13,14, 14,15, 15,16, 0,17, 17,18, 18,19, 19,20]
|
|
pairs = np.array(pairs).reshape(-1,2)
|
|
|
|
colors = [100., 100., 100.,
|
|
100., 0., 0.,
|
|
150., 0., 0.,
|
|
200., 0., 0.,
|
|
255., 0., 0.,
|
|
100., 100., 0.,
|
|
150., 150., 0.,
|
|
200., 200., 0.,
|
|
255., 255., 0.,
|
|
0., 100., 50.,
|
|
0., 150., 75.,
|
|
0., 200., 100.,
|
|
0., 255., 125.,
|
|
0., 50., 100.,
|
|
0., 75., 150.,
|
|
0., 100., 200.,
|
|
0., 125., 255.,
|
|
100., 0., 100.,
|
|
150., 0., 150.,
|
|
200., 0., 200.,
|
|
255., 0., 255.]
|
|
colors = np.array(colors).reshape(-1,3)
|
|
#colors = np.zeros_like(colors)
|
|
poseScales = [1]
|
|
#img = render_keypoints(img, left_hand_keypoints, pairs, colors, thicknessCircleRatioLeft, thicknessLineRatioWRTCircle, poseScales, threshold, alpha=alpha)
|
|
img = render_keypoints(img, right_hand_keypoints, pairs, colors, thicknessCircleRatioRight, thicknessLineRatioWRTCircle, poseScales, threshold, alpha=alpha)
|
|
#img = render_keypoints(img, right_hand_keypoints, pairs, colors, thickness_circle_ratio, thickness_line_ratio_wrt_circle, pose_scales, 0.1)
|
|
return img
|
|
|
|
def render_openpose(img: np.array,
|
|
hand_keypoints: np.array) -> np.array:
|
|
"""
|
|
Render keypoints in the OpenPose format on input image.
|
|
Args:
|
|
img (np.array): Input image of shape (H, W, 3) with pixel values in the [0,255] range.
|
|
body_keypoints (np.array): Keypoint array of shape (N, 3); 3 <====> (x, y, confidence).
|
|
Returns:
|
|
(np.array): Image of shape (H, W, 3) with keypoints drawn on top of the original image.
|
|
"""
|
|
img = render_hand_keypoints(img, hand_keypoints)
|
|
return img
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description='HaMeR demo code')
|
|
parser.add_argument('--checkpoint', type=str, default='/data/home/zy3023/code/hand/Dyn-HaMR/', help='Path to pretrained model checkpoint')
|
|
parser.add_argument('--img_folder', type=str, default='images', help='Folder with input images')
|
|
parser.add_argument('--out_folder', type=str, default='out_demo', help='Output folder to save rendered results')
|
|
parser.add_argument('--side_view', dest='side_view', action='store_true', default=False, help='If set, render side view also')
|
|
parser.add_argument('--full_frame', dest='full_frame', action='store_true', default=True, help='If set, render all people together also')
|
|
parser.add_argument('--save_mesh', dest='save_mesh', action='store_true', default=False, help='If set, save meshes to disk also')
|
|
parser.add_argument('--batch_size', type=int, default=1, help='Batch size for inference/fitting')
|
|
parser.add_argument('--rescale_factor', type=float, default=2.0, help='Factor for padding the bbox')
|
|
parser.add_argument('--body_detector', type=str, default='vitdet', choices=['vitdet', 'regnety'], help='Using regnety improves runtime and reduces memory')
|
|
parser.add_argument('--file_type', nargs='+', default=['*.jpg', '*.png'], help='List of file extensions to consider')
|
|
|
|
args = parser.parse_args()
|
|
|
|
# Download and load checkpoints
|
|
download_models(CACHE_DIR_HAMER)
|
|
# CACHE_DIR_HAMER = args.checkpoint
|
|
model, model_cfg = load_hamer(args.checkpoint)
|
|
|
|
# Setup HaMeR model
|
|
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
|
|
model = model.to(device)
|
|
model.eval()
|
|
|
|
# Load detector
|
|
from hamer.utils.utils_detectron2 import DefaultPredictor_Lazy
|
|
if args.body_detector == 'vitdet':
|
|
from detectron2.config import LazyConfig
|
|
import hamer
|
|
cfg_path = Path(hamer.__file__).parent/'configs'/'cascade_mask_rcnn_vitdet_h_75ep.py'
|
|
detectron2_cfg = LazyConfig.load(str(cfg_path))
|
|
detectron2_cfg.train.init_checkpoint = "https://dl.fbaipublicfiles.com/detectron2/ViTDet/COCO/cascade_mask_rcnn_vitdet_h/f328730692/model_final_f05665.pkl"
|
|
for i in range(3):
|
|
detectron2_cfg.model.roi_heads.box_predictors[i].test_score_thresh = 0.25
|
|
detector = DefaultPredictor_Lazy(detectron2_cfg)
|
|
elif args.body_detector == 'regnety':
|
|
from detectron2 import model_zoo
|
|
from detectron2.config import get_cfg
|
|
detectron2_cfg = model_zoo.get_config('new_baselines/mask_rcnn_regnety_4gf_dds_FPN_400ep_LSJ.py', trained=True)
|
|
detectron2_cfg.model.roi_heads.box_predictor.test_score_thresh = 0.5
|
|
detectron2_cfg.model.roi_heads.box_predictor.test_nms_thresh = 0.4
|
|
detector = DefaultPredictor_Lazy(detectron2_cfg)
|
|
|
|
# keypoint detector
|
|
cpm = ViTPoseModel(device)
|
|
|
|
# Setup the renderer
|
|
renderer = Renderer(model_cfg, faces=model.mano.faces)
|
|
|
|
# Make output directory if it does not exist
|
|
os.makedirs(args.out_folder, exist_ok=True)
|
|
|
|
# Get all demo images ends with .jpg or .png
|
|
img_paths = [img for end in args.file_type for img in Path(args.img_folder).glob(end)]
|
|
|
|
# Iterate over all images in folder
|
|
for img_path in img_paths:
|
|
img_cv2 = cv2.imread(str(img_path))
|
|
|
|
# Detect humans in image
|
|
det_out = detector(img_cv2)
|
|
img = img_cv2.copy()[:, :, ::-1]
|
|
|
|
det_instances = det_out['instances']
|
|
valid_idx = (det_instances.pred_classes==0) & (det_instances.scores > 0.5)
|
|
pred_bboxes=det_instances.pred_boxes.tensor[valid_idx].cpu().numpy()
|
|
pred_scores=det_instances.scores[valid_idx].cpu().numpy()
|
|
|
|
# Detect human keypoints for each person
|
|
vitposes_out = cpm.predict_pose(
|
|
img,
|
|
[np.concatenate([pred_bboxes, pred_scores[:, None]], axis=1)],
|
|
)
|
|
|
|
bboxes = []
|
|
is_right = []
|
|
|
|
# Use hands based on hand keypoint detections
|
|
for vitposes in vitposes_out:
|
|
left_hand_keyp = vitposes['keypoints'][-42:-21]
|
|
right_hand_keyp = vitposes['keypoints'][-21:]
|
|
|
|
# Rejecting not confident detections
|
|
keyp = left_hand_keyp
|
|
valid = keyp[:,2] > 0.5
|
|
if sum(valid) > 3:
|
|
bbox = [keyp[valid,0].min(), keyp[valid,1].min(), keyp[valid,0].max(), keyp[valid,1].max()]
|
|
bboxes.append(bbox)
|
|
is_right.append(0)
|
|
keyp = right_hand_keyp
|
|
valid = keyp[:,2] > 0.5
|
|
if sum(valid) > 3:
|
|
bbox = [keyp[valid,0].min(), keyp[valid,1].min(), keyp[valid,0].max(), keyp[valid,1].max()]
|
|
bboxes.append(bbox)
|
|
is_right.append(1)
|
|
|
|
if len(bboxes) == 0:
|
|
continue
|
|
|
|
boxes = np.stack(bboxes)
|
|
right = np.stack(is_right)
|
|
|
|
# Run reconstruction on all detected hands
|
|
dataset = ViTDetDataset(model_cfg, img_cv2, boxes, right, rescale_factor=args.rescale_factor)
|
|
dataloader = torch.utils.data.DataLoader(dataset, batch_size=8, shuffle=False, num_workers=0)
|
|
|
|
all_verts = []
|
|
all_cam_t = []
|
|
all_right = []
|
|
all_pred_2d = []
|
|
all_bboxes = []
|
|
for batch in dataloader:
|
|
batch = recursive_to(batch, device)
|
|
with torch.no_grad():
|
|
out = model(batch)
|
|
|
|
multiplier = (2*batch['right']-1)
|
|
pred_cam = out['pred_cam']
|
|
pred_cam[:,1] = multiplier*pred_cam[:,1]
|
|
box_center = batch["box_center"].float()
|
|
box_size = batch["box_size"].float()
|
|
img_size = batch["img_size"].float()
|
|
multiplier = (2*batch['right']-1)
|
|
scaled_focal_length = model_cfg.EXTRA.FOCAL_LENGTH / model_cfg.MODEL.IMAGE_SIZE * img_size.max()
|
|
pred_cam_t_full = cam_crop_to_full(pred_cam, box_center, box_size, img_size, scaled_focal_length).detach().cpu().numpy()
|
|
|
|
# Render the result
|
|
batch_size = batch['img'].shape[0]
|
|
for n in range(batch_size):
|
|
# Get filename from path img_path
|
|
img_fn, _ = os.path.splitext(os.path.basename(img_path))
|
|
person_id = int(batch['personid'][n])
|
|
white_img = (torch.ones_like(batch['img'][n]).cpu() - DEFAULT_MEAN[:,None,None]/255) / (DEFAULT_STD[:,None,None]/255)
|
|
input_patch = batch['img'][n].cpu() * (DEFAULT_STD[:,None,None]/255) + (DEFAULT_MEAN[:,None,None]/255)
|
|
input_patch = input_patch.permute(1,2,0).numpy()
|
|
|
|
# regression_img = renderer(out['pred_vertices'][n].detach().cpu().numpy(),
|
|
# out['pred_cam_t'][n].detach().cpu().numpy(),
|
|
# batch['img'][n],
|
|
# mesh_base_color=LIGHT_BLUE,
|
|
# scene_bg_color=(1, 1, 1),
|
|
# )
|
|
|
|
# if args.side_view:
|
|
# side_img = renderer(out['pred_vertices'][n].detach().cpu().numpy(),
|
|
# out['pred_cam_t'][n].detach().cpu().numpy(),
|
|
# white_img,
|
|
# mesh_base_color=LIGHT_BLUE,
|
|
# scene_bg_color=(1, 1, 1),
|
|
# side_view=True)
|
|
# print('side_img: ', side_img.shape, regression_img.shape, input_patch.shape)
|
|
# final_img = np.concatenate([input_patch, regression_img, side_img], axis=1)
|
|
# else:
|
|
# final_img = np.concatenate([input_patch, regression_img], axis=1)
|
|
|
|
# cv2.imwrite(os.path.join(args.out_folder, f'{img_fn}_{person_id}.png'), 255*final_img[:, :, ::-1])
|
|
|
|
# Add all verts and cams to list
|
|
is_right = batch['right'][n].cpu().numpy()
|
|
verts = out['pred_vertices'][n].detach().cpu().numpy()
|
|
pred_joints = out['pred_keypoints_2d'][n].detach().cpu().numpy()
|
|
pred_joints[:,0] = (2*is_right-1)*pred_joints[:,0]
|
|
bbox = batch['bbox'][n].detach().cpu().numpy()
|
|
|
|
verts[:,0] = (2*is_right-1)*verts[:,0]
|
|
cam_t = pred_cam_t_full[n]
|
|
all_verts.append(verts)
|
|
all_pred_2d.append(pred_joints)
|
|
all_bboxes.append(bbox)
|
|
all_cam_t.append(cam_t)
|
|
all_right.append(is_right)
|
|
|
|
# Save all meshes to disk
|
|
if args.save_mesh:
|
|
camera_translation = cam_t.copy()
|
|
tmesh = renderer.vertices_to_trimesh(verts, camera_translation, LIGHT_BLUE, is_right=is_right)
|
|
tmesh.export(os.path.join(args.out_folder, f'{img_fn}_{person_id}.obj'))
|
|
|
|
# Render front view
|
|
if args.full_frame and len(all_verts) > 0:
|
|
all_pred_2d = np.stack(all_pred_2d)
|
|
all_bboxes = np.stack(all_bboxes)
|
|
# print('vit_2d: ', all_vit_2d.shape)
|
|
print('pred_2d: ', all_pred_2d.shape)
|
|
print('all_bboxes: ', all_bboxes.shape)
|
|
all_pred_2d = model_cfg.MODEL.IMAGE_SIZE * (all_pred_2d + 0.5)
|
|
all_pred_2d = convert_crop_coords_to_orig_img(bbox=all_bboxes, keypoints=all_pred_2d, crop_size=model_cfg.MODEL.IMAGE_SIZE)
|
|
all_pred_2d = np.concatenate([all_pred_2d, np.ones((all_pred_2d.shape[0], all_pred_2d.shape[1], 1))], axis=-1)
|
|
|
|
assert len(all_pred_2d.shape) == 3
|
|
assert all_pred_2d.shape[-1] == 3
|
|
all_pred_2d[:,:,-1] = 1
|
|
|
|
misc_args = dict(
|
|
mesh_base_color=LIGHT_BLUE,
|
|
scene_bg_color=(1, 1, 1),
|
|
focal_length=scaled_focal_length,
|
|
)
|
|
cam_view = renderer.render_rgba_multiple(all_verts, cam_t=all_cam_t, render_res=img_size[n], is_right=all_right, **misc_args)
|
|
|
|
# Overlay image
|
|
input_img = img_cv2.astype(np.float32)[:,:,::-1]/255.0
|
|
input_img = np.concatenate([input_img, np.ones_like(input_img[:,:,:1])], axis=2) # Add alpha channel
|
|
input_img_overlay = input_img[:,:,:3] * (1-cam_view[:,:,3:]) + cam_view[:,:,:3] * cam_view[:,:,3:]
|
|
|
|
cv2.imwrite(os.path.join(args.out_folder, f'{img_fn}_all.jpg'), 255*input_img_overlay[:, :, ::-1])
|
|
|
|
# render 2d keypoints
|
|
pred_img = input_img.copy()[:,:,:-1][:,:,::-1] * 255
|
|
for i in range(len(all_verts)):
|
|
body_keypoints_2d = all_pred_2d[i, :21].copy()
|
|
assert (body_keypoints_2d == all_pred_2d[i, :21]).all()
|
|
pred_img = render_openpose(pred_img, body_keypoints_2d)
|
|
|
|
# draw 2d keypoints
|
|
cv2.imwrite(os.path.join(args.out_folder, f'{img_fn}_all_2d.jpg'), pred_img)
|
|
|
|
if __name__ == '__main__':
|
|
main()
|