diff --git a/.DS_Store b/.DS_Store new file mode 100644 index 0000000..558c56c Binary files /dev/null and b/.DS_Store differ diff --git a/.gitignore b/.gitignore index 81e543a..bb39035 100644 --- a/.gitignore +++ b/.gitignore @@ -1,37 +1,6 @@ -*.swp -*.jpg -*.png -smpl/ -tmp/ -pretrained_models/ -.vscode/ -test_result/ -log/ -log.txt -checkpoints/ -evaluate_results/ -*.pyc -*.npy -test_log.txt -*.pkl -pretrained/ -render_result/ -lightweight-human-pose-estimation.pytorch/ -data/* -webvideos/ -samples/ -*.pth -*.obj -data/ -detectors/ -data -samples_all/ -*.zip -extra_data/ -sample_data/ +*.so +tests.py mocap_output/ -demo.sh -demo_multi.sh -mocap_utils/frame_to_video.py -mocap_utils/frame_to_gif.py -mocap_utils/select_epick_kitchen.py +sample_data/ +extra_data/ +*.onnx diff --git a/__pycache__/render.cpython-39.pyc b/__pycache__/render.cpython-39.pyc new file mode 100644 index 0000000..e92e659 Binary files /dev/null and b/__pycache__/render.cpython-39.pyc differ diff --git a/__pycache__/render_ctypes.cpython-39.pyc b/__pycache__/render_ctypes.cpython-39.pyc new file mode 100644 index 0000000..4da8c74 Binary files /dev/null and b/__pycache__/render_ctypes.cpython-39.pyc differ diff --git a/bodymocap/__pycache__/body_bbox_detector.cpython-39.pyc b/bodymocap/__pycache__/body_bbox_detector.cpython-39.pyc new file mode 100644 index 0000000..094c669 Binary files /dev/null and b/bodymocap/__pycache__/body_bbox_detector.cpython-39.pyc differ diff --git a/bodymocap/__pycache__/body_bbox_detector_ort.cpython-39.pyc b/bodymocap/__pycache__/body_bbox_detector_ort.cpython-39.pyc new file mode 100644 index 0000000..7f846b7 Binary files /dev/null and b/bodymocap/__pycache__/body_bbox_detector_ort.cpython-39.pyc differ diff --git a/bodymocap/__pycache__/body_mocap_api.cpython-39.pyc b/bodymocap/__pycache__/body_mocap_api.cpython-39.pyc new file mode 100644 index 0000000..9ad81f4 Binary files /dev/null and b/bodymocap/__pycache__/body_mocap_api.cpython-39.pyc differ diff --git a/bodymocap/__pycache__/constants.cpython-39.pyc b/bodymocap/__pycache__/constants.cpython-39.pyc new file mode 100644 index 0000000..87049cf Binary files /dev/null and b/bodymocap/__pycache__/constants.cpython-39.pyc differ diff --git a/bodymocap/body_bbox_detector.py b/bodymocap/body_bbox_detector.py index 0aac428..17f8d74 100644 --- a/bodymocap/body_bbox_detector.py +++ b/bodymocap/body_bbox_detector.py @@ -20,6 +20,7 @@ from detectors.body_pose_estimator.val import normalize, pad_width from detectors.body_pose_estimator.modules.pose import Pose, track_poses from detectors.body_pose_estimator.modules.keypoints import extract_keypoints, group_keypoints +from alfred.dl.torch.common import device class BodyPoseEstimator(object): @@ -38,7 +39,7 @@ def __load_body_estimator(self): checkpoint = torch.load(pose2d_checkpoint, map_location='cpu') load_state(net, checkpoint) net = net.eval() - net = net.cuda() + # net = net.cuda() self.model = net @@ -54,8 +55,7 @@ def __infer_fast(self, img, input_height_size, stride, upsample_ratio, 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.cuda() + tensor_img = tensor_img.to(device) stages_output = self.model(tensor_img) diff --git a/bodymocap/body_bbox_detector_ort.py b/bodymocap/body_bbox_detector_ort.py new file mode 100644 index 0000000..d662f3e --- /dev/null +++ b/bodymocap/body_bbox_detector_ort.py @@ -0,0 +1,114 @@ +import argparse +import cv2 +import numpy as np +from alfred.utils.timer import ATimer +from .utils.utils import ORTWrapper +import math +from operator import itemgetter +from alfred.vis.image.pose import vis_pose_result, vis_pose_by_joints +from alfred.utils.file_io import ImageSourceIter +from .utils.utils import normalize, extract_keypoints, connections_nms, group_keypoints + + +class LightweightedPoseDetector: + + def __init__(self, onnx_f='data/human-pose-estimation.onnx') -> None: + self.onnx_model = ORTWrapper(onnx_f) + + self.stride = 8 + self.upsample_ratio = 4 + self.num_keypoints = 18 + self.img_mean = np.array([128, 128, 128]).astype(np.float32) + self.img_scale = np.float32(1/256) + + def infer(self, img): + height, width, _ = img.shape + + net_in_height = 256 + net_in_width = 288 + + scale = min(net_in_height / height, net_in_width/width) + + scaled_img = cv2.resize( + img, (0, 0), fx=scale, fy=scale, interpolation=cv2.INTER_LINEAR) + # scaled_img = normalize(scaled_img, img_mean, img_scale) + s_h, s_w, _ = scaled_img.shape + in_img = np.ones([net_in_height, net_in_width, 3] + ).astype(np.uint8) * 128 + top = (net_in_height - s_h) // 2 + left = (net_in_width - s_w) // 2 + in_img[top: top + s_h, left: left + s_w] = scaled_img + + in_img = normalize(in_img, self.img_mean, self.img_scale, ) + inp_img = np.expand_dims(in_img.transpose((2, 0, 1)), axis=0) + print(inp_img.shape) + stages_output = self.onnx_model.infer(inp_img) + # print(stages_output) + heatmaps = stages_output['stage_1_output_1_heatmaps'] + pafs = stages_output['stage_1_output_0_pafs'] + + heatmaps = heatmaps.squeeze(0) + pafs = pafs.squeeze(0) + return heatmaps, pafs, scale, [top, left] + + def run_one_img(self, img): + heatmaps, pafs, scale, pad = self.infer(img) + + all_keypoints_by_type = [] + total_keypoints_num = 0 + for kpt_idx in range(self.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) + + # h, w + for kpt_id in range(all_keypoints.shape[0]): + all_keypoints[kpt_id, 0] = ( + all_keypoints[kpt_id, 0] * self.stride / self.upsample_ratio - pad[1]) / scale + all_keypoints[kpt_id, 1] = ( + all_keypoints[kpt_id, 1] * self.stride / self.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( + (self.num_keypoints, 2), dtype=np.int32) * -1 + for kpt_id in range(self.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]) + # print(pose_entries[n][18]) + current_poses.append(pose_keypoints) + + if len(current_poses) > 0: + return np.stack(current_poses) + return current_poses + + def get_enlarged_boxes_from_poses(self, poses, img_h, img_w): + current_bbox = [] + for pose in poses: + found_keypoints = np.zeros( + (np.count_nonzero(pose[:, 0] != -1), 2), dtype=np.int32) + found_kpt_id = 0 + for kpt_id in range(self.num_keypoints): + if pose[kpt_id, 0] == -1: + continue + found_keypoints[found_kpt_id] = pose[kpt_id] + found_kpt_id += 1 + bb = cv2.boundingRect(found_keypoints) + current_bbox.append(bb) + for i, bbox in enumerate(current_bbox): + x, y, w, h = bbox + margin = 0.05 + 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, img_w) + y1 = min(y+h+y_margin, img_h) + current_bbox[i] = np.array((x0, y0, x1-x0, y1-y0)).astype(np.int32) + return current_bbox diff --git a/bodymocap/body_mocap_api.py b/bodymocap/body_mocap_api.py index e929d2e..97ee80a 100644 --- a/bodymocap/body_mocap_api.py +++ b/bodymocap/body_mocap_api.py @@ -4,7 +4,7 @@ import sys import torch import numpy as np -import pickle +import pickle from torchvision.transforms import Normalize from bodymocap.models import hmr, SMPL, SMPLX @@ -12,34 +12,54 @@ from bodymocap.utils.imutils import crop, crop_bboxInfo, process_image_bbox, process_image_keypoints, bbox_from_keypoints from mocap_utils.coordconv import convert_smpl_to_bbox, convert_bbox_to_oriIm import mocap_utils.geometry_utils as gu +from alfred.dl.torch.common import device +import os +from .utils.utils import ORTWrapper + +torch.set_grad_enabled(False) class BodyMocap(object): def __init__(self, regressor_checkpoint, smpl_dir, device=torch.device('cuda'), use_smplx=False): - - self.device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') + self.device = device # Load parametric model (SMPLX or SMPL) if use_smplx: smplModelPath = smpl_dir + '/SMPLX_NEUTRAL.pkl' self.smpl = SMPLX(smpl_dir, - batch_size=1, - num_betas = 10, - use_pca = False, - create_transl=False).to(self.device) + batch_size=1, + num_betas=10, + use_pca=False, + create_transl=False).to(self.device) self.use_smplx = True else: smplModelPath = smpl_dir + '/basicModel_neutral_lbs_10_207_0_v1.0.0.pkl' - self.smpl = SMPL(smplModelPath, batch_size=1, create_transl=False).to(self.device) + self.smpl = SMPL(smplModelPath, batch_size=1, + create_transl=False).to(self.device) self.use_smplx = False - - #Load pre-trained neural network + + # Load pre-trained neural network SMPL_MEAN_PARAMS = './extra_data/body_module/data_from_spin/smpl_mean_params.npz' - self.model_regressor = hmr(SMPL_MEAN_PARAMS).to(self.device) - checkpoint = torch.load(regressor_checkpoint) - self.model_regressor.load_state_dict(checkpoint['model'], strict=False) - self.model_regressor.eval() - + self.model_regressor_onnx_file = regressor_checkpoint.replace( + 'pt', 'onnx') + if os.path.exists(self.model_regressor_onnx_file): + self.onnx = True + self.model_regressor = ORTWrapper(self.model_regressor_onnx_file) + else: + self.onnx = False + self.model_regressor = hmr(SMPL_MEAN_PARAMS).to(self.device) + checkpoint = torch.load(regressor_checkpoint, map_location='cpu') + self.model_regressor.load_state_dict( + checkpoint['model'], strict=False) + self.model_regressor.eval() + + # export to onnx + a = torch.rand([1, 3, 224, 224]).to(device) + torch.onnx.export(self.model_regressor, a, + self.model_regressor_onnx_file, + output_names=['rotmat', 'shape', 'cam'], + opset_version=12) + print('an onnx saved into: ', self.model_regressor_onnx_file) def regress(self, img_original, body_bbox_list): """ @@ -68,79 +88,101 @@ def regress(self, img_original, body_bbox_list): pred_output_list.append(None) continue - with torch.no_grad(): - # model forward - pred_rotmat, pred_betas, pred_camera = self.model_regressor(norm_img.to(self.device)) - - #Convert rot_mat to aa since hands are always in aa - # pred_aa = rotmat3x3_to_angle_axis(pred_rotmat) - pred_aa = gu.rotation_matrix_to_angle_axis(pred_rotmat).cuda() - pred_aa = pred_aa.reshape(pred_aa.shape[0], 72) - smpl_output = self.smpl( - betas=pred_betas, - body_pose=pred_aa[:,3:], - global_orient=pred_aa[:,:3], - pose2rot=True) - pred_vertices = smpl_output.vertices - pred_joints_3d = smpl_output.joints - - pred_vertices = pred_vertices[0].cpu().numpy() - - pred_camera = pred_camera.cpu().numpy().ravel() - camScale = pred_camera[0] # *1.15 - camTrans = pred_camera[1:] - - pred_output = dict() - # Convert mesh to original image space (X,Y are aligned to image) - # 1. SMPL -> 2D bbox - # 2. 2D bbox -> original 2D image - pred_vertices_bbox = convert_smpl_to_bbox(pred_vertices, camScale, camTrans) - pred_vertices_img = convert_bbox_to_oriIm( - pred_vertices_bbox, boxScale_o2n, bboxTopLeft, img_original.shape[1], img_original.shape[0]) - - # Convert joint to original image space (X,Y are aligned to image) - pred_joints_3d = pred_joints_3d[0].cpu().numpy() # (1,49,3) - pred_joints_vis = pred_joints_3d[:,:3] # (49,3) - pred_joints_vis_bbox = convert_smpl_to_bbox(pred_joints_vis, camScale, camTrans) - pred_joints_vis_img = convert_bbox_to_oriIm( - pred_joints_vis_bbox, boxScale_o2n, bboxTopLeft, img_original.shape[1], img_original.shape[0]) - - # Output - pred_output['img_cropped'] = img[:, :, ::-1] - pred_output['pred_vertices_smpl'] = smpl_output.vertices[0].cpu().numpy() # SMPL vertex in original smpl space - pred_output['pred_vertices_img'] = pred_vertices_img # SMPL vertex in image space - pred_output['pred_joints_img'] = pred_joints_vis_img # SMPL joints in image space - - pred_aa_tensor = gu.rotation_matrix_to_angle_axis(pred_rotmat.detach().cpu()[0]) - pred_output['pred_body_pose'] = pred_aa_tensor.cpu().numpy().reshape(1, 72) # (1, 72) - - pred_output['pred_rotmat'] = pred_rotmat.detach().cpu().numpy() # (1, 24, 3, 3) - pred_output['pred_betas'] = pred_betas.detach().cpu().numpy() # (1, 10) - - pred_output['pred_camera'] = pred_camera - pred_output['bbox_top_left'] = bboxTopLeft - pred_output['bbox_scale_ratio'] = boxScale_o2n - pred_output['faces'] = self.smpl.faces - - if self.use_smplx: - img_center = np.array((img_original.shape[1], img_original.shape[0]) ) * 0.5 - # right hand - pred_joints = smpl_output.right_hand_joints[0].cpu().numpy() - pred_joints_bbox = convert_smpl_to_bbox(pred_joints, camScale, camTrans) - pred_joints_img = convert_bbox_to_oriIm( - pred_joints_bbox, boxScale_o2n, bboxTopLeft, img_original.shape[1], img_original.shape[0]) - pred_output['right_hand_joints_img_coord'] = pred_joints_img - # left hand - pred_joints = smpl_output.left_hand_joints[0].cpu().numpy() - pred_joints_bbox = convert_smpl_to_bbox(pred_joints, camScale, camTrans) - pred_joints_img = convert_bbox_to_oriIm( - pred_joints_bbox, boxScale_o2n, bboxTopLeft, img_original.shape[1], img_original.shape[0]) - pred_output['left_hand_joints_img_coord'] = pred_joints_img - - pred_output_list.append(pred_output) + # model forward + if self.onnx: + preds = self.model_regressor.infer(norm_img.cpu().numpy()) + pred_rotmat = preds['rotmat'] + pred_betas = preds['shape'] + pred_camera = preds['cam'] + else: + pred_rotmat, pred_betas, pred_camera = self.model_regressor( + norm_img.to(self.device)) + + pred_rotmat = torch.as_tensor(pred_rotmat).to(device) + pred_betas = torch.as_tensor(pred_betas).to(device) + pred_camera = torch.as_tensor(pred_camera).to(device) + + # Convert rot_mat to aa since hands are always in aa + # pred_aa = rotmat3x3_to_angle_axis(pred_rotmat) + # pred_aa = gu.rotation_matrix_to_angle_axis(pred_rotmat).cuda() + pred_aa = gu.rotation_matrix_to_angle_axis(pred_rotmat).to(device) + pred_aa = pred_aa.reshape(pred_aa.shape[0], 72) + smpl_output = self.smpl( + betas=pred_betas, + body_pose=pred_aa[:, 3:], + global_orient=pred_aa[:, :3], + pose2rot=True) + pred_vertices = smpl_output.vertices + pred_joints_3d = smpl_output.joints + + pred_vertices = pred_vertices[0].cpu().numpy() + + pred_camera = pred_camera.cpu().numpy().ravel() + camScale = pred_camera[0] # *1.15 + camTrans = pred_camera[1:] + + pred_output = dict() + # Convert mesh to original image space (X,Y are aligned to image) + # 1. SMPL -> 2D bbox + # 2. 2D bbox -> original 2D image + pred_vertices_bbox = convert_smpl_to_bbox( + pred_vertices, camScale, camTrans) + pred_vertices_img = convert_bbox_to_oriIm( + pred_vertices_bbox, boxScale_o2n, bboxTopLeft, img_original.shape[1], img_original.shape[0]) + + # Convert joint to original image space (X,Y are aligned to image) + pred_joints_3d = pred_joints_3d[0].cpu().numpy() # (1,49,3) + pred_joints_vis = pred_joints_3d[:, :3] # (49,3) + pred_joints_vis_bbox = convert_smpl_to_bbox( + pred_joints_vis, camScale, camTrans) + pred_joints_vis_img = convert_bbox_to_oriIm( + pred_joints_vis_bbox, boxScale_o2n, bboxTopLeft, img_original.shape[1], img_original.shape[0]) + + # Output + pred_output['img_cropped'] = img[:, :, ::-1] + # SMPL vertex in original smpl space + pred_output['pred_vertices_smpl'] = smpl_output.vertices[0].cpu().numpy() + # SMPL vertex in image space + pred_output['pred_vertices_img'] = pred_vertices_img + # SMPL joints in image space + pred_output['pred_joints_img'] = pred_joints_vis_img + + pred_aa_tensor = gu.rotation_matrix_to_angle_axis( + pred_rotmat.detach().cpu()[0]) + pred_output['pred_body_pose'] = pred_aa_tensor.cpu( + ).numpy().reshape(1, 72) # (1, 72) + + pred_output['pred_rotmat'] = pred_rotmat.detach( + ).cpu().numpy() # (1, 24, 3, 3) + pred_output['pred_betas'] = pred_betas.detach( + ).cpu().numpy() # (1, 10) + + pred_output['pred_camera'] = pred_camera + pred_output['bbox_top_left'] = bboxTopLeft + pred_output['bbox_scale_ratio'] = boxScale_o2n + pred_output['faces'] = self.smpl.faces + + if self.use_smplx: + img_center = np.array( + (img_original.shape[1], img_original.shape[0])) * 0.5 + # right hand + pred_joints = smpl_output.right_hand_joints[0].cpu().numpy() + pred_joints_bbox = convert_smpl_to_bbox( + pred_joints, camScale, camTrans) + pred_joints_img = convert_bbox_to_oriIm( + pred_joints_bbox, boxScale_o2n, bboxTopLeft, img_original.shape[1], img_original.shape[0]) + pred_output['right_hand_joints_img_coord'] = pred_joints_img + # left hand + pred_joints = smpl_output.left_hand_joints[0].cpu().numpy() + pred_joints_bbox = convert_smpl_to_bbox( + pred_joints, camScale, camTrans) + pred_joints_img = convert_bbox_to_oriIm( + pred_joints_bbox, boxScale_o2n, bboxTopLeft, img_original.shape[1], img_original.shape[0]) + pred_output['left_hand_joints_img_coord'] = pred_joints_img + + pred_output_list.append(pred_output) return pred_output_list - def get_hand_bboxes(self, pred_body_list, img_shape): """ @@ -153,8 +195,8 @@ def get_hand_bboxes(self, pred_body_list, img_shape): hand_bbox_list = list() for pred_body in pred_body_list: hand_bbox = dict( - left_hand = None, - right_hand = None + left_hand=None, + right_hand=None ) if pred_body is None: hand_bbox_list.append(hand_bbox) @@ -165,8 +207,10 @@ def get_hand_bboxes(self, pred_body_list, img_shape): if pred_joints_vis_img is not None: # get initial bbox - x0, x1 = np.min(pred_joints_vis_img[:, 0]), np.max(pred_joints_vis_img[:, 0]) - y0, y1 = np.min(pred_joints_vis_img[:, 1]), np.max(pred_joints_vis_img[:, 1]) + x0, x1 = np.min(pred_joints_vis_img[:, 0]), np.max( + pred_joints_vis_img[:, 0]) + y0, y1 = np.min(pred_joints_vis_img[:, 1]), np.max( + pred_joints_vis_img[:, 1]) width, height = x1-x0, y1-y0 # extend the obtained bbox margin = int(max(height, width) * 0.2) @@ -176,7 +220,8 @@ def get_hand_bboxes(self, pred_body_list, img_shape): x1 = min(x1 + margin, img_width) y1 = min(y1 + margin, img_height) # result bbox in (x0, y0, w, h) format - hand_bbox[hand_type] = np.array([x0, y0, x1-x0, y1-y0]) # in (x, y, w, h ) format + hand_bbox[hand_type] = np.array( + [x0, y0, x1-x0, y1-y0]) # in (x, y, w, h ) format hand_bbox_list.append(hand_bbox) diff --git a/bodymocap/models/__pycache__/__init__.cpython-39.pyc b/bodymocap/models/__pycache__/__init__.cpython-39.pyc new file mode 100644 index 0000000..daf0ec2 Binary files /dev/null and b/bodymocap/models/__pycache__/__init__.cpython-39.pyc differ diff --git a/bodymocap/models/__pycache__/hmr.cpython-39.pyc b/bodymocap/models/__pycache__/hmr.cpython-39.pyc new file mode 100644 index 0000000..bedaa9a Binary files /dev/null and b/bodymocap/models/__pycache__/hmr.cpython-39.pyc differ diff --git a/bodymocap/models/__pycache__/smpl.cpython-39.pyc b/bodymocap/models/__pycache__/smpl.cpython-39.pyc new file mode 100644 index 0000000..c4a3c25 Binary files /dev/null and b/bodymocap/models/__pycache__/smpl.cpython-39.pyc differ diff --git a/bodymocap/models/hmr.py b/bodymocap/models/hmr.py index 00c08ee..7bf93dc 100755 --- a/bodymocap/models/hmr.py +++ b/bodymocap/models/hmr.py @@ -152,7 +152,6 @@ def forward(self, x, init_pose=None, init_shape=None, init_cam=None, n_iter=3): pred_cam = self.deccam(xc) + pred_cam pred_rotmat = rot6d_to_rotmat(pred_pose).view(batch_size, 24, 3, 3) - return pred_rotmat, pred_shape, pred_cam def hmr(smpl_mean_params, pretrained=True, **kwargs): diff --git a/bodymocap/utils/__pycache__/geometry.cpython-39.pyc b/bodymocap/utils/__pycache__/geometry.cpython-39.pyc new file mode 100644 index 0000000..798a029 Binary files /dev/null and b/bodymocap/utils/__pycache__/geometry.cpython-39.pyc differ diff --git a/bodymocap/utils/__pycache__/imutils.cpython-39.pyc b/bodymocap/utils/__pycache__/imutils.cpython-39.pyc new file mode 100644 index 0000000..6e38420 Binary files /dev/null and b/bodymocap/utils/__pycache__/imutils.cpython-39.pyc differ diff --git a/bodymocap/utils/__pycache__/utils.cpython-39.pyc b/bodymocap/utils/__pycache__/utils.cpython-39.pyc new file mode 100644 index 0000000..4dfcc20 Binary files /dev/null and b/bodymocap/utils/__pycache__/utils.cpython-39.pyc differ diff --git a/bodymocap/utils/geometry.py b/bodymocap/utils/geometry.py index ba2dff8..1741b9f 100644 --- a/bodymocap/utils/geometry.py +++ b/bodymocap/utils/geometry.py @@ -86,7 +86,14 @@ def rot6d_to_rotmat(x): a2 = x[:, :, 1] b1 = F.normalize(a1) b2 = F.normalize(a2 - torch.einsum('bi,bi->b', b1, a2).unsqueeze(-1) * b1) - b3 = torch.cross(b1, b2) + if torch.onnx.is_in_onnx_export(): + # cross not support by onnx + b3 = torch.stack(( + b1[...,1]*b2[...,2] - b1[...,2]*b2[...,1], + b1[...,2]*b2[...,0] - b1[...,0]*b2[...,2], + b1[...,0]*b2[...,1] - b1[...,1]*b2[...,0]), dim=-1) + else: + b3 = torch.cross(b1, b2) return torch.stack((b1, b2, b3), dim=-1) diff --git a/bodymocap/utils/utils.py b/bodymocap/utils/utils.py new file mode 100644 index 0000000..53eda68 --- /dev/null +++ b/bodymocap/utils/utils.py @@ -0,0 +1,229 @@ + +import cv2 +import math +from operator import itemgetter +import numpy as np +try: + import onnxruntime +except ImportError: + onnxruntime = None + + +class ORTWrapper: + + def __init__(self, onnx_f) -> None: + self.onnx_f = onnx_f + so = onnxruntime.SessionOptions() + so.intra_op_num_threads = 6 + so.execution_mode = onnxruntime.ExecutionMode.ORT_PARALLEL + so.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL + self.ort_session = onnxruntime.InferenceSession( + onnx_f, sess_options=so) + + # todo: extract possible input hw for vision models + + def infer(self, imgs): + inputs = [imgs] + assert len(inputs) == len(self.ort_session.get_inputs() + ), 'inputs must same with model.' + ort_inputs = dict( + (self.ort_session.get_inputs()[i].name, inpt) for i, inpt in enumerate(inputs)) + ort_outs = self.ort_session.run(None, ort_inputs) + + outs_dict = dict() + for i, oo in enumerate(self.ort_session.get_outputs()): + n = oo.name + outs_dict[n] = ort_outs[i] + return outs_dict + + +BODY_PARTS_KPT_IDS = [[1, 2], [1, 5], [2, 3], [3, 4], [5, 6], [6, 7], [1, 8], [8, 9], [9, 10], [1, 11], + [11, 12], [12, 13], [1, 0], [0, 14], [14, 16], [0, 15], [15, 17], [2, 16], [5, 17]] +BODY_PARTS_PAF_IDS = ([12, 13], [20, 21], [14, 15], [16, 17], [22, 23], [24, 25], [0, 1], [2, 3], [4, 5], + [6, 7], [8, 9], [10, 11], [28, 29], [30, 31], [34, 35], [32, 33], [36, 37], [18, 19], [26, 27]) + + +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 + 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 connections_nms(a_idx, b_idx, affinity_scores): + # From all retrieved connections that share the same starting/ending keypoints leave only the top-scoring ones. + order = affinity_scores.argsort()[::-1] + affinity_scores = affinity_scores[order] + a_idx = a_idx[order] + b_idx = b_idx[order] + idx = [] + has_kpt_a = set() + has_kpt_b = set() + for t, (i, j) in enumerate(zip(a_idx, b_idx)): + if i not in has_kpt_a and j not in has_kpt_b: + idx.append(t) + has_kpt_a.add(i) + has_kpt_b.add(j) + idx = np.asarray(idx, dtype=np.int32) + return a_idx[idx], b_idx[idx], affinity_scores[idx] + + +def group_keypoints(all_keypoints_by_type, pafs, pose_entry_size=20, min_paf_score=0.05): + pose_entries = [] + all_keypoints = np.array( + [item for sublist in all_keypoints_by_type for item in sublist]) + points_per_limb = 10 + grid = np.arange(points_per_limb, dtype=np.float32).reshape(1, -1, 1) + all_keypoints_by_type = [np.array(keypoints, np.float32) + for keypoints in all_keypoints_by_type] + + for part_id in range(len(BODY_PARTS_PAF_IDS)): + part_pafs = pafs[:, :, BODY_PARTS_PAF_IDS[part_id]] + kpts_a = all_keypoints_by_type[BODY_PARTS_KPT_IDS[part_id][0]] + kpts_b = all_keypoints_by_type[BODY_PARTS_KPT_IDS[part_id][1]] + n = len(kpts_a) + m = len(kpts_b) + if n == 0 or m == 0: + continue + + # Get vectors between all pairs of keypoints, i.e. candidate limb vectors. + a = kpts_a[:, :2] + a = np.broadcast_to(a[None], (m, n, 2)) + b = kpts_b[:, :2] + vec_raw = (b[:, None, :] - a).reshape(-1, 1, 2) + + # Sample points along every candidate limb vector. + steps = (1 / (points_per_limb - 1) * vec_raw) + points = steps * grid + a.reshape(-1, 1, 2) + points = points.round().astype(dtype=np.int32) + x = points[..., 0].ravel() + y = points[..., 1].ravel() + + # Compute affinity score between candidate limb vectors and part affinity field. + field = part_pafs[y, x].reshape(-1, points_per_limb, 2) + vec_norm = np.linalg.norm(vec_raw, ord=2, axis=-1, keepdims=True) + vec = vec_raw / (vec_norm + 1e-6) + affinity_scores = (field * vec).sum(-1).reshape(-1, points_per_limb) + valid_affinity_scores = affinity_scores > min_paf_score + valid_num = valid_affinity_scores.sum(1) + affinity_scores = (affinity_scores * + valid_affinity_scores).sum(1) / (valid_num + 1e-6) + success_ratio = valid_num / points_per_limb + + # Get a list of limbs according to the obtained affinity score. + valid_limbs = np.where(np.logical_and( + affinity_scores > 0, success_ratio > 0.8))[0] + if len(valid_limbs) == 0: + continue + b_idx, a_idx = np.divmod(valid_limbs, n) + affinity_scores = affinity_scores[valid_limbs] + + # Suppress incompatible connections. + a_idx, b_idx, affinity_scores = connections_nms( + a_idx, b_idx, affinity_scores) + connections = list(zip(kpts_a[a_idx, 3].astype(np.int32), + kpts_b[b_idx, 3].astype(np.int32), + affinity_scores)) + if len(connections) == 0: + continue + + if part_id == 0: + pose_entries = [np.ones(pose_entry_size) + * -1 for _ in range(len(connections))] + for i in range(len(connections)): + pose_entries[i][BODY_PARTS_KPT_IDS[0][0]] = connections[i][0] + pose_entries[i][BODY_PARTS_KPT_IDS[0][1]] = connections[i][1] + pose_entries[i][-1] = 2 + pose_entries[i][-2] = np.sum( + all_keypoints[connections[i][0:2], 2]) + connections[i][2] + elif part_id == 17 or part_id == 18: + kpt_a_id = BODY_PARTS_KPT_IDS[part_id][0] + kpt_b_id = BODY_PARTS_KPT_IDS[part_id][1] + for i in range(len(connections)): + for j in range(len(pose_entries)): + if pose_entries[j][kpt_a_id] == connections[i][0] and pose_entries[j][kpt_b_id] == -1: + pose_entries[j][kpt_b_id] = connections[i][1] + elif pose_entries[j][kpt_b_id] == connections[i][1] and pose_entries[j][kpt_a_id] == -1: + pose_entries[j][kpt_a_id] = connections[i][0] + continue + else: + kpt_a_id = BODY_PARTS_KPT_IDS[part_id][0] + kpt_b_id = BODY_PARTS_KPT_IDS[part_id][1] + for i in range(len(connections)): + num = 0 + for j in range(len(pose_entries)): + if pose_entries[j][kpt_a_id] == connections[i][0]: + pose_entries[j][kpt_b_id] = connections[i][1] + num += 1 + pose_entries[j][-1] += 1 + pose_entries[j][-2] += all_keypoints[connections[i] + [1], 2] + connections[i][2] + if num == 0: + pose_entry = np.ones(pose_entry_size) * -1 + pose_entry[kpt_a_id] = connections[i][0] + pose_entry[kpt_b_id] = connections[i][1] + pose_entry[-1] = 2 + pose_entry[-2] = np.sum(all_keypoints[connections[i] + [0:2], 2]) + connections[i][2] + pose_entries.append(pose_entry) + + filtered_entries = [] + for i in range(len(pose_entries)): + if pose_entries[i][-1] < 3 or (pose_entries[i][-2] / pose_entries[i][-1] < 0.2): + continue + filtered_entries.append(pose_entries[i]) + pose_entries = np.asarray(filtered_entries) + return pose_entries, all_keypoints + + +def extract_keypoints(heatmap, all_keypoints, total_keypoint_num): + heatmap[heatmap < 0.1] = 0 + heatmap_with_borders = np.pad(heatmap, [(2, 2), (2, 2)], mode='constant') + heatmap_center = heatmap_with_borders[1:heatmap_with_borders.shape[0] - + 1, 1:heatmap_with_borders.shape[1]-1] + heatmap_left = heatmap_with_borders[1:heatmap_with_borders.shape[0] - + 1, 2:heatmap_with_borders.shape[1]] + heatmap_right = heatmap_with_borders[1:heatmap_with_borders.shape[0] - + 1, 0:heatmap_with_borders.shape[1]-2] + heatmap_up = heatmap_with_borders[2:heatmap_with_borders.shape[0], + 1:heatmap_with_borders.shape[1]-1] + heatmap_down = heatmap_with_borders[0:heatmap_with_borders.shape[0] - + 2, 1:heatmap_with_borders.shape[1]-1] + + heatmap_peaks = (heatmap_center > heatmap_left) &\ + (heatmap_center > heatmap_right) &\ + (heatmap_center > heatmap_up) &\ + (heatmap_center > heatmap_down) + heatmap_peaks = heatmap_peaks[1:heatmap_center.shape[0] - + 1, 1:heatmap_center.shape[1]-1] + keypoints = list(zip(np.nonzero(heatmap_peaks)[ + 1], np.nonzero(heatmap_peaks)[0])) # (w, h) + keypoints = sorted(keypoints, key=itemgetter(0)) + + suppressed = np.zeros(len(keypoints), np.uint8) + keypoints_with_score_and_id = [] + keypoint_num = 0 + for i in range(len(keypoints)): + if suppressed[i]: + continue + for j in range(i+1, len(keypoints)): + if math.sqrt((keypoints[i][0] - keypoints[j][0]) ** 2 + + (keypoints[i][1] - keypoints[j][1]) ** 2) < 6: + suppressed[j] = 1 + keypoint_with_score_and_id = (keypoints[i][0], keypoints[i][1], heatmap[keypoints[i][1], keypoints[i][0]], + total_keypoint_num + keypoint_num) + keypoints_with_score_and_id.append(keypoint_with_score_and_id) + keypoint_num += 1 + all_keypoints.append(keypoints_with_score_and_id) + return keypoint_num diff --git a/demo/__pycache__/demo_bodymocap.cpython-39.pyc b/demo/__pycache__/demo_bodymocap.cpython-39.pyc new file mode 100644 index 0000000..56b8d7e Binary files /dev/null and b/demo/__pycache__/demo_bodymocap.cpython-39.pyc differ diff --git a/demo/__pycache__/demo_bodymocap_ort.cpython-39.pyc b/demo/__pycache__/demo_bodymocap_ort.cpython-39.pyc new file mode 100644 index 0000000..76dac78 Binary files /dev/null and b/demo/__pycache__/demo_bodymocap_ort.cpython-39.pyc differ diff --git a/demo/__pycache__/demo_bodymocap_sim.cpython-39.pyc b/demo/__pycache__/demo_bodymocap_sim.cpython-39.pyc new file mode 100644 index 0000000..167b721 Binary files /dev/null and b/demo/__pycache__/demo_bodymocap_sim.cpython-39.pyc differ diff --git a/demo/__pycache__/demo_options.cpython-39.pyc b/demo/__pycache__/demo_options.cpython-39.pyc new file mode 100644 index 0000000..0c4448d Binary files /dev/null and b/demo/__pycache__/demo_options.cpython-39.pyc differ diff --git a/demo/demo_bodymocap.py b/demo/demo_bodymocap.py index 704b77f..2776b1b 100644 --- a/demo/demo_bodymocap.py +++ b/demo/demo_bodymocap.py @@ -109,6 +109,7 @@ def run_body_mocap(args, body_bbox_detector, body_mocap, visualizer): if args.single_person and len(body_bbox_list)>0: body_bbox_list = [body_bbox_list[0], ] + print('start regress.') # Body Pose Regression pred_output_list = body_mocap.regress(img_original_bgr, body_bbox_list) assert len(body_bbox_list) == len(pred_output_list) @@ -153,7 +154,7 @@ def main(): args = DemoOptions().parse() device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') - assert torch.cuda.is_available(), "Current version only supports GPU" + # assert torch.cuda.is_available(), "Current version only supports GPU" # Set bbox detector body_bbox_detector = BodyPoseEstimator() diff --git a/demo/demo_bodymocap_ort.py b/demo/demo_bodymocap_ort.py new file mode 100644 index 0000000..2b82975 --- /dev/null +++ b/demo/demo_bodymocap_ort.py @@ -0,0 +1,204 @@ +# Copyright (c) Facebook, Inc. and its affiliates. + +import os +from re import M +import sys +import os.path as osp +import torch +from torchvision.transforms import Normalize +import numpy as np +import cv2 +import argparse +import json +import pickle +from datetime import datetime + +from demo.demo_options import DemoOptions +from bodymocap.body_mocap_api import BodyMocap +from bodymocap.body_bbox_detector import BodyPoseEstimator +from bodymocap.body_bbox_detector_ort import LightweightedPoseDetector +import mocap_utils.demo_utils as demo_utils +import mocap_utils.general_utils as gnu +from mocap_utils.timer import Timer + +import renderer.image_utils as imu +from renderer.viewer2D import ImShow +from alfred.utils.timer import ATimer +import open3d as o3d +# from renderer.render import render +from realrender.render import render_human_mesh + + +def run_body_mocap(args, body_bbox_detector, body_mocap, visualizer): + # Setup input data to handle different types of inputs + input_type, input_data = demo_utils.setup_input(args) + + cur_frame = args.start_frame + video_frame = 0 + timer = Timer() + # vis = o3d.visualization.Visualizer() + # vis.create_window("Test", width=1280, height=900) + + while True: + timer.tic() + # load data + load_bbox = False + + if input_type == 'image_dir': + if cur_frame < len(input_data): + image_path = input_data[cur_frame] + img_original_bgr = cv2.imread(image_path) + else: + img_original_bgr = None + + elif input_type == 'bbox_dir': + if cur_frame < len(input_data): + print("Use pre-computed bounding boxes") + image_path = input_data[cur_frame]['image_path'] + hand_bbox_list = input_data[cur_frame]['hand_bbox_list'] + body_bbox_list = input_data[cur_frame]['body_bbox_list'] + img_original_bgr = cv2.imread(image_path) + load_bbox = True + else: + img_original_bgr = None + + elif input_type == 'video': + _, img_original_bgr = input_data.read() + if video_frame < cur_frame: + video_frame += 1 + continue + # save the obtained video frames + image_path = osp.join(args.out_dir, "frames", + f"{cur_frame:05d}.jpg") + if img_original_bgr is not None: + video_frame += 1 + if args.save_frame: + gnu.make_subdir(image_path) + cv2.imwrite(image_path, img_original_bgr) + + elif input_type == 'webcam': + _, img_original_bgr = input_data.read() + + if video_frame < cur_frame: + video_frame += 1 + continue + # save the obtained video frames + image_path = osp.join(args.out_dir, "frames", + f"scene_{cur_frame:05d}.jpg") + if img_original_bgr is not None: + video_frame += 1 + if args.save_frame: + gnu.make_subdir(image_path) + cv2.imwrite(image_path, img_original_bgr) + else: + assert False, "Unknown input_type" + + cur_frame += 1 + if img_original_bgr is None or cur_frame > args.end_frame: + break + print("--------------------------------------") + + if load_bbox: + body_pose_list = None + else: + with ATimer('pose'): + keypoints_2d = body_bbox_detector.run_one_img(img_original_bgr) + body_bbox_list = body_bbox_detector.get_enlarged_boxes_from_poses( + keypoints_2d, img_original_bgr.shape[0], img_original_bgr.shape[1]) + + hand_bbox_list = [None, ] * len(body_bbox_list) + + # save the obtained body & hand bbox to json file + if args.save_bbox_output: + demo_utils.save_info_to_json( + args, image_path, body_bbox_list, hand_bbox_list) + + if len(body_bbox_list) < 1: + print(f"No body deteced: {image_path}") + continue + + # Sort the bbox using bbox size + # (to make the order as consistent as possible without tracking) + bbox_size = [(x[2] * x[3]) for x in body_bbox_list] + idx_big2small = np.argsort(bbox_size)[::-1] + body_bbox_list = [body_bbox_list[i] for i in idx_big2small] + if args.single_person and len(body_bbox_list) > 0: + body_bbox_list = [body_bbox_list[0], ] + + # extract mesh for rendering (vertices in image space and faces) from pred_output_list + with ATimer('regress'): + print('start regress.') + # Body Pose Regression + pred_output_list = body_mocap.regress( + img_original_bgr, body_bbox_list) + assert len(body_bbox_list) == len(pred_output_list) + pred_mesh_list = demo_utils.extract_mesh_from_output( + pred_output_list) + + # visualization + with ATimer('vis'): + # sim3dr + pred_mesh_list = pred_mesh_list[0] + vertices = pred_mesh_list['vertices'] + vertices[:, 2] = - vertices[:, 2] + res_img = render_human_mesh(img_original_bgr, [ + vertices], pred_mesh_list['faces'], alpha=0.9, with_bg_flag=True) + + # vis.poll_events() + # vis.update_renderer() + + cv2.imshow('original', img_original_bgr) + cv2.imshow('res_img', res_img) + cv2.waitKey(1) + timer.toc(bPrint=True, title="Time") + + # vis.destroy_window() + # save images as a video + if not args.no_video_out and input_type in ['video', 'webcam']: + demo_utils.gen_video_out(args.out_dir, args.seq_name) + + if input_type == 'webcam' and input_data is not None: + input_data.release() + cv2.destroyAllWindows() + + +def mesh_list_to_o3d_mesh(pred_mesh_list): + res = [] + for i in pred_mesh_list: + m = o3d.geometry.TriangleMesh() + m.vertices = o3d.utility.Vector3dVector(i['vertices']) + m.triangles = o3d.utility.Vector3iVector(i['faces']) + res.append(m) + return res + + +def main(): + args = DemoOptions().parse() + + device = torch.device( + 'cuda') if torch.cuda.is_available() else torch.device('cpu') + # assert torch.cuda.is_available(), "Current version only supports GPU" + + # Set bbox detector + # body_bbox_detector = BodyPoseEstimator() + body_bbox_detector = LightweightedPoseDetector( + 'extra_data/body_module/human-pose-estimation.onnx') + + # Set mocap regressor + use_smplx = args.use_smplx + checkpoint_path = args.checkpoint_body_smplx if use_smplx else args.checkpoint_body_smpl + print("use_smplx", use_smplx) + body_mocap = BodyMocap(checkpoint_path, args.smpl_dir, device, use_smplx) + + # Set Visualizer + if args.renderer_type in ['pytorch3d', 'opendr']: + from renderer.screen_free_visualizer import Visualizer + else: + from renderer.visualizer import Visualizer + visualizer = Visualizer(args.renderer_type) + + run_body_mocap(args, body_bbox_detector, body_mocap, visualizer) + + +if __name__ == '__main__': + main() diff --git a/demo/demo_options.py b/demo/demo_options.py index fdf4132..406f229 100644 --- a/demo/demo_options.py +++ b/demo/demo_options.py @@ -53,7 +53,7 @@ def __init__(self): parser.add_argument('--frankmocap_fast_mode', action='store_true', help="Use fast hand detection mode for whole body motion capture (frankmocap)") # renderer - parser.add_argument("--renderer_type", type=str, default="opengl", + parser.add_argument("--renderer_type", type=str, default="pytorch3d", choices=['pytorch3d', 'opendr', 'opengl_gui', 'opengl'], help="type of renderer to use") self.parser = parser diff --git a/detectors/body_pose_estimator b/detectors/body_pose_estimator new file mode 160000 index 0000000..f2bd587 --- /dev/null +++ b/detectors/body_pose_estimator @@ -0,0 +1 @@ +Subproject commit f2bd5877ba13221145939638265f96715688013f diff --git a/detectors/hand_object_detector b/detectors/hand_object_detector new file mode 160000 index 0000000..27d9577 --- /dev/null +++ b/detectors/hand_object_detector @@ -0,0 +1 @@ +Subproject commit 27d9577c8e60b9cae3b8eefc0d23262af5e39bfd diff --git a/detectors/hand_only_detector b/detectors/hand_only_detector new file mode 160000 index 0000000..ad6f6b1 --- /dev/null +++ b/detectors/hand_only_detector @@ -0,0 +1 @@ +Subproject commit ad6f6b176cdcfa5510a9382175824402e7f3ecb6 diff --git a/docs/requirements.txt b/docs/requirements.txt index b372c21..4431668 100755 --- a/docs/requirements.txt +++ b/docs/requirements.txt @@ -21,3 +21,6 @@ smplx sklearn opendr chumpy +alfred-py +realrender + diff --git a/mocap_utils/__pycache__/coordconv.cpython-39.pyc b/mocap_utils/__pycache__/coordconv.cpython-39.pyc new file mode 100644 index 0000000..c992aab Binary files /dev/null and b/mocap_utils/__pycache__/coordconv.cpython-39.pyc differ diff --git a/mocap_utils/__pycache__/demo_utils.cpython-39.pyc b/mocap_utils/__pycache__/demo_utils.cpython-39.pyc new file mode 100644 index 0000000..fcfdbf4 Binary files /dev/null and b/mocap_utils/__pycache__/demo_utils.cpython-39.pyc differ diff --git a/mocap_utils/__pycache__/general_utils.cpython-39.pyc b/mocap_utils/__pycache__/general_utils.cpython-39.pyc new file mode 100644 index 0000000..6f537be Binary files /dev/null and b/mocap_utils/__pycache__/general_utils.cpython-39.pyc differ diff --git a/mocap_utils/__pycache__/geometry_utils.cpython-39.pyc b/mocap_utils/__pycache__/geometry_utils.cpython-39.pyc new file mode 100644 index 0000000..a73a673 Binary files /dev/null and b/mocap_utils/__pycache__/geometry_utils.cpython-39.pyc differ diff --git a/mocap_utils/__pycache__/geometry_utils_torch.cpython-39.pyc b/mocap_utils/__pycache__/geometry_utils_torch.cpython-39.pyc new file mode 100644 index 0000000..58181b4 Binary files /dev/null and b/mocap_utils/__pycache__/geometry_utils_torch.cpython-39.pyc differ diff --git a/mocap_utils/__pycache__/timer.cpython-39.pyc b/mocap_utils/__pycache__/timer.cpython-39.pyc new file mode 100644 index 0000000..a4f0086 Binary files /dev/null and b/mocap_utils/__pycache__/timer.cpython-39.pyc differ diff --git a/renderer/__pycache__/glRenderer.cpython-39.pyc b/renderer/__pycache__/glRenderer.cpython-39.pyc new file mode 100644 index 0000000..796d11a Binary files /dev/null and b/renderer/__pycache__/glRenderer.cpython-39.pyc differ diff --git a/renderer/__pycache__/glViewer.cpython-39.pyc b/renderer/__pycache__/glViewer.cpython-39.pyc new file mode 100644 index 0000000..d23383c Binary files /dev/null and b/renderer/__pycache__/glViewer.cpython-39.pyc differ diff --git a/renderer/__pycache__/image_utils.cpython-39.pyc b/renderer/__pycache__/image_utils.cpython-39.pyc new file mode 100644 index 0000000..5f91f5b Binary files /dev/null and b/renderer/__pycache__/image_utils.cpython-39.pyc differ diff --git a/renderer/__pycache__/meshRenderer.cpython-39.pyc b/renderer/__pycache__/meshRenderer.cpython-39.pyc new file mode 100644 index 0000000..3f7a2e4 Binary files /dev/null and b/renderer/__pycache__/meshRenderer.cpython-39.pyc differ diff --git a/renderer/__pycache__/od_renderer.cpython-39.pyc b/renderer/__pycache__/od_renderer.cpython-39.pyc new file mode 100644 index 0000000..c89c783 Binary files /dev/null and b/renderer/__pycache__/od_renderer.cpython-39.pyc differ diff --git a/renderer/__pycache__/p3d_renderer.cpython-39.pyc b/renderer/__pycache__/p3d_renderer.cpython-39.pyc new file mode 100644 index 0000000..f13fe08 Binary files /dev/null and b/renderer/__pycache__/p3d_renderer.cpython-39.pyc differ diff --git a/renderer/__pycache__/render.cpython-39.pyc b/renderer/__pycache__/render.cpython-39.pyc new file mode 100644 index 0000000..3798be3 Binary files /dev/null and b/renderer/__pycache__/render.cpython-39.pyc differ diff --git a/renderer/__pycache__/render_utils.cpython-39.pyc b/renderer/__pycache__/render_utils.cpython-39.pyc new file mode 100644 index 0000000..8cb1ce2 Binary files /dev/null and b/renderer/__pycache__/render_utils.cpython-39.pyc differ diff --git a/renderer/__pycache__/screen_free_visualizer.cpython-39.pyc b/renderer/__pycache__/screen_free_visualizer.cpython-39.pyc new file mode 100644 index 0000000..ee2bf08 Binary files /dev/null and b/renderer/__pycache__/screen_free_visualizer.cpython-39.pyc differ diff --git a/renderer/__pycache__/viewer2D.cpython-39.pyc b/renderer/__pycache__/viewer2D.cpython-39.pyc new file mode 100644 index 0000000..dab4d29 Binary files /dev/null and b/renderer/__pycache__/viewer2D.cpython-39.pyc differ diff --git a/renderer/__pycache__/visualizer.cpython-39.pyc b/renderer/__pycache__/visualizer.cpython-39.pyc new file mode 100644 index 0000000..ca5cc48 Binary files /dev/null and b/renderer/__pycache__/visualizer.cpython-39.pyc differ diff --git a/renderer/p3d_renderer.py b/renderer/p3d_renderer.py index 789f01f..56eca18 100644 --- a/renderer/p3d_renderer.py +++ b/renderer/p3d_renderer.py @@ -20,11 +20,14 @@ MeshRasterizer, SoftPhongShader, ) +from alfred.dl.torch.common import device + class Pytorch3dRenderer(object): def __init__(self, img_size, mesh_color): - self.device = torch.device("cuda:0") + # self.device = torch.device("cuda:0") + self.device = device # self.render_size = 1920 self.img_size = img_size @@ -149,8 +152,8 @@ def render(self, verts, faces, bg_img): verts[:, 2] /= 112 verts[:, 2] += 5 - verts_tensor = torch.from_numpy(verts).float().unsqueeze(0).cuda() - faces_tensor = torch.from_numpy(faces.copy()).long().unsqueeze(0).cuda() + verts_tensor = torch.from_numpy(verts).float().unsqueeze(0).to(self.device) + faces_tensor = torch.from_numpy(faces.copy()).long().unsqueeze(0).to(self.device) # set color mesh_color = self.mesh_color.repeat(1, verts.shape[0], 1) diff --git a/renderer/shaders/__pycache__/framework.cpython-39.pyc b/renderer/shaders/__pycache__/framework.cpython-39.pyc new file mode 100644 index 0000000..a0e6ee5 Binary files /dev/null and b/renderer/shaders/__pycache__/framework.cpython-39.pyc differ diff --git a/scripts/install_frankmocap.sh b/scripts/install_frankmocap.sh old mode 100644 new mode 100755