diff --git a/__init__.py b/__init__.py index 23849c1..d475e51 100644 --- a/__init__.py +++ b/__init__.py @@ -6,7 +6,7 @@ import cv2 import numpy as np import folder_paths import os - +import math connect_color = [ [ 0, 0, 255], @@ -42,8 +42,489 @@ skeleton = [ [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 +] -def gen_skeleton(pose_keypoints_2d, canvas_width, canvas_height,landmarkType): +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: @@ -137,7 +618,7 @@ class Pose_Inter: for i, keypoints in enumerate(keypoints_array): # 显示图像 - image = gen_skeleton( + image = draw_skeleton( keypoints, openpose_dict_2["canvas_width"], openpose_dict_2["canvas_height"], @@ -150,14 +631,252 @@ class Pose_Inter: #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": 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": "Pose Interpolation", + "Pose_Inter_V2": "Pose Interpolation_V2", + "PoseModify": "Pose Modify by factor", + "GenTPose": "GenTPose" } diff --git a/poseInterV2.jpg b/poseInterV2.jpg new file mode 100644 index 0000000..e3170a1 Binary files /dev/null and b/poseInterV2.jpg differ diff --git a/poseModify.jpg b/poseModify.jpg new file mode 100644 index 0000000..858d78e Binary files /dev/null and b/poseModify.jpg differ