import torch import comfy.utils import cv2 import numpy as np import folder_paths import os import math connect_color = [ [ 0, 0, 255], [255, 0, 0], [255, 170, 0], [255, 255, 0], [255, 85, 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], [ 85, 0, 255], [170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85] ] for i, (R, G, B) in enumerate(connect_color): connect_color[i] = [B,G,R] # 骨架连接的关节对 skeleton = [ [0, 1], [1, 2], [2, 3], [3, 4], [1, 5], [5, 6], [6, 7], [1, 8], [8, 9], [9, 10], [1, 11], [11, 12], [12, 13], [14, 0], [14, 16], [15, 0], [15, 17] ] bonesPoints = [ [1,0],# neck [0,14 ],#right eye [0,15 ],#left eye [14,16 ],# right ear [15,17 ],# left ear [1,2 ],#right shoulder [1,5 ],#left shoulder [2,3 ],#right elbow [5,6 ],#left elbow [3,4 ],#right wrist [6,7 ],#left wrist [1,8 ],#right hip [1,11 ],#left hip [8,9 ],#right knee [11,12 ],# left knee [9,10 ],#right ankle [12,13 ] # left ankle ] class HumanPart: def __init__(self, index): self.index = index self.abs_x = 0.0 #绝对坐标 self.abs_y = 0.0 # 绝对坐标 self.length = 0.0 #长度 self.score=1.0 #置信度 self.abs_angle=0.0 #绝对角度 self.ref_angle=0.0 #相对角度 self.children = [] #子物体 self.ref_scale=1.0 self.abs_scale=1.0 self.parent = None #父物体 def SetParent(self, parent): self.parent = parent parent.children.append(self) #根据相对角度和长度计算绝对坐标 def Update(self): # 如果有父节点，则更新绝对位置和角度 if self.parent is not None: if self.score > 0.0: # 更新绝对缩放 self.abs_scale = self.parent.abs_scale * self.ref_scale # 更新绝对角度 self.abs_angle = self.parent.abs_angle + self.ref_angle # 更新绝对x坐标 self.abs_x = self.parent.abs_x + self.length * math.cos(self.abs_angle)* self.abs_scale # 更新绝对y坐标 self.abs_y = self.parent.abs_y + self.length * math.sin(self.abs_angle)* self.abs_scale else: # 如果没有父节点，则绝对角度和缩放等于参考角度和缩放 self.abs_angle = self.ref_angle self.abs_scale = self.ref_scale #self.score=1.0 # 遍历子节点，递归调用Update方法 for child in self.children: child.Update() #根据绝对位置，计算角度、相对位置和相对角度 def CalculateRef(self): # 如果有父节点，则更新相对位置和角度 if self.parent is not None: if self.score > 0.0: self.length = math.sqrt((self.abs_x - self.parent.abs_x)**2 + (self.abs_y - self.parent.abs_y)**2) self.abs_angle =math.atan2(self.abs_y - self.parent.abs_y, self.abs_x - self.parent.abs_x) self.ref_angle = self.abs_angle - self.parent.abs_angle self.ref_scale = self.abs_scale / self.parent.abs_scale #如果父节点score为0，而本节点score为1，那么这个节点就只有abs_x,abs_y是可靠的，其他都是不可靠的 #除非使用IK算法，否则无法计算可靠的角度和长度。超出了这个插件编写的本意。留着以后优化吧。 #openpose的线条，在两个端点都可靠的情况下才会绘制。因此本插件即使不计算这个骨骼，也不影响绘制。 #本插件还有猜测不可靠节点的位置的功能，不放在这个函数里处理，是因为这个函数是递归调用的，放在这里的话，猜测的位置会不断被覆盖。 #猜测功能用在pose interpolation里，因为设计本意里，pose interpolation的输入来自两张不同的图片，里面的可靠节点数量不一致，所以需要猜测不可靠节点的位置。 else: self.length = 0.0 self.ref_angle = 0.0 self.abs_angle= 0.0 # 遍历子节点，递归调用Update方法 for child in self.children: child.CalculateRef() class HumanDoll: def __init__(self, canvas_width, canvas_height): self.canvas_width = canvas_width self.canvas_height = canvas_height self.parts = [None] * (len(bonesPoints) + 1) self.keyPoints=[] self.RootPart = HumanPart(1) # 根部 固定长度为0，对应keypoints[1]的位置 self.abs_x = canvas_width / 2 self.abs_y = canvas_height / 3 self.RootPart.abs_x = canvas_width / 2 self.RootPart.abs_y = canvas_height / 3 self.RootPart.ref_angle = 0.0 self.RootPart.score = 1.0 self.parts[1] = self.RootPart for (p,c) in bonesPoints: # p为父节点，c为子节点 child = HumanPart(c) child.SetParent(self.parts[p]) self.parts[c] = child def UpdateByKeypoints(self, keypoints): self.keyPoints = keypoints for i, keypoint in enumerate(keypoints): if keypoint is not None: self.parts[i].abs_x = keypoint[0] self.parts[i].abs_y = keypoint[1] self.parts[i].score = keypoint[2] else: self.parts[i].score = 0.0 self.RootPart.CalculateRef() def UpdateByParts(self): self.RootPart.Update() for i in range(len(self.keyPoints)): self.keyPoints[i] = (self.parts[i].abs_x, self.parts[i].abs_y, self.parts[i].score) def Draw(self, width, height,draw_all=False): # 在这里实现绘制函数 return keyPoints_2_image(self.keyPoints, width, height,draw_all) def guessMissingParts(self): # 在这里实现猜测缺失部分的功能。用于pose interpolation doll_T_Pose = HumanDoll(self.canvas_width,self.canvas_height) keyPoints = gen_t_pose_keypoints(self.canvas_width,self.canvas_height) doll_T_Pose.UpdateByKeypoints(keyPoints) # 定义各部位缺失时的参考顺序数组 # 格式为：(目标部位索引, [参考部位列表按优先级排序]) reference_order = [ (0, [2,5,8,11]), # 补充头部（部位0）时参考的顺序 (2, [5,8,11,0]), # 补充右肩（部位2）时参考的顺序 (5, [2,8,11,0]), # 补充左肩（部位5）时参考的顺序 (8, [11,2,5,0]), # 补充右髋（部位8）时参考的顺序 (11, [8,2,5,0]), # 补充左髋（部位11）时参考的顺序 (3, [6,2,8,11,0]), # 补充右上臂（部位3）时参考的顺序 (6, [3,2,11,8,0]), # 补充左上臂（部位6）时参考的顺序 (4, [3,7,2,5,8,11,0]), # 补充右前臂（部位4）时参考的顺序 (7, [4,6,5,3,2,8,11,0]), # 补充左前臂（部位7）时参考的顺序 (9, [12,8,11,2,5,0]), # 补充右大腿（部位9）时参考的顺序 (12, [9,11,8,5,2,0]), # 补充左大腿（部位12）时参考的顺序 (10, [13,9,12,8,11,2,5,0]), # 补充右小腿（部位10）时参考的顺序 (13, [10,12,9,11,8,2,5,0]), # 补充左小腿（部位13）时参考的顺序 #以上都是以父节点长度参考缺失节点长度，而角度则从t-pose中复制 #但是脸上的部位，需要另外的算法。实际运行中出现过耳朵可靠而眼睛不可靠的情况，因此需要另外的算法 # (14, [15,0]), # 补充右眼睛（部位14）时参考的顺序 # (15, [14,0]), # 补充左眼睛（部位15）时参考的顺序 # (16, [17,14,15,0]), # 补充右耳朵（部位16）时参考的顺序 # (17, [16,15,14,0]), # 补充左耳朵（部位17）时参考的顺序 # 其他部位的参考规则可以继续添加 # ... ] # 遍历所有定义好的参考规则 for target_part, reference_parts in reference_order: # 跳过1号部位（根据需求描述1号部位不会缺失） if target_part == 1: continue # 如果目标部位缺失（分数为0） if self.parts[target_part].score == 0.0: # 标记为已尝试修复 self.parts[target_part].score = 1.0 # 复制T姿态下的参考角度 self.parts[target_part].ref_angle = doll_T_Pose.parts[target_part].ref_angle # 遍历参考部位列表，按优先级尝试补充数据 for ref_part in reference_parts: if self.parts[ref_part].score > 0.0: # 使用参考部位数据补充目标部位长度 # 比例因子：目标部位在T姿态下的长度 / 参考部位在T姿态下的长度 ratio = doll_T_Pose.parts[target_part].length / doll_T_Pose.parts[ref_part].length #if target_part == 16 or target_part == 17: self.parts[target_part].length = self.parts[ref_part].length * ratio #print(f"[{target_part}]-[{ref_part}]-ratio:{ratio}- self.parts[ref_part].length:{self.parts[ref_part].length} - result_length:{self.parts[target_part].length}") break # 找到第一个有效参考后退出循环 else: # 如果所有参考部位都无效，重置分数 self.parts[target_part].score = 0.0 #现在处理右眼。如果右眼缺失，则根据左耳朵和鼻子的位置计算 if self.parts[14].score == 0.0: self.parts[14].score=1.0 if self.parts[16].score > 0.0 and self.parts[0].score > 0.0: self.parts[14].abs_x = (self.parts[16].abs_x + self.parts[0].abs_x)/2 self.parts[14].abs_y = (self.parts[16].abs_y + self.parts[0].abs_y)/2 self.parts[14].CalculateRef() elif self.parts[0].score > 0.0: ratio = doll_T_Pose.parts[14].length / doll_T_Pose.parts[0].length self.parts[14].length = self.parts[0].length * ratio else: self.parts[14].score = 0.0 #现在处理左眼。如果左眼缺失，则根据右耳朵和鼻子的位置计算 if self.parts[15].score == 0.0: self.parts[15].score=1.0 if self.parts[17].score > 0.0 and self.parts[0].score > 0.0: self.parts[15].abs_x = (self.parts[17].abs_x + self.parts[0].abs_x)/2 self.parts[15].abs_y = (self.parts[17].abs_y + self.parts[0].abs_y)/2 self.parts[15].CalculateRef() elif self.parts[0].score > 0.0: ratio = doll_T_Pose.parts[15].length / doll_T_Pose.parts[0].length self.parts[15].length = self.parts[0].length * ratio else: self.parts[15].score = 0.0 reference_order_part2 = [ (16, [17,14,15,0]), # 补充右耳朵（部位16）时参考的顺序 (17, [16,15,14,0]), # 补充左耳朵（部位17）时参考的顺序 ] for target_part, reference_parts in reference_order_part2: # 跳过1号部位（根据需求描述1号部位不会缺失） if target_part == 1: continue # 如果目标部位缺失（分数为0） if self.parts[target_part].score == 0.0: # 标记为已尝试修复 self.parts[target_part].score = 1.0 # 复制T姿态下的参考角度 self.parts[target_part].ref_angle = doll_T_Pose.parts[target_part].ref_angle # 遍历参考部位列表，按优先级尝试补充数据 for ref_part in reference_parts: if self.parts[ref_part].score > 0.0: # 使用参考部位数据补充目标部位长度 # 比例因子：目标部位在T姿态下的长度 / 参考部位在T姿态下的长度 ratio = doll_T_Pose.parts[target_part].length / doll_T_Pose.parts[ref_part].length #if target_part == 16 or target_part == 17: self.parts[target_part].length = self.parts[ref_part].length * ratio #print(f"[{target_part}]-[{ref_part}]-ratio:{ratio}- self.parts[ref_part].length:{self.parts[ref_part].length} - result_length:{self.parts[target_part].length}") break # 找到第一个有效参考后退出循环 else: # 如果所有参考部位都无效，重置分数 self.parts[target_part].score = 0.0 #现在处理右眼。如果右眼缺失，则根据左耳朵和鼻子的位置计算 print("修复完成") def gen_t_pose_keypoints(canvas_width, canvas_height): # 初始化关键点数组 keyPoints = [None] * (len(bonesPoints) + 1) # 定义起始点（1号关键点，颈部）的位置 start_x = canvas_width / 2 start_y = canvas_height / 4 keyPoints[1] = (start_x, start_y, 1.0) # 定义骨骼长度 neck_length = canvas_height / 10 shoulder_length = canvas_width / 8 arm_length = canvas_width / 6 hip_width = shoulder_length hip_height = canvas_height / 4 leg_length = canvas_height / 5 head_radius = canvas_height / 15 # 头部半径 eye_distance = canvas_width / 10 # 眼睛间距 ear_distance = canvas_width / 10 # 耳朵间距 # 计算0号关键点（头部中心）的位置 head_center_x = start_x head_center_y = start_y - neck_length keyPoints[0] = (head_center_x, head_center_y, 1.0) # 计算右肩（2号关键点）的位置 right_shoulder_x = start_x - shoulder_length right_shoulder_y = start_y keyPoints[2] = (right_shoulder_x, right_shoulder_y, 1.0) # 计算左肩（5号关键点）的位置 left_shoulder_x = start_x + shoulder_length left_shoulder_y = start_y keyPoints[5] = (left_shoulder_x, left_shoulder_y, 1.0) # 计算右肘（3号关键点）的位置 right_elbow_x = right_shoulder_x - arm_length right_elbow_y = right_shoulder_y keyPoints[3] = (right_elbow_x, right_elbow_y, 1.0) # 计算左肘（6号关键点）的位置 left_elbow_x = left_shoulder_x + arm_length left_elbow_y = left_shoulder_y keyPoints[6] = (left_elbow_x, left_elbow_y, 1.0) # 计算右手腕（4号关键点）的位置 right_wrist_x = right_elbow_x - arm_length right_wrist_y = right_elbow_y keyPoints[4] = (right_wrist_x, right_wrist_y, 1.0) # 计算左手腕（7号关键点）的位置 left_wrist_x = left_elbow_x+ arm_length left_wrist_y = left_elbow_y keyPoints[7] = (left_wrist_x, left_wrist_y, 1.0) # 计算右髋（8号关键点）的位置 right_hip_x = start_x - hip_width right_hip_y = start_y + hip_height keyPoints[8] = (right_hip_x, right_hip_y, 1.0) # 计算左髋（11号关键点）的位置 left_hip_x = start_x + hip_width left_hip_y = start_y + hip_height keyPoints[11] = (left_hip_x, left_hip_y, 1.0) # 计算右膝（9号关键点）的位置 right_knee_x = right_hip_x right_knee_y = right_hip_y + leg_length keyPoints[9] = (right_knee_x, right_knee_y, 1.0) # 计算左膝（12号关键点）的位置 left_knee_x = left_hip_x left_knee_y = left_hip_y + leg_length keyPoints[12] = (left_knee_x, left_knee_y, 1.0) # 计算右脚踝（10号关键点）的位置 right_ankle_x = right_knee_x right_ankle_y = right_knee_y + leg_length keyPoints[10] = (right_ankle_x, right_ankle_y, 1.0) # 计算左脚踝（13号关键点）的位置 left_ankle_x = left_knee_x left_ankle_y = left_knee_y + leg_length keyPoints[13] = (left_ankle_x, left_ankle_y, 1.0) # 计算右眼（14号关键点）的位置 right_eye_x = head_center_x - eye_distance / 2 right_eye_y = head_center_y- eye_distance / 2 keyPoints[14] = (right_eye_x, right_eye_y, 1.0) # 计算左眼（15号关键点）的位置 left_eye_x = head_center_x + eye_distance / 2 left_eye_y = head_center_y- eye_distance / 2 keyPoints[15] = (left_eye_x, left_eye_y, 1.0) # 计算右耳（16号关键点）的位置 right_ear_x = head_center_x - ear_distance right_ear_y = head_center_y keyPoints[16] = (right_ear_x, right_ear_y, 1.0) # 计算左耳（17号关键点）的位置 left_ear_x = head_center_x + ear_distance left_ear_y = head_center_y keyPoints[17] = (left_ear_x, left_ear_y, 1.0) return keyPoints def poseKeypoints2D_2_keypoints(pose_keypoints_2d, canvas_width, canvas_height,landmarkType): keyPoints = [pose_keypoints_2d[i:i + 3] for i in range(0, len(pose_keypoints_2d), 3)] if landmarkType == "DWPose": canvas_height=1 canvas_width=1 for i in range(len(keyPoints)): x, y, z = keyPoints[i] keyPoints[i] = (x * canvas_width, y * canvas_height, z) return keyPoints def keypoints_2_bones(keyPoints): bones=[] for i, (a, b) in enumerate(bonesPoints): a_x, a_y, a_z = keyPoints[a] b_x, b_y, b_z = keyPoints[b] # 现在 (a_x,a_y) 是骨骼起点的坐标，(b_x,b_y) 是骨骼终点的坐标 # 然后计算骨骼的长度和角度，连同坐标一起生成一个骨骼对象，放入bones数组 # 如果a_z=0或者b_z=0，则忽略这个骨骼，放入一个空对象 if a_z == 0 or b_z == 0: bone = { 'start': (a_x, a_y), 'end': (b_x, b_y), 'length': -1, 'angle': 0 } bones.append(bone) continue # 计算骨骼的长度 bone_length = math.sqrt((b_x - a_x) ** 2 + (b_y - a_y) ** 2) # 计算骨骼的角度 bone_angle = math.atan2(b_y - a_y, b_x - a_x) # 生成骨骼对象 bone = { 'start': (a_x, a_y), 'end': (b_x, b_y), 'length': bone_length, 'angle': bone_angle } bones.append(bone) return bones def bones_2_keypoints(bones): keyPoints = [None] * 18 # 从1号关键点开始 start_bone = bones[0] start_x, start_y = start_bone['start'] neck_length = start_bone['length'] neck_angle = start_bone['angle'] keyPoints[1] = (start_x ,start_y,1.0) # 计算新的0号关键点的位置 new_x = start_x + neck_length * math.cos(neck_angle) new_y = start_y + neck_length * math.sin(neck_angle) keyPoints[0] = (new_x ,new_y,1.0) # 依次计算后续关键点的位置 for i, (a, b) in enumerate(bonesPoints): if i == 0: continue # 跳过第一个骨骼，因为它已经处理过了 # 获取前一个关键点的位置 prev_x,prev_y,score = keyPoints[a] # 获取当前骨骼的信息 bone = bones[i] bone_length = bone['length'] bone_angle = bone['angle'] if bone_length == -1: keyPoints[b]=(0,0,0) continue # 计算新的关键点位置 new_x = prev_x + bone_length * math.cos(bone_angle) new_y = prev_y + bone_length * math.sin(bone_angle) # 将新的关键点位置存入数组 keyPoints[b]=(new_x ,new_y,1.0) return keyPoints def keyPoints_2_image(keyPoints, canvas_width, canvas_height ,draw_all=False): image = np.zeros((canvas_height, canvas_width, 3), dtype=np.uint8) # 绘制骨架 for i, (a, b) in enumerate(skeleton): a_x, a_y, a_z = keyPoints[a] b_x, b_y, b_z = keyPoints[b] if draw_all: if a_z != 0 and b_z != 0: cv2.line(image, (int(a_x), int(a_y)), (int(b_x), int(b_y)), connect_color[i] + [0], 4) else: cv2.line(image, (int(a_x), int(a_y)), (int(b_x), int(b_y)), (255,255,255), 4) else: # 检查是否在画面内（假设a_z和b_z为0表示不在画面内） a_in_canvas = a_z != 0 and 0 <= a_x < canvas_width and 0 <= a_y < canvas_height b_in_canvas = b_z != 0 and 0 <= b_x < canvas_width and 0 <= b_y < canvas_height if a_in_canvas and b_in_canvas: cv2.line(image, (int(a_x), int(a_y)), (int(b_x), int(b_y)), connect_color[i] + [0], 4) # 绘制关键点 for i, (x, y, z) in enumerate(keyPoints): if z!=0: try: cv2.circle(image, (int(x), int(y)), int(canvas_width/100), connect_color[i], -1) except: pass # cv2.putText(image, str(i), (int(x), int(y)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1) return image def draw_skeleton(pose_keypoints_2d, canvas_width, canvas_height,landmarkType): # 加载背景图片或创建一个空白画布 image = None #cv2.imread('background.jpg') # 使用实际的背景图片路径 if image is None: image = np.zeros((canvas_height, canvas_width, 3), dtype=np.uint8) if landmarkType == "DWPose": canvas_height=1 canvas_width=1 tri_tuples = [pose_keypoints_2d[i:i + 3] for i in range(0, len(pose_keypoints_2d), 3)] # 绘制骨架 for i, (a, b) in enumerate(skeleton): a_x, a_y, a_z = tri_tuples[a] a_x, a_y = ( a_x * canvas_width, a_y * canvas_height ) b_x, b_y, b_z = tri_tuples[b] b_x, b_y = ( b_x * canvas_width, b_y * canvas_height ) if a_z != 0 and b_z != 0: cv2.line(image, (int(a_x), int(a_y)), (int(b_x), int(b_y)), connect_color[i] + [0], 4) # 绘制关键点 for i, (x, y, z) in enumerate(tri_tuples): if z!=0: cv2.circle(image, (int(x * canvas_width), int(y * canvas_height)), 6, connect_color[i], -1) #cv2.putText(image, str(i), (int(x * canvas_width), int(y * canvas_height)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1) return image def transform_keypoints(keypoints_1, keypoints_2, frames): tri_tuples_1 = [keypoints_1[i:i + 3] for i in range(0, len(keypoints_1), 3)] tri_tuples_2 = [keypoints_2[i:i + 3] for i in range(0, len(keypoints_2), 3)] keypoints_array = [keypoints_1] for j in range(1, frames): kp = [] for i in range(len(tri_tuples_1)): x1, y1, z1 = tri_tuples_1[i] x2, y2, z2 = tri_tuples_2[i] if z1 == 0 and z2 == 0: new_x, new_y, new_z = (0.0, 0.0, 0.0) elif z1 == 0: new_x, new_y, new_z = (x2, y2, z2) elif z2 == 0: new_x, new_y, new_z = (x1, y1, z1) else: new_x, new_y, new_z = ( x1 + (x2-x1) * j/frames, y1 + (y2-y1) * j/frames , 1.0) kp.append( new_x) kp.append( new_y) kp.append( new_z) keypoints_array.append(kp) #keypoints_array.append(keypoints_2) return keypoints_array class Pose_Inter: def __init__(self): pass @classmethod def INPUT_TYPES(s): return { "required": { "pose_from": ("POSE_KEYPOINT", ), "pose_to": ("POSE_KEYPOINT", ), "interpolate_frames": ("INT", {"default": 10, "min": 2, "max": 100, "step": 1}), "landmarkType": (["OpenPose","DWPose"], ), }, } RETURN_TYPES = ("IMAGE",) FUNCTION = "run" CATEGORY = "Pose Interpolation" def run(self,pose_from,pose_to,interpolate_frames,landmarkType): openpose_dict_2 = pose_from[0] openpose_dict = pose_to[0] keypoints_array = transform_keypoints( openpose_dict_2["people"][0]["pose_keypoints_2d"], openpose_dict["people"][0]["pose_keypoints_2d"], interpolate_frames ) output=[] #print("image shape") #print(image.shape) for i, keypoints in enumerate(keypoints_array): # 显示图像 image = draw_skeleton( keypoints, openpose_dict_2["canvas_width"], openpose_dict_2["canvas_height"], landmarkType ) image = torch.from_numpy(image.astype(np.float32) / 255.0)#.unsqueeze(0) output.append(image) tensor_stacked = torch.stack(output) #print("shape of tensor_stacked") #print(tensor_stacked.shape) return (tensor_stacked,) class PoseModify: def __init__(self): pass @classmethod def INPUT_TYPES(s): return { "required": { "pose_keypoint": ("POSE_KEYPOINT", ), "landmarkType": (["OpenPose","DWPose"], ), "globalScale": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10.0}), "head": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10.0,"step":0.1}), "neck": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10.0,"step":0.1}), "neck": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10.0,"step":0.1}), "torsoWidth": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10.0,"step":0.1}), "torsoHeight": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10.0,"step":0.1}), "upperArmLength": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10.0,"step":0.1}), "lowerArmLength": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10.0,"step":0.1}), "upperLegLength": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10.0,"step":0.1}), "lowerLegLength": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10.0,"step":0.1}), }, } RETURN_TYPES = ("IMAGE",) FUNCTION = "run" CATEGORY = "Pose Interpolation" def run(self,pose_keypoint,landmarkType,globalScale,head,neck,torsoWidth,torsoHeight,upperArmLength,lowerArmLength,upperLegLength,lowerLegLength): # if head==1.0 & neck==1.0 & torsoWidth==1.0 & torsoHeight==1.0 & upperArmLength==1.0 & lowerArmLength==1.0 & upperLegLength==1.0 & lowerLegLength==1.0: # lengthFactors = [1.0] * 18 # else: lengthFactors = [neck,1.0, torsoWidth, upperArmLength, lowerArmLength, torsoWidth, upperArmLength, lowerArmLength, torsoHeight, upperLegLength, lowerLegLength, torsoHeight, upperLegLength, lowerLegLength, head, head, head, head] openpose_dict_2 = pose_keypoint[0] width = openpose_dict_2["canvas_width"] height = openpose_dict_2["canvas_height"] output=[] openpose_dict_2 = pose_keypoint[0] width = openpose_dict_2["canvas_width"] height = openpose_dict_2["canvas_height"] keyPoints = gen_t_pose_keypoints(width,height) doll = HumanDoll(width,height) doll.UpdateByKeypoints(keyPoints) for openpose_dict_2 in pose_keypoint: human1Keypoints = poseKeypoints2D_2_keypoints( openpose_dict_2["people"][0]["pose_keypoints_2d"], width, height, landmarkType ) doll.UpdateByKeypoints(human1Keypoints) doll.RootPart.ref_scale = globalScale for i in range(len(doll.parts)): if doll.parts[i].score > 0.0: doll.parts[i].length *= lengthFactors[i] doll.UpdateByParts() image=doll.Draw(width,height) image = torch.from_numpy(image.astype(np.float32) / 255.0)#.unsqueeze(0) output.append(image) tensor_stacked = torch.stack(output) #print("shape of tensor_stacked") #print(tensor_stacked.shape) return (tensor_stacked,) class GenTPose: def __init__(self): pass @classmethod def INPUT_TYPES(s): return { "required": { "width": ("INT", {"default": 480, "min": 2, "max": 1024, "step": 1}), "height": ("INT", {"default": 640, "min": 2, "max": 1024, "step": 1}), }, } RETURN_TYPES = ("IMAGE",) FUNCTION = "run" CATEGORY = "Pose Interpolation" def run(self,width,height): keyPoints = gen_t_pose_keypoints(width,height) doll = HumanDoll(width, height) doll.UpdateByKeypoints(keyPoints) doll.UpdateByParts() image=doll.Draw(width, height) output=[] image = torch.from_numpy(image.astype(np.float32) / 255.0)#.unsqueeze(0) output.append(image) tensor_stacked = torch.stack(output) #print("shape of tensor_stacked") #print(tensor_stacked.shape) return (tensor_stacked,) class Pose_Inter_V2: def __init__(self): pass @classmethod def INPUT_TYPES(s): return { "required": { "pose_from": ("POSE_KEYPOINT", ), "pose_to": ("POSE_KEYPOINT", ), "interpolate_frames": ("INT", {"default": 10, "min": 2, "max": 100, "step": 1}), "landmarkType": (["OpenPose","DWPose"], ), "draw_all": ("BOOLEAN", {"default": False}), }, } RETURN_TYPES = ("IMAGE",) FUNCTION = "run" CATEGORY = "Pose Interpolation" def run(self,pose_from,pose_to,interpolate_frames,landmarkType,draw_all): output=[] openpose_dict = pose_from[0] openpose_dict_2 = pose_to[0] width = openpose_dict_2["canvas_width"] height = openpose_dict_2["canvas_height"] dollFrom = HumanDoll(width,height) fromKeys = poseKeypoints2D_2_keypoints(openpose_dict["people"][0]["pose_keypoints_2d"],width,height,landmarkType) dollFrom.UpdateByKeypoints(fromKeys) dollTo = HumanDoll(width,height) toKeys = poseKeypoints2D_2_keypoints(openpose_dict_2["people"][0]["pose_keypoints_2d"],width,height,landmarkType) dollTo.UpdateByKeypoints(toKeys) dollFrom.guessMissingParts() dollTo.guessMissingParts() image=dollFrom.Draw(width,height,draw_all) image = torch.from_numpy(image.astype(np.float32) / 255.0) output.append(image) position_x_arr = np.full(interpolate_frames+2, -1.0, dtype=np.float32) position_y_arr = np.full(interpolate_frames+2,-1.0, dtype=np.float32) position_x_arr[0] = dollFrom.RootPart.abs_x position_y_arr[0] = dollFrom.RootPart.abs_y position_x_arr[interpolate_frames+1] = dollTo.RootPart.abs_x position_y_arr[interpolate_frames+1] = dollTo.RootPart.abs_y for i in range(1, interpolate_frames+1): position_x_arr[i] = position_x_arr[0] + (position_x_arr[interpolate_frames+1] - position_x_arr[0]) * (i - 1) / interpolate_frames position_y_arr[i] = position_y_arr[0] + (position_y_arr[interpolate_frames+1] - position_y_arr[0]) * (i - 1) / interpolate_frames #18个骨骼，每个骨骼2个数组，每个数组有interpolate_frames+2个元素表示骨骼每一帧的长度和角度 curveDatas = np.full((18, 2, interpolate_frames+2), -1.0, dtype=np.float32) for i in range(len(dollFrom.parts)): if dollFrom.parts[i].score > 0.0: curveDatas[i][0][0] = dollFrom.parts[i].length curveDatas[i][1][0] = dollFrom.parts[i].ref_angle dollFrom.parts[i].score=1.0 for i in range(len(dollTo.parts)): if dollTo.parts[i].score > 0.0: curveDatas[i][0][interpolate_frames+1] = dollTo.parts[i].length curveDatas[i][1][interpolate_frames+1] = dollTo.parts[i].ref_angle dollTo.parts[i].score=1.0 for i in range(18): j=0 #骨骼长度，线性插值 for k in range(1, interpolate_frames+1): curveDatas[i][j][k] = curveDatas[i][j][0] + (curveDatas[i][j][interpolate_frames+1] - curveDatas[i][j][0]) * (k) / (interpolate_frames+1) for i in range(18): j=1 a = curveDatas[i][j][0] b = curveDatas[i][j][interpolate_frames+1] # 将 b 调整到与 a 变化量最小的等效弧度值 b_adjusted = b - 2 * np.pi * np.round((b - a) / (2 * np.pi)) curveDatas[i][j][interpolate_frames+1] = b_adjusted for k in range(1, interpolate_frames+1): curveDatas[i][j][k] = a + (b_adjusted - a) * (k) / (interpolate_frames+1) curveDatas[i][j] = np.mod(curveDatas[i][j], 2 * np.pi) for i in range(1,interpolate_frames+1): dollFrom.RootPart.abs_x = position_x_arr[i] dollFrom.RootPart.abs_y = position_y_arr[i] for k in range(18): dollFrom.parts[k].length = curveDatas[k][0][i] dollFrom.parts[k].ref_angle = curveDatas[k][1][i] dollFrom.UpdateByParts() image=dollFrom.Draw(width,height,draw_all) image = torch.from_numpy(image.astype(np.float32) / 255.0)#.unsqueeze(0) output.append(image) image=dollTo.Draw(width,height,draw_all) image = torch.from_numpy(image.astype(np.float32) / 255.0)#.unsqueeze(0) output.append(image) tensor_stacked = torch.stack(output) return (tensor_stacked,) # A dictionary that contains all nodes you want to export with their names # NOTE: names should be globally unique NODE_CLASS_MAPPINGS = { "Pose_Inter": Pose_Inter, "Pose_Inter_V2": Pose_Inter_V2, "PoseModify": PoseModify, "GenTPose":GenTPose } # A dictionary that contains the friendly/humanly readable titles for the nodes NODE_DISPLAY_NAME_MAPPINGS = { "Pose_Inter": "Pose Interpolation", "Pose_Inter_V2": "Pose Interpolation_V2", "PoseModify": "Pose Modify by factor", "GenTPose": "GenTPose" }