From 3f26363748aaefa65b262f05f8afdecf3746b3db Mon Sep 17 00:00:00 2001 From: Hacker 17082006 Date: Mon, 1 Apr 2024 23:09:18 +0700 Subject: [PATCH] Add SMPL2OpenPose --- .../hmr2/utils/render_openpose.py | 9 ++-- nodes/human4d_nodes.py | 46 +++++++++++++++---- nodes/smpl_nodes.py | 18 ++++---- 3 files changed, 52 insertions(+), 21 deletions(-) diff --git a/motiondiff_modules/hmr2/utils/render_openpose.py b/motiondiff_modules/hmr2/utils/render_openpose.py index a4a26e3..2da4245 100644 --- a/motiondiff_modules/hmr2/utils/render_openpose.py +++ b/motiondiff_modules/hmr2/utils/render_openpose.py @@ -52,8 +52,7 @@ def render_keypoints(img: np.array, 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] + width, height = img.shape[1], img.shape[0] area = width * height lineType = 8 @@ -76,8 +75,8 @@ def render_keypoints(img: np.array, thicknessLineScaled = int(round(min(thicknessLine[index1], thicknessLine[index2]) * pose_scales[0])) colorIndex = index2 color = colors[colorIndex % numberColors] - keypoint1 = keypoints[index1, :-1].astype(np.int) - keypoint2 = keypoints[index2, :-1].astype(np.int) + 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 @@ -86,7 +85,7 @@ def render_keypoints(img: np.array, thicknessCircleScaled = int(round(thicknessCircle[faceIndex] * pose_scales[0])) colorIndex = part color = colors[colorIndex % numberColors] - center = keypoints[faceIndex, :-1].astype(np.int) + center = keypoints[faceIndex, :-1].astype(np.int32) cv2.circle(img, tuple(center.tolist()), radiusScaled, tuple(color.tolist()), thicknessCircleScaled, lineType, shift) return img diff --git a/nodes/human4d_nodes.py b/nodes/human4d_nodes.py index 5db4aa4..e4c3119 100644 --- a/nodes/human4d_nodes.py +++ b/nodes/human4d_nodes.py @@ -20,6 +20,8 @@ from comfy.model_management import get_torch_device from types import SimpleNamespace from torch.utils.data import DataLoader from ..md_config import get_smpl_models_dict +from motiondiff_modules.hmr2.utils.render_openpose import render_openpose as _render_openpose +from functools import partial smpl_models_dict = get_smpl_models_dict() @@ -28,7 +30,7 @@ class Humans4DLoader: def INPUT_TYPES(s): return { "required": { - "det_filename": ("STRING", {"default": "rtdetr-x.pt"}), + "det_filename": ("STRING", {"default": "yolov8x.pt"}), "fp16": ("BOOLEAN", {"default": False}) } } @@ -52,6 +54,33 @@ class Humans4DLoader: detector = detector_cls(str(Path(CACHE_DIR_4DHUMANS) / det_filename)) return (SimpleNamespace(human4d=model, model_cfg=model_cfg, detector=detector, fp16=fp16), ) +# kps_2d_frames: #List of [num_subjects, 44, 3] +def render_openpose(kps_2d_frames, boxes_frames, frame_width, frame_height): + openpose_frames = [] + for subjects_kps, xyxy_boxes_batch in zip(kps_2d_frames, boxes_frames): + print(xyxy_boxes_batch.shape) + canvas = np.zeros([frame_height, frame_width, 3], dtype=np.uint8) + subjects_kps = subjects_kps.numpy() # [num_subjects, 44, 3] + subjects_kps = np.concatenate((subjects_kps, np.ones_like(subjects_kps)[:, :, [0]]), axis=-1) + keypoint_matches = [(1, 12), (2, 8), (3, 7), (4, 6), (5, 9), (6, 10), (7, 11), (8, 14), (9, 2), (10, 1), (11, 0), (12, 3), (13, 4), (14, 5)] + for i in range(subjects_kps.shape[0]): + subject_xyxy_box = xyxy_boxes_batch[i] + print(subject_xyxy_box.shape) + x0, y0, x1, y1 = subject_xyxy_box.astype(np.int32) + _width, _height = x1-x0+1, y1-y0+1 + subjects_kps[i, :, 0] = _width * (subjects_kps[i, :, 0] + 0.5) + subjects_kps[i, :, 1] = _height * (subjects_kps[i, :, 1] + 0.5) + + pred_keypoints_img = np.zeros([_height, _width, 3], dtype=np.uint8) + body_keypoints = subjects_kps[i, :25] + extra_keypoints = subjects_kps[i, -19:] + for pair in keypoint_matches: + body_keypoints[pair[0], :] = extra_keypoints[pair[1], :] + pred_keypoints_img = _render_openpose(pred_keypoints_img, body_keypoints) + canvas[y0:y1+1, x0:x1+1, :] = pred_keypoints_img + openpose_frames.append(canvas) + return torch.from_numpy(np.stack(openpose_frames)) + class Human4D_Img2SMPL: @classmethod def INPUT_TYPES(s): @@ -113,16 +142,15 @@ class Human4D_Img2SMPL: scaled_focal_length = models.model_cfg.EXTRA.FOCAL_LENGTH / models.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() - # Render the result batch_size = batch['img'].shape[0] for n in range(batch_size): verts = out['pred_vertices'][n].detach().cpu() #Shape [num_verts, 3] cam_t = pred_cam_t_full[n] - kps_2d = out['pred_keypoints_2d'][n] + kps_2d = out['pred_keypoints_2d'][n].detach().cpu() #Shape [44, 3] _all_verts.append(verts) _all_cam_t.append(cam_t) _all_kps_2d.append(kps_2d) - + verts_frames.append( torch.stack(_all_verts) #Shape [num_subjects, num_verts, 3] ) @@ -130,15 +158,17 @@ class Human4D_Img2SMPL: torch.stack(_all_cam_t) #Shape [num_subjects, 3] ) kps_2d_frames.append( - torch.stack(_all_kps_2d) #Shape [num_subjects, 44, 2] + torch.stack(_all_kps_2d) #Shape [num_subjects, 44, 3] ) - verts_frames #List of [num_subjects, num_verts, 3] cam_t_frames #List of [num_frames, num_subjects, 3] - kps_2d_frames #List of [num_subjects, 44, 2] + kps_2d_frames #List of [num_subjects, 44, 3] return (( smpl_models_dict["SMPL_NEUTRAL.pkl"], verts_frames, - {"normalized_to_vertices": True, 'cam': cam_t_frames, "frame_width": img_size[0, 0], "frame_height": img_size[0, 1], "focal_length": scaled_focal_length, "keypoints_2d": kps_2d_frames} + {"normalized_to_vertices": True, 'cam': cam_t_frames, + "frame_width": int(img_size[0, 0].item()), "frame_height": int(img_size[0, 1].item()), + "focal_length": scaled_focal_length, + "render_openpose": partial(render_openpose, kps_2d_frames, boxes_images, int(img_size[0, 0].item()), int(img_size[0, 1].item()))} # In Comfy, IMAGE is a batched Tensor so all frames always share the same size ), ) diff --git a/nodes/smpl_nodes.py b/nodes/smpl_nodes.py index e4ad3f9..d760bcd 100644 --- a/nodes/smpl_nodes.py +++ b/nodes/smpl_nodes.py @@ -323,7 +323,6 @@ class SMPLShapeParameters: smpl[2]["shape_parameters"] = shape_parameters return (smpl,) -""" class Render_OpenPose_From_SMPL_Mesh_Multiple_Subjects: @classmethod def INPUT_TYPES(s): @@ -334,12 +333,13 @@ class Render_OpenPose_From_SMPL_Mesh_Multiple_Subjects: } RETURN_TYPES = ("IMAGE",) CATEGORY = "MotionDiff/smpl" - + FUNCTION = "render" def render(self, smpl_multi_subjects): - meta = smpl_multi_subjects[2] - kps_2d_frames = meta['keypoints_2d'] - -""" + render_openpose = smpl_multi_subjects[2].get("render_openpose", None) + if render_openpose is None: + raise NotImplementedError("render_openpose") + return (render_openpose().float() / 255., ) + NODE_CLASS_MAPPINGS = { "SmplifyMotionData": SmplifyMotionData, @@ -348,7 +348,8 @@ NODE_CLASS_MAPPINGS = { "SaveSMPL": SaveSMPL, "ExportSMPLTo3DSoftware": ExportSMPLTo3DSoftware, "SMPLShapeParameters": SMPLShapeParameters, - "RenderMultipleSubjectsSMPLMesh": RenderMultipleSubjectsSMPLMesh + "RenderMultipleSubjectsSMPLMesh": RenderMultipleSubjectsSMPLMesh, + "Render_OpenPose_From_SMPL_Mesh_Multiple_Subjects": Render_OpenPose_From_SMPL_Mesh_Multiple_Subjects } NODE_DISPLAY_NAME_MAPPINGS = { @@ -358,5 +359,6 @@ NODE_DISPLAY_NAME_MAPPINGS = { "SaveSMPL": "Save SMPL", "ExportSMPLTo3DSoftware": "Export SMPL to 3DCGI Software", "SMPLShapeParameters": "SMPL Shape Parameters", - "RenderMultipleSubjectsSMPLMesh": "Render Mutiple Subjects from SMPL Mesh" + "RenderMultipleSubjectsSMPLMesh": "Render Mutiple Subjects from SMPL Mesh", + "Render_OpenPose_From_SMPL_Mesh_Multiple_Subjects": "Render Multiple OpenPose from SMPL Mesh" } \ No newline at end of file