From 2dcb704e0b9c331cd03a7be37ddd42849d9da4ee Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=8D=8E=E6=88=90=E7=82=80?= <49787377+350800786@users.noreply.github.com> Date: Thu, 20 Feb 2025 19:40:18 +0800 Subject: [PATCH] 0.0.2 --- __init__.py | 171 ++++++++++++++++++++++++---- utils/dwpose/dwpose_detector.py | 58 ++++++++++ utils/dwpose/skeleton_extraction.py | 133 ++++++++++++++++++++++ utils/dwpose/wholebody.py | 54 +++++++++ utils/image_utils.py | 43 ++++--- 5 files changed, 421 insertions(+), 38 deletions(-) create mode 100644 utils/dwpose/dwpose_detector.py create mode 100644 utils/dwpose/skeleton_extraction.py create mode 100644 utils/dwpose/wholebody.py diff --git a/__init__.py b/__init__.py index 9b08263..62e0191 100644 --- a/__init__.py +++ b/__init__.py @@ -1,7 +1,6 @@ import os import torch import numpy as np -import time from diffusers.models.attention_processor import XFormersAttnProcessor from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection @@ -18,7 +17,12 @@ from .StableAnimator.animation.modules.pose_net import PoseNet from .StableAnimator.animation.modules.unet import UNetSpatioTemporalConditionModel from .StableAnimator.animation.pipelines.inference_pipeline_animation import InferenceAnimationPipeline -from .utils.image_utils import tensor_to_pil, tensor_to_np, np_to_tensor, load_images_from_folder, save_frames_as_mp4 +from .utils.image_utils import tensor_to_pil, tensor_to_np, np_to_tensor, load_images_from_folder + + +# from .StableAnimator.DWPose.dwpose_utils.dwpose_detector import dwpose_detector_aligned +from .utils.dwpose.dwpose_detector import DWposeDetectorAligned # 这里重写StableAnimator.DWPose.dwpose_utils.dwpose_detector.dwpose_detector_aligned 函数，以适应ComfyUI的模型加载机制 +from .utils.dwpose.skeleton_extraction import draw_pose GLOBAL_CATEGORY = "HJH_StableAnimatorNode🪅" @@ -189,7 +193,8 @@ class StableAnimatorNode: "required": { "stable_animator_models":("STABLEANIMATORMODELS",), "reference_image": ("IMAGE",), - "pose_images_dir": ("STRING",{"default":"",}), # 输入为图像序列文件夹路径 + # "pose_images_dir": ("STRING",{"default":"",}), # 输入为图像序列文件夹路径 + "pose_images": ("IMAGE",), "format":(["512x512","576x1024"],{}), "tile_size":("INT",{"default":16,"min":16,"max":64,"step":16}), "frames_overlap":("INT",{"default":4,"min":4,"max":64,"step":4}), @@ -231,8 +236,8 @@ class StableAnimatorNode: FUNCTION = "generate" CATEGORY = GLOBAL_CATEGORY - RETURN_TYPES = ("STRING",) - RETURN_NAMES = ("video_path",) + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("video_frames",) # OUTPUT_IS_LIST = (True,) def generate(self, **kwargs): @@ -244,8 +249,13 @@ class StableAnimatorNode: models = kwargs["stable_animator_models"] reference_image = tensor_to_pil(kwargs["reference_image"]) reference_image_np = tensor_to_np(kwargs["reference_image"]) - pose_images_dir = kwargs["pose_images_dir"] - pose_images = load_images_from_folder(pose_images_dir, width=width, height=height) + # pose_images_dir = kwargs["pose_images_dir"] + # pose_images = load_images_from_folder(pose_images_dir, width=width, height=height) + tensor_pose_images = kwargs["pose_images"] + pose_images = [] + for pose_image in tensor_pose_images: + pose_images.append(tensor_to_pil(pose_image.unsqueeze(0)).resize((width,height))) + seed = kwargs["seed"] if kwargs["seed"] != -1 else torch.randint(0, 2**32-1, (1,)).item() generator = torch.Generator(device=self.device).manual_seed(seed) num_frames = len(pose_images) @@ -293,7 +303,7 @@ class StableAnimatorNode: tile_overlap=kwargs["frames_overlap"], decode_chunk_size=kwargs["decode_chunk_size"], motion_bucket_id=127., - fps=7, + fps=kwargs["fps"], min_guidance_scale=kwargs["guidance_scale"], max_guidance_scale=kwargs["guidance_scale"], noise_aug_strength=kwargs["noise_aug_strength"], @@ -305,28 +315,141 @@ class StableAnimatorNode: # 转换为张量输出 output_tensors = [] - np_frames=np.array(video_frames) - print("********************") - print(np_frames.shape) - print(np_frames[0].shape) - # for frame in video_frames: - # np_frames = np.array(frame) - # output_tensors.append(np_to_tensor(np.array(frame))) + for frame in video_frames: + output_tensors.append(np_to_tensor(np.array(frame))) - output_root = folder_paths.get_output_directory() - output_dir = os.path.join(output_root, "stable_animator_output") - os.makedirs(output_dir, exist_ok=True) - output_file = os.path.join(output_dir, f"{time.time()}.mp4") - save_frames_as_mp4(np_frames, output_file, fps=kwargs["fps"]) - return (output_file,) - # return (output_tensors,) + return torch.cat(output_tensors, dim=0), + + +class StableAnimatorDWPoseDetectorAlignedModels: + """ + StableAnimator加载DWPose模型, + """ + def __init__(self): + pass + @classmethod + def INPUT_TYPES(cls): + """定义输入参数""" + return { + "required": { + }, + } + + FUNCTION = "load" + CATEGORY = GLOBAL_CATEGORY + RETURN_TYPES = ("DWPOSEDETECTORALIGNED",) + RETURN_NAMES = ("dwpose_detector_aligned",) + + def load(self, ): + """执行模型加载""" + return DWposeDetectorAligned(), + +class StableAnimatorSkeletonNode: + """ + StableAnimator生成视频POSE骨架 + """ + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + """定义输入参数""" + return { + "required": { + "dwpose_detector_aligned":("DWPOSEDETECTORALIGNED",), + "reference_image":("IMAGE",), + "video_frames":("IMAGE",), + }, + } + + FUNCTION = "extraction" + CATEGORY = GLOBAL_CATEGORY + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("pose_frames",) + + def extraction(self, **kwargs): + ''' + 以下代码参考 StableAnimator.DWPose.dwpose_utils.dwpose_detector.get_video_pose, 以适应ComfyUI的参数输入 + ''' + dwpose_detector_aligned = kwargs["dwpose_detector_aligned"] + ref_image = tensor_to_np(kwargs["reference_image"]) + height, width, _ = ref_image.shape + ref_pose = dwpose_detector_aligned(ref_image) + ref_keypoint_id = [0, 1, 2, 5, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17] + ref_keypoint_id = [i for i in ref_keypoint_id \ + if len(ref_pose['bodies']['subset']) > 0 and ref_pose['bodies']['subset'][0][i] >= .0] + ref_body = ref_pose['bodies']['candidate'][ref_keypoint_id] + + video_frames = kwargs["video_frames"] + + detected_poses = [] + for video_frame in video_frames: + frame = tensor_to_np(video_frame.unsqueeze(0)) + pose = dwpose_detector_aligned(frame) + detected_poses.append(pose) + + detected_bodies = np.stack([p['bodies']['candidate'] for p in detected_poses if p['bodies']['candidate'].shape[0] == 18])[:,ref_keypoint_id] + ay, by = np.polyfit(detected_bodies[:, :, 1].flatten(), np.tile(ref_body[:, 1], len(detected_bodies)), 1) + fh = height + fw = width + ax = ay / (fh / fw / height * width) + bx = np.mean(np.tile(ref_body[:, 0], len(detected_bodies)) - detected_bodies[:, :, 0].flatten() * ax) + a = np.array([ax, ay]) + b = np.array([bx, by]) + + output_pose = [] + for detected_pose in detected_poses: + detected_pose['bodies']['candidate'] = detected_pose['bodies']['candidate'] * a + b + detected_pose['faces'] = detected_pose['faces'] * a + b + detected_pose['hands'] = detected_pose['hands'] * a + b + im = draw_pose(detected_pose, height, width) + output_pose.append(np_to_tensor(np.array(im))) + + output = torch.cat(output_pose, dim=0) + return output, + +class StableAnimatorLoadFramesFromFolderNode: + """ + StableAnimator从文件夹加载帧 + """ + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + """定义输入参数""" + return { + "required": { + "pose_images_folder": ("STRING",{"default":"",}), # 输入为图像序列文件夹路径 + }, + } + + FUNCTION = "load" + CATEGORY = GLOBAL_CATEGORY + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("frames",) + + def load(self, **kwargs): + frames = [] + folder = kwargs["pose_images_folder"] + frames = load_images_from_folder(folder,) + + return frames, NODE_CLASS_MAPPINGS = { "StableAnimatorNode": StableAnimatorNode, - "StableAnimatorModels": StableAnimatorModels + "StableAnimatorModels": StableAnimatorModels, + "StableAnimatorSkeletonNode": StableAnimatorSkeletonNode, + "StableAnimatorDWPoseDetectorAlignedModels": StableAnimatorDWPoseDetectorAlignedModels, + "StableAnimatorLoadFramesFromFolderNode": StableAnimatorLoadFramesFromFolderNode, } NODE_DISPLAY_NAME_MAPPINGS = { "StableAnimatorNode": "HJH-StableAnimator - Video Generation", - "StableAnimatorModels":"HJH-StableAnimator - Load Models" + "StableAnimatorModels":"HJH-StableAnimator - Load Models", + "StableAnimatorSkeletonNode": "HJH-StableAnimator - Pose Extraction", + "StableAnimatorDWPoseDetectorAlignedModels": "HJH-StableAnimator - Load DWPose Models", + "StableAnimatorLoadFramesFromFolderNode": "HJH-StableAnimator - Load Pose Frames From Folder", } \ No newline at end of file diff --git a/utils/dwpose/dwpose_detector.py b/utils/dwpose/dwpose_detector.py new file mode 100644 index 0000000..b4e5b26 --- /dev/null +++ b/utils/dwpose/dwpose_detector.py @@ -0,0 +1,58 @@ +import os + +import numpy as np +import torch + +from .wholebody import Wholebody + +os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" +# device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + +class DWposeDetectorAligned: + def __init__(self,): + self.pose_estimation = Wholebody() + + + def release_memory(self): + if hasattr(self, 'pose_estimation'): + del self.pose_estimation + import gc; gc.collect() + + def __call__(self, oriImg): + oriImg = oriImg.copy() + H, W, C = oriImg.shape + with torch.no_grad(): + candidate, score = self.pose_estimation(oriImg) + nums, _, locs = candidate.shape + candidate[..., 0] /= float(W) + candidate[..., 1] /= float(H) + body = candidate[:, :18].copy() + body = body.reshape(nums * 18, locs) + subset = score[:, :18].copy() + for i in range(len(subset)): + for j in range(len(subset[i])): + if subset[i][j] > 0.3: + subset[i][j] = int(18 * i + j) + else: + subset[i][j] = -1 + + # un_visible = subset < 0.3 + # candidate[un_visible] = -1 + + # foot = candidate[:, 18:24] + + faces = candidate[:, 24:92] + + hands = candidate[:, 92:113] + hands = np.vstack([hands, candidate[:, 113:]]) + + faces_score = score[:, 24:92] + hands_score = np.vstack([score[:, 92:113], score[:, 113:]]) + + bodies = dict(candidate=body, subset=subset, score=score[:, :18]) + pose = dict(bodies=bodies, hands=hands, hands_score=hands_score, faces=faces, faces_score=faces_score) + + return pose + + +# dwpose_detector_aligned = DWposeDetectorAligned(device=device) \ No newline at end of file diff --git a/utils/dwpose/skeleton_extraction.py b/utils/dwpose/skeleton_extraction.py new file mode 100644 index 0000000..281d7f9 --- /dev/null +++ b/utils/dwpose/skeleton_extraction.py @@ -0,0 +1,133 @@ +import math +import matplotlib +import cv2 +import numpy as np + +eps = 0.01 + +def alpha_blend_color(color, alpha): + """blend color according to point conf + """ + return [int(c * alpha) for c in color] + +def draw_bodypose(canvas, candidate, subset, score): + H, W, C = canvas.shape + candidate = np.array(candidate) + subset = np.array(subset) + + stickwidth = 4 + + limbSeq = [[2, 3], [2, 6], [3, 4], [4, 5], [6, 7], [7, 8], [2, 9], [9, 10], \ + [10, 11], [2, 12], [12, 13], [13, 14], [2, 1], [1, 15], [15, 17], \ + [1, 16], [16, 18], [3, 17], [6, 18]] + + colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0], \ + [0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255], \ + [170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85]] + + for i in range(17): + for n in range(len(subset)): + index = subset[n][np.array(limbSeq[i]) - 1] + conf = score[n][np.array(limbSeq[i]) - 1] + if conf[0] < 0.3 or conf[1] < 0.3: + continue + Y = candidate[index.astype(int), 0] * float(W) + X = candidate[index.astype(int), 1] * float(H) + mX = np.mean(X) + mY = np.mean(Y) + length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5 + angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1])) + polygon = cv2.ellipse2Poly((int(mY), int(mX)), (int(length / 2), stickwidth), int(angle), 0, 360, 1) + cv2.fillConvexPoly(canvas, polygon, alpha_blend_color(colors[i], conf[0] * conf[1])) + + canvas = (canvas * 0.6).astype(np.uint8) + + for i in range(18): + for n in range(len(subset)): + index = int(subset[n][i]) + if index == -1: + continue + x, y = candidate[index][0:2] + conf = score[n][i] + x = int(x * W) + y = int(y * H) + cv2.circle(canvas, (int(x), int(y)), 4, alpha_blend_color(colors[i], conf), thickness=-1) + + return canvas + +def draw_handpose(canvas, all_hand_peaks, all_hand_scores): + H, W, C = canvas.shape + + edges = [[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]] + + for peaks, scores in zip(all_hand_peaks, all_hand_scores): + + for ie, e in enumerate(edges): + x1, y1 = peaks[e[0]] + x2, y2 = peaks[e[1]] + x1 = int(x1 * W) + y1 = int(y1 * H) + x2 = int(x2 * W) + y2 = int(y2 * H) + score = int(scores[e[0]] * scores[e[1]] * 255) + if x1 > eps and y1 > eps and x2 > eps and y2 > eps: + cv2.line(canvas, (x1, y1), (x2, y2), + matplotlib.colors.hsv_to_rgb([ie / float(len(edges)), 1.0, 1.0]) * score, thickness=2) + + for i, keyponit in enumerate(peaks): + x, y = keyponit + x = int(x * W) + y = int(y * H) + score = int(scores[i] * 255) + if x > eps and y > eps: + cv2.circle(canvas, (x, y), 4, (0, 0, score), thickness=-1) + return canvas + +def draw_facepose(canvas, all_lmks, all_scores): + H, W, C = canvas.shape + for lmks, scores in zip(all_lmks, all_scores): + for lmk, score in zip(lmks, scores): + x, y = lmk + x = int(x * W) + y = int(y * H) + conf = int(score * 255) + if x > eps and y > eps: + cv2.circle(canvas, (x, y), 3, (conf, conf, conf), thickness=-1) + return canvas + +def draw_pose(pose, H, W, ref_w=2160): + """vis dwpose outputs + + Args: + pose (List): DWposeDetector outputs in dwpose_detector.py + H (int): height + W (int): width + ref_w (int, optional) Defaults to 2160. + + Returns: + np.ndarray: image pixel value in RGB mode + """ + bodies = pose['bodies'] + faces = pose['faces'] + hands = pose['hands'] + candidate = bodies['candidate'] + subset = bodies['subset'] + + sz = min(H, W) + sr = (ref_w / sz) if sz != ref_w else 1 + + ########################################## create zero canvas ################################################## + canvas = np.zeros(shape=(int(H*sr), int(W*sr), 3), dtype=np.uint8) + + ########################################### draw body pose ##################################################### + canvas = draw_bodypose(canvas, candidate, subset, score=bodies['score']) + + ########################################### draw hand pose ##################################################### + canvas = draw_handpose(canvas, hands, pose['hands_score']) + + ########################################### draw face pose ##################################################### + canvas = draw_facepose(canvas, faces, pose['faces_score']) + + # return cv2.cvtColor(cv2.resize(canvas, (W, H)), cv2.COLOR_BGR2RGB).transpose(2, 0, 1) + return cv2.cvtColor(cv2.resize(canvas, (W, H)), cv2.COLOR_BGR2RGB) \ No newline at end of file diff --git a/utils/dwpose/wholebody.py b/utils/dwpose/wholebody.py new file mode 100644 index 0000000..ccb439c --- /dev/null +++ b/utils/dwpose/wholebody.py @@ -0,0 +1,54 @@ +import numpy as np + +import onnxruntime as ort +from ...StableAnimator.DWPose.dwpose_utils.onnxdet import inference_detector +from ...StableAnimator.DWPose.dwpose_utils.onnxpose import inference_pose + +import folder_paths + +class Wholebody: + def __init__(self): + + root_dir = folder_paths.get_folder_paths("stable_animator")[0] + + device = 'cuda:0' + # device = 'cuda' + providers = ['CPUExecutionProvider' + ] if device == 'cpu' else ['CUDAExecutionProvider'] + onnx_det = f'{root_dir}/DWPose/yolox_l.onnx' + onnx_pose = f'{root_dir}/DWPose/dw-ll_ucoco_384.onnx' + + self.session_det = ort.InferenceSession(path_or_bytes=onnx_det, providers=providers) + self.session_pose = ort.InferenceSession(path_or_bytes=onnx_pose, providers=providers) + + def __call__(self, oriImg): + det_result = inference_detector(self.session_det, oriImg) + keypoints, scores = inference_pose(self.session_pose, det_result, oriImg) + + keypoints_info = np.concatenate( + (keypoints, scores[..., None]), axis=-1) + # compute neck joint + neck = np.mean(keypoints_info[:, [5, 6]], axis=1) + # neck score when visualizing pred + neck[:, 2:4] = np.logical_and( + keypoints_info[:, 5, 2:4] > 0.3, + keypoints_info[:, 6, 2:4] > 0.3).astype(int) + new_keypoints_info = np.insert( + keypoints_info, 17, neck, axis=1) + mmpose_idx = [ + 17, 6, 8, 10, 7, 9, 12, 14, 16, 13, 15, 2, 1, 4, 3 + ] + openpose_idx = [ + 1, 2, 3, 4, 6, 7, 8, 9, 10, 12, 13, 14, 15, 16, 17 + ] + new_keypoints_info[:, openpose_idx] = \ + new_keypoints_info[:, mmpose_idx] + keypoints_info = new_keypoints_info + + keypoints, scores = keypoints_info[ + ..., :2], keypoints_info[..., 2] + + return keypoints, scores + + + diff --git a/utils/image_utils.py b/utils/image_utils.py index 2a8d280..2ca01f3 100644 --- a/utils/image_utils.py +++ b/utils/image_utils.py @@ -6,6 +6,7 @@ import cv2 from PIL import ImageOps, Image + def pil_to_tensor(image_pil): i = ImageOps.exif_transpose(image_pil) image = i.convert("RGB") @@ -30,25 +31,39 @@ def np_to_tensor(image_np): image_tensor = torch.from_numpy(image)[None,] return image_tensor -def load_images_from_folder(folder, width, height): +# def load_images_from_folder(folder, width, height): +# """从目录读取图片""" +# images = [] +# files = os.listdir(folder) +# png_files = [f for f in files if f.endswith('.png')] +# png_files.sort(key=lambda x: int(x.split('_')[1].split('.')[0])) +# for filename in png_files: +# img = Image.open(os.path.join(folder, filename)).convert('RGB') +# img = img.resize((width, height)) +# images.append(img) + +# return images + +def load_images_from_folder(folder, ): """从目录读取图片""" images = [] files = os.listdir(folder) png_files = [f for f in files if f.endswith('.png')] - png_files.sort(key=lambda x: int(x.split('_')[1].split('.')[0])) + # png_files.sort(key=lambda x: int(x.split('_')[1].split('.')[0])) for filename in png_files: img = Image.open(os.path.join(folder, filename)).convert('RGB') - img = img.resize((width, height)) - images.append(img) + # img = img.resize((width, height)) + images.append(pil_to_tensor(img)) - return images + return torch.cat(images, dim=0) -def save_frames_as_mp4(frames, output_mp4_path, fps): - print("Starting saving the frames as mp4") - height, width, _ = frames[0].shape - fourcc = cv2.VideoWriter_fourcc(*'mp4v') # 'H264' for better quality - out = cv2.VideoWriter(output_mp4_path, fourcc, float(fps), (width, height)) - for frame in frames: - frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR) - out.write(frame_bgr) - out.release() \ No newline at end of file +# def save_frames_as_mp4(frames, output_mp4_path, fps): +# print("Starting saving the frames as mp4") +# height, width, _ = frames[0].shape +# # fourcc = cv2.VideoWriter_fourcc(*'mp4v') # 'H264' for better quality +# fourcc = cv2.VideoWriter_fourcc(*'avc1') +# out = cv2.VideoWriter(output_mp4_path, fourcc, float(fps), (width, height)) +# for frame in frames: +# frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR) +# out.write(frame_bgr) +# out.release() \ No newline at end of file