save code

This commit is contained in:
xsl
2025-08-09 14:50:05 +08:00
parent 62ef148f72
commit 2f1e03ca15
2 changed files with 67 additions and 26 deletions
+26 -25
View File
@@ -241,19 +241,20 @@ def distance(p1, p2):
def addBase(points, keypoints, H, W, unit_len, closh_short):
add_point(points, keypoints, 2, H, W, -unit_len*0.4, -unit_len*0.1)
add_point(points, keypoints, 1, H, W, 0, -unit_len*0.25)
add_point(points, keypoints, 5, H, W, unit_len*0.4, -unit_len*0.1)
add_point(points, keypoints, 2, H, W, -unit_len*0.4, -unit_len*0.2)
add_point(points, keypoints, 1, H, W, 0, -unit_len*0.4)
add_point(points, keypoints, 5, H, W, unit_len*0.4, -unit_len*0.2)
add_point(points, keypoints, 6, H, W, unit_len*0.2)
add_point(points, keypoints, 3, H, W, -unit_len*0.2)
if not closh_short:
add_point(points, keypoints, 4, H, W, 0, 0)
add_point(points, keypoints, 7, H, W, 0, 0)
# if not closh_short:
add_point(points, keypoints, 4, H, W, 0, 0)
add_point(points, keypoints, 7, H, W, 0, 0)
def draw_pose_mask(mask: torch.Tensor, keypoints, H, W, cloth_len, line_len, closh_short=True):
points = []
unit_len = distance(keypoints[0], keypoints[1])
print(f'unit_len:{unit_len} cloth_len:{cloth_len}')
offsetX = 0.8
if cloth_len == "":
addBase(points, keypoints, H, W, unit_len, closh_short)
add_point(points, keypoints, 11, H, W, unit_len*0.2, -unit_len*0.3)
@@ -274,8 +275,8 @@ def draw_pose_mask(mask: torch.Tensor, keypoints, H, W, cloth_len, line_len, clo
add_point(points, keypoints, 11, H, W, unit_len*0.2, unit_len*2.2)
add_point(points, keypoints, 8, H, W, -unit_len*0.2, unit_len*2.2)
else:
add_point(points, keypoints, 12, H, W, unit_len*0.2, 0)
add_point(points, keypoints, 9, H, W, -unit_len*0.2, 0)
add_point(points, keypoints, 12, H, W, unit_len*offsetX, 0)
add_point(points, keypoints, 9, H, W, -unit_len*offsetX, 0)
elif cloth_len == '小腿':
addBase(points, keypoints, H, W, unit_len, closh_short)
add_point(points, keypoints, 11, H, W, unit_len*0.2, -unit_len*1.5)
@@ -284,10 +285,10 @@ def draw_pose_mask(mask: torch.Tensor, keypoints, H, W, cloth_len, line_len, clo
add_point(points, keypoints, 11, H, W, unit_len*0.2, unit_len*3)
add_point(points, keypoints, 8, H, W, -unit_len*0.2, unit_len*3)
else:
add_point(points, keypoints, 12, H, W, unit_len*0.2, 0)
add_point(points, keypoints, 9, H, W, -unit_len*0.2, 0)
add_point(points, keypoints, 12, H, W, unit_len*0.2, unit_len*1.2)
add_point(points, keypoints, 9, H, W, -unit_len*0.2, unit_len*1.2)
add_point(points, keypoints, 12, H, W, unit_len*offsetX, 0)
add_point(points, keypoints, 9, H, W, -unit_len*offsetX, 0)
add_point(points, keypoints, 12, H, W, unit_len*offsetX, unit_len*1.2)
add_point(points, keypoints, 9, H, W, -unit_len*offsetX, unit_len*1.2)
elif cloth_len == '脚踝':
addBase(points, keypoints, H, W, unit_len, closh_short)
add_point(points, keypoints, 11, H, W, unit_len*0.2, -unit_len*1.5)
@@ -296,14 +297,14 @@ def draw_pose_mask(mask: torch.Tensor, keypoints, H, W, cloth_len, line_len, clo
add_point(points, keypoints, 11, H, W, unit_len*0.2, unit_len*6)
add_point(points, keypoints, 8, H, W, -unit_len*0.2, unit_len*6)
else:
add_point(points, keypoints, 12, H, W, unit_len*0.2, 0)
add_point(points, keypoints, 9, H, W, -unit_len*0.2, 0)
add_point(points, keypoints, 12, H, W, unit_len*offsetX, 0)
add_point(points, keypoints, 9, H, W, -unit_len*offsetX, 0)
if keypoints[10] == None or keypoints[13] == None:
add_point(points, keypoints, 12, H, W, unit_len*0.2, unit_len*2.2)
add_point(points, keypoints, 9, H, W, -unit_len*0.2, unit_len*2.2)
add_point(points, keypoints, 12, H, W, unit_len*offsetX, unit_len*2.2)
add_point(points, keypoints, 9, H, W, -unit_len*offsetX, unit_len*2.2)
else:
add_point(points, keypoints, 13, H, W, unit_len*0.2, unit_len*1.2)
add_point(points, keypoints, 10, H, W, -unit_len*0.2, unit_len*1.2)
add_point(points, keypoints, 13, H, W, unit_len*offsetX, unit_len*1.2)
add_point(points, keypoints, 10, H, W, -unit_len*offsetX, unit_len*1.2)
elif cloth_len == '拖地':
addBase(points, keypoints, H, W, unit_len, closh_short)
add_point(points, keypoints, 11, H, W, unit_len*0.2, -unit_len*1.5)
@@ -312,21 +313,21 @@ def draw_pose_mask(mask: torch.Tensor, keypoints, H, W, cloth_len, line_len, clo
add_point(points, keypoints, 11, H, W, unit_len*0.2, unit_len*7)
add_point(points, keypoints, 8, H, W, -unit_len*0.2, unit_len*7)
else:
add_point(points, keypoints, 12, H, W, unit_len*0.2, 0)
add_point(points, keypoints, 9, H, W, -unit_len*0.2, 0)
add_point(points, keypoints, 12, H, W, unit_len*offsetX, 0)
add_point(points, keypoints, 9, H, W, -unit_len*offsetX, 0)
if keypoints[10] == None or keypoints[13] == None:
add_point(points, keypoints, 12, H, W, unit_len*0.2, unit_len*3.2)
add_point(points, keypoints, 9, H, W, -unit_len*0.2, unit_len*3.2)
add_point(points, keypoints, 12, H, W, unit_len*offsetX, unit_len*3.2)
add_point(points, keypoints, 9, H, W, -unit_len*offsetX, unit_len*3.2)
else:
add_point(points, keypoints, 13, H, W, unit_len*0.2, unit_len*1)
add_point(points, keypoints, 10, H, W, -unit_len*0.2, unit_len*1)
add_point(points, keypoints, 13, H, W, unit_len*offsetX, unit_len*1)
add_point(points, keypoints, 10, H, W, -unit_len*offsetX, unit_len*1)
else:
addBase(points, keypoints, H, W, unit_len, closh_short)
add_point(points, keypoints, 11, H, W, unit_len*0.2, 0)
add_point(points, keypoints, 8, H, W, -unit_len*0.2, 0)
points = convex_hull(points)
line_len = 0.7
# line_len = 0.7
for i in range(len(points)):
if i == len(points) - 1:
p1 = points[i]