diff --git a/NLFPoseExtract/nlf_render.py b/NLFPoseExtract/nlf_render.py index 9abb43c..b0fe7e6 100644 --- a/NLFPoseExtract/nlf_render.py +++ b/NLFPoseExtract/nlf_render.py @@ -146,7 +146,7 @@ def collect_smpl_poses_samurai(data): -def render_nlf_as_images(smpl_poses, dw_poses, height, width, video_length, intrinsic_matrix=None, draw_2d=True): +def render_nlf_as_images(smpl_poses, dw_poses, height, width, video_length, intrinsic_matrix=None, draw_2d=True, draw_face=True, draw_hands=True): """ return a list of images """ base_colors_255_dict = { @@ -244,7 +244,7 @@ def render_nlf_as_images(smpl_poses, dw_poses, height, width, video_length, intr frames_np_rgba = render_whole(cylinder_specs_list, H=height, W=width, fx=focal_x, fy=focal_y, cx=princpt[0], cy=princpt[1]) if dw_poses is not None and draw_2d: - canvas_2d = draw_pose_to_canvas_np(aligned_poses, pool=None, H=height, W=width, reshape_scale=0, show_feet_flag=False, show_body_flag=False, show_cheek_flag=True, dw_hand=True) + canvas_2d = draw_pose_to_canvas_np(aligned_poses, pool=None, H=height, W=width, reshape_scale=0, show_feet_flag=False, show_body_flag=False, show_cheek_flag=True, dw_hand=True, show_face_flag=draw_face, show_hand_flag=draw_hands) for i in range(len(frames_np_rgba)): frame_img = frames_np_rgba[i] diff --git a/nodes.py b/nodes.py index a8b2078..4aa2cb0 100644 --- a/nodes.py +++ b/nodes.py @@ -164,6 +164,8 @@ class RenderNLFPoses: "optional": { "dw_poses": ("DWPOSES", {"default": None, "tooltip": "Optional DW pose model for 2D drawing"}), "ref_dw_pose": ("DWPOSES", {"default": None, "tooltip": "Optional reference DW pose model for alignment"}), + "draw_face": ("BOOLEAN", {"default": True, "tooltip": "Whether to draw face keypoints"}), + "draw_hands": ("BOOLEAN", {"default": True, "tooltip": "Whether to draw hand keypoints"}), } } @@ -172,7 +174,7 @@ class RenderNLFPoses: FUNCTION = "predict" CATEGORY = "WanVideoWrapper" - def predict(self, nlf_poses, width, height, dw_poses=None, ref_dw_pose=None): + def predict(self, nlf_poses, width, height, dw_poses=None, ref_dw_pose=None, draw_face=True, draw_hands=True): from .NLFPoseExtract.nlf_render import render_nlf_as_images, shift_dwpose_according_to_nlf, process_data_to_COCO_format, intrinsic_matrix_from_field_of_view from .NLFPoseExtract.align3d import solve_new_camera_params_central, solve_new_camera_params_down @@ -224,9 +226,11 @@ class RenderNLFPoses: logging.info(f"Scale - m: {scale_m}, face: {scale_face}") shift_dwpose_according_to_nlf(pose_input, dw_pose_input, ori_camera_pose, new_camera_intrinsics, height, width) - frames_np = render_nlf_as_images(pose_input, dw_pose_input, height, width, len(pose_input), intrinsic_matrix=new_camera_intrinsics) + intrinsic_matrix = new_camera_intrinsics else: - frames_np = render_nlf_as_images(pose_input, dw_pose_input, height, width, len(pose_input), intrinsic_matrix=ori_camera_pose) + intrinsic_matrix = ori_camera_pose + + frames_np = render_nlf_as_images(pose_input, dw_pose_input, height, width, len(pose_input), intrinsic_matrix=intrinsic_matrix, draw_face=draw_face, draw_hands=draw_hands) frames_tensor = torch.from_numpy(np.stack(frames_np, axis=0)).contiguous() / 255.0 frames_tensor, mask = frames_tensor[..., :3], frames_tensor[..., -1] > 0.5