save code

src/custom_controlnet_aux/dwpose/__init__.py

@@ -105,10 +105,12 @@ def my_draw_poses(poses: List[PoseResult], H, W, draw_body=True, draw_hand=True,
 import torch
 import math
 
-def draw_pose_mask_line(mask: torch.Tensor, W, H, p1, p2, line_width=50):
+def draw_pose_mask_line(mask: torch.Tensor, H, W, p1, p2, line_width=50):
     x1, y1 = p1.x, p1.y
     x2, y2 = p2.x, p2.y
     
+    print(f'draw line p1({x1},{y1})  p2({x2},{y2})')
+
     # Create coordinate grids
     y_grid, x_grid = torch.meshgrid(torch.arange(H), torch.arange(W))
     x_grid = x_grid.float()
@@ -163,19 +165,103 @@ def get_point_on_line(start, end, t):
     
     return Keypoint(x, y)
 
-def draw_pose_mask(mask: torch.Tensor, keypoints, H, W, changdu):
-    points = []
-    if changdu == "yao":
-        points.append(keypoints[3-1])
-        points.append(get_point_on_line(keypoints[2-1], keypoints[1-1], 0.2))
-        points.append(keypoints[6-1])
-        points.append(keypoints[7-1])
-        points.append(get_point_on_line(keypoints[12-1], keypoints[2-1], 0.2)) 
-        points.append(get_point_on_line(keypoints[9-1], keypoints[2-1], 0.2)) 
-        points.append(keypoints[4-1])
+def add_point_old(points, keypoints, index1, index2 = None, offset = 0):
+    if index1 in keypoints and index2 in keypoints:
+        p1 = keypoints[index1]
+        p2 = keypoints[index2]
+        outpoint = get_point_on_line(p1, p2, offset)
+    elif index1 in keypoints:
+        outpoint = keypoints[index1]
+    elif index2 in keypoints:
+        outpoint = keypoints[index2]
     else:
-        # 如果不是"yao"，使用所有关键点
-        points = keypoints
+        outpoint = None
+
+    if outpoint != None:
+        points.append(outpoint)
+
+from functools import cmp_to_key
+def convex_hull(points):
+    """返回Point对象的凸包（顺时针顺序）"""
+    if len(points) <= 1:
+        return points.copy()
+    
+    # 1. 找基准点（y最小，相同时x最小）
+    pivot = min(points, key=lambda p: (p.y, p.x))
+    
+    # 2. 极角排序比较函数（避免浮点误差）
+    def compare(a, b):
+        # 计算叉积
+        cross = (a.x - pivot.x) * (b.y - pivot.y) - (a.y - pivot.y) * (b.x - pivot.x)
+        if cross == 0:  # 共线时按距离排序
+            dist_a = (a.x - pivot.x)**2 + (a.y - pivot.y)**2
+            dist_b = (b.x - pivot.x)**2 + (b.y - pivot.y)**2
+            return dist_a - dist_b
+        return -cross  # 逆时针方向
+    
+    # 排序（跳过基准点）
+    sorted_points = sorted([p for p in points if p != pivot], key=cmp_to_key(compare))
+    
+    # 3. Graham扫描
+    stack = [pivot, sorted_points[0]]
+    for p in sorted_points[1:]:
+        while len(stack) >= 2:
+            # 取栈顶两个点和当前点计算叉积
+            o, a = stack[-2], stack[-1]
+            cross = (a.x - o.x) * (p.y - o.y) - (a.y - o.y) * (p.x - o.x)
+            if cross <= 0:  # 非左转则弹出
+                stack.pop()
+            else:
+                break
+        stack.append(p)
+    
+    return stack
+
+def clipOffset(v, min, max):
+    if v < min:
+        return min
+    if v > max:
+        return max
+    return v
+
+def add_point(points, keypoints, index, H, W, offsetX=0, offsetY=0):
+    if keypoints[index] != None:
+        p = keypoints[index]
+        sx = p.x + offsetX
+        sy = p.y + offsetY
+        x = clipOffset(sx, 0, W-1)
+        y = clipOffset(sy, 0, H-1)
+        print(f'add src:({p.x} {p.y}) index:{index} point:({x} {y})')
+        points.append(Keypoint(x,y))
+    
+        
+
+def distance(p1, p2):
+    return math.sqrt((p1.x - p2.x)**2 + (p1.y - p2.y)**2)
+
+
+def draw_pose_mask(mask: torch.Tensor, keypoints, H, W, cloth_len):
+    points = []
+    unit_len = distance(keypoints[0], keypoints[1])
+    print(f'unit_len {unit_len}')
+    if cloth_len == "腰":
+        add_point(points, keypoints, 2, H, W, -unit_len*0.2, 0)
+        add_point(points, keypoints, 1, H, W, 0, -unit_len*0.2)
+        add_point(points, keypoints, 5, H, W, unit_len*0.2, 0)
+        add_point(points, keypoints, 6, H, W)
+        add_point(points, keypoints, 11, H, W, unit_len*0.2, -unit_len*0.3)
+        add_point(points, keypoints, 8, H, W, -unit_len*0.2, -unit_len*0.3)
+        add_point(points, keypoints, 3, H, W)
+    else:
+        add_point(points, keypoints, 2, H, W, -unit_len*0.2, 0)
+        add_point(points, keypoints, 1, H, W, 0, -unit_len*0.2)
+        add_point(points, keypoints, 5, H, W, unit_len*0.2, 0)
+        add_point(points, keypoints, 6, H, W)
+        add_point(points, keypoints, 11, H, W, unit_len*0.2, -unit_len*0.3)
+        add_point(points, keypoints, 8, H, W, -unit_len*0.2, -unit_len*0.3)
+        add_point(points, keypoints, 3, H, W)
+
+    points = convex_hull(points)
 
     for i in range(len(points)):
         if i == len(points) - 1:
@@ -186,16 +272,16 @@ def draw_pose_mask(mask: torch.Tensor, keypoints, H, W, changdu):
             p2 = points[i + 1]
 
         if p1 is not None and p2 is not None:
-            draw_pose_mask_line(mask, W, H, p1, p2, line_width=100)
+            draw_pose_mask_line(mask, H, W, p1, p2, line_width=(unit_len*0.1))
 
     return mask
 
-def draw_poses_mask(poses: List[PoseResult], H, W, changdu):
+def draw_poses_mask(poses: List[PoseResult], H, W, cloth_len):
 
     mask = torch.zeros((H, W), dtype=torch.float32)
 
     for pose in poses:
-        mask = draw_pose_mask(mask, pose.body.keypoints, H, W, changdu)
+        mask = draw_pose_mask(mask, pose.body.keypoints, H, W, cloth_len)
             
     return mask