1+ import pyvista as pv
2+ import numpy as np
3+ import vtk
4+
5+ # 生成更像人形的仿真机器人模型（优化比例+结构）
6+ def create_robot (return_parts = False ):
7+ # 1. 躯干（更修长，接近人形）
8+ torso = pv .Cube (center = (0 , 0 , 0.8 ), x_length = 0.4 , y_length = 0.3 , z_length = 1.2 )
9+ torso ["part_id" ] = np .zeros (torso .n_points , dtype = int ) # 躯干ID=0
10+
11+ # 2. 头部（新增，更像人形）
12+ head = pv .Cube (center = (0 , 0 , 1.6 ), x_length = 0.3 , y_length = 0.3 , z_length = 0.3 )
13+ head ["part_id" ] = np .full (head .n_points , 5 , dtype = int ) # 头部ID=5
14+
15+ # 3. 右胳膊（上臂+前臂，比例优化）
16+ right_upper_arm = pv .Cube (center = (0.5 , 0 , 1.0 ), x_length = 0.2 , y_length = 0.2 , z_length = 0.5 )
17+ right_upper_arm ["part_id" ] = np .ones (right_upper_arm .n_points , dtype = int ) # 右上臂ID=1
18+ right_forearm = pv .Cube (center = (0.9 , 0 , 1.0 ), x_length = 0.2 , y_length = 0.2 , z_length = 0.5 )
19+ right_forearm ["part_id" ] = np .full (right_forearm .n_points , 2 , dtype = int ) # 右前臂ID=2
20+
21+ # 4. 左胳膊（新增，对称结构）
22+ left_upper_arm = pv .Cube (center = (- 0.5 , 0 , 1.0 ), x_length = 0.2 , y_length = 0.2 , z_length = 0.5 )
23+ left_upper_arm ["part_id" ] = np .full (left_upper_arm .n_points , 6 , dtype = int ) # 左上臂ID=6
24+ left_forearm = pv .Cube (center = (- 0.9 , 0 , 1.0 ), x_length = 0.2 , y_length = 0.2 , z_length = 0.5 )
25+ left_forearm ["part_id" ] = np .full (left_forearm .n_points , 7 , dtype = int ) # 左前臂ID=7
26+
27+ # 5. 右大腿+小腿（比例优化）
28+ right_thigh = pv .Cube (center = (0.2 , 0 , - 0.2 ), x_length = 0.2 , y_length = 0.2 , z_length = 0.6 )
29+ right_thigh ["part_id" ] = np .full (right_thigh .n_points , 8 , dtype = int ) # 右大腿ID=8
30+ right_calf = pv .Cube (center = (0.2 , 0 , - 0.8 ), x_length = 0.2 , y_length = 0.2 , z_length = 0.6 )
31+ right_calf ["part_id" ] = np .full (right_calf .n_points , 9 , dtype = int ) # 右小腿ID=9
32+
33+ # 6. 左大腿+小腿（比例优化）
34+ left_thigh = pv .Cube (center = (- 0.2 , 0 , - 0.2 ), x_length = 0.2 , y_length = 0.2 , z_length = 0.6 )
35+ left_thigh ["part_id" ] = np .full (left_thigh .n_points , 3 , dtype = int ) # 左大腿ID=3
36+ left_calf = pv .Cube (center = (- 0.2 , 0 , - 0.8 ), x_length = 0.2 , y_length = 0.2 , z_length = 0.6 )
37+ left_calf ["part_id" ] = np .full (left_calf .n_points , 4 , dtype = int ) # 左小腿ID=4
38+
39+ # 用MultiBlock合并所有部件
40+ block = pv .MultiBlock ()
41+ block .append (torso )
42+ block .append (head )
43+ block .append (right_upper_arm )
44+ block .append (right_forearm )
45+ block .append (left_upper_arm )
46+ block .append (left_forearm )
47+ block .append (right_thigh )
48+ block .append (right_calf )
49+ block .append (left_thigh )
50+ block .append (left_calf )
51+
52+ # 转换为单个UnstructuredGrid
53+ robot = block .combine ()
54+
55+ # 添加关节标记（右肘、左膝）
56+ robot .point_data ["joints" ] = np .zeros (robot .n_points , dtype = int )
57+ # 右肘关节（右上臂+右前臂连接点）
58+ elbow_idx = robot .find_closest_point ((0.7 , 0 , 1.0 ))
59+ robot .point_data ["joints" ][elbow_idx ] = 1
60+ # 左膝关节（左大腿+左小腿连接点）
61+ knee_idx = robot .find_closest_point ((- 0.2 , 0 , - 0.5 ))
62+ robot .point_data ["joints" ][knee_idx ] = 1
63+
64+ if return_parts :
65+ return robot , [torso , head , right_upper_arm , right_forearm , left_upper_arm , left_forearm , right_thigh ,
66+ right_calf , left_thigh , left_calf ]
67+ return robot
68+
69+
70+ # 关节旋转函数（保持逻辑，适配新部件ID）
71+ def rotate_joint (robot , joint_name , angle_deg ):
72+ joint_config = {
73+ "right_elbow" : {
74+ "center" : (0.7 , 0 , 1.0 ), # 右肘关节中心点
75+ "target_part" : 2 # 右前臂ID=2
76+ },
77+ "left_knee" : {
78+ "center" : (- 0.2 , 0 , - 0.5 ), # 左膝关节中心点
79+ "target_part" : 4 # 左小腿ID=4
80+ }
81+ }
82+
83+ if joint_name not in joint_config :
84+ print (f"错误：关节名称{ joint_name } )
85+ return robot
86+
87+ config = joint_config [joint_name ]
88+ part_mask = robot ["part_id" ] == config ["target_part" ]
89+ if not np .any (part_mask ):
90+ print (f"错误：未找到part_id={ config ['target_part' ]} )
91+ return robot
92+
93+ # 旋转逻辑（与之前一致）
94+ target_points = robot .points [part_mask ].copy ()
95+ vtk_points = vtk .vtkPoints ()
96+ for point in target_points :
97+ vtk_points .InsertNextPoint (point )
98+
99+ transform = vtk .vtkTransform ()
100+ transform .Translate (config ["center" ][0 ], config ["center" ][1 ], config ["center" ][2 ])
101+ transform .RotateY (angle_deg )
102+ transform .Translate (- config ["center" ][0 ], - config ["center" ][1 ], - config ["center" ][2 ])
103+
104+ transform_filter = vtk .vtkTransformFilter ()
105+ transform_filter .SetTransform (transform )
106+ poly_data = vtk .vtkPolyData ()
107+ poly_data .SetPoints (vtk_points )
108+ transform_filter .SetInputData (poly_data )
109+ transform_filter .Update ()
110+
111+ rotated_vtk_points = transform_filter .GetOutput ().GetPoints ()
112+ rotated_points = np .array ([rotated_vtk_points .GetPoint (i ) for i in range (rotated_vtk_points .GetNumberOfPoints ())])
113+ robot .points [part_mask ] = rotated_points
114+ return robot
115+
116+
117+ # 本地测试
118+ if __name__ == "__main__" :
119+ robot_initial = create_robot ()
120+ print (f"初始模型：{ robot_initial .n_points } { robot_initial .n_cells } )
121+
122+ # 测试右肘+左膝旋转
123+ robot_rotated = rotate_joint (robot_initial .copy (), "right_elbow" , 90 )
124+ robot_rotated = rotate_joint (robot_rotated , "left_knee" , 60 )
125+
126+ # 渲染对比
127+ plotter = pv .Plotter (shape = (1 , 2 ), window_size = (1200 , 600 ))
128+ # 左窗口：初始姿态
129+ plotter .subplot (0 , 0 )
130+ plotter .add_mesh (
131+ robot_initial ,
132+ scalars = "part_id" ,
133+ cmap = "tab10" , # 更多颜色区分部件
134+ edge_color = "black" ,
135+ show_edges = True ,
136+ scalar_bar_args = {"title" : "部件ID" }
137+ )
138+ plotter .add_axes ()
139+ grid = pv .Plane (center = (0 , 0 , - 1.2 ), i_size = 3 , j_size = 3 , i_resolution = 10 , j_resolution = 10 )
140+ plotter .add_mesh (grid , color = "lightgray" , opacity = 0.5 )
141+ plotter .add_text ("初始人形机器人" , font_size = 14 , position = "upper_left" )
142+
143+ # 右窗口：旋转后姿态
144+ plotter .subplot (0 , 1 )
145+ plotter .add_mesh (
146+ robot_rotated ,
147+ scalars = "part_id" ,
148+ cmap = "tab10" ,
149+ edge_color = "black" ,
150+ show_edges = True ,
151+ scalar_bar_args = {"title" : "部件ID" }
152+ )
153+ plotter .add_axes ()
154+ plotter .add_mesh (grid , color = "lightgray" , opacity = 0.5 )
155+ plotter .add_text ("抬臂+屈膝动作" , font_size = 14 , position = "upper_left" )
156+
157+ plotter .show ()
0 commit comments