# worm_gear_v15_ug_quality_plus_lines_optimized_v2.py
# 1. 补 STL 导入
# 2. 主线程 GLFW + Tkinter 非阻塞
# 3. 错误路径清理 GLFW
# 4. 启用面剔除
# 5. 热更新网格(不重启窗口)
import numpy as np
import math, glm, numba as nb, logging, time, os, glfw, OpenGL.GL as gl, OpenGL.GL.shaders as shaders
import tkinter as tk, tkinter.messagebox as msg, tkinter.filedialog as fd
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
import matplotlib.pyplot as plt
from collections import namedtuple
from typing import Tuple, List, Optional
from stl import mesh # <<< NEW 1
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s: %(message)s')
# ---------------- 参数结构 ----------------
class MachiningError:
__slots__ = ('tool_rx_err_deg', 'tool_ry_err_deg', 'tool_rz_err_deg',
'pitch_err_um', 'tool_dx_um', 'tool_dy_um')
def __init__(self, rx=0.0, ry=0.0, rz=0.0, pitch=0.0, dx=0.0, dy=0.0):
self.tool_rx_err_deg = rx
self.tool_ry_err_deg = ry
self.tool_rz_err_deg = rz
self.pitch_err_um = pitch
self.tool_dx_um = dx
self.tool_dy_um = dy
class WormParams:
__slots__ = ('module', 'teeth', 'length', 'pressure_angle', 'helix_angle',
'backlash', 'root_factor', 'pitch_radius', 'base_radius',
'lead', 'root_radius', 'tip_radius', 'axial_res', 'tooth_res')
def __init__(self, m=2, z=2, L=50, alpha=20, beta=5, bl=0.1, rf=1.0):
self.module, self.teeth, self.length = m, z, L
self.pressure_angle = math.radians(alpha)
self.helix_angle = math.radians(beta)
self.backlash, self.root_factor = bl, rf
self._derive()
def _derive(self):
self.pitch_radius = self.module * self.teeth / 2
self.base_radius = self.pitch_radius * math.cos(self.pressure_angle)
self.lead = 2 * math.pi * self.pitch_radius / math.tan(self.helix_angle)
self.root_radius = self.pitch_radius - 1.25 * self.module * self.root_factor
self.tip_radius = self.pitch_radius + self.module
complexity = max(1, min(5, int(10 / self.module)))
self.axial_res = 100 * complexity
self.tooth_res = 100 * complexity
# ---------------- 几何计算 ----------------
@nb.njit(cache=True, fastmath=True)
def theory_profile(p: WormParams) -> np.ndarray:
rb, ra, rf = p.base_radius, p.tip_radius, p.root_radius
n = p.tooth_res // 2
theta_max = math.sqrt((ra/rb)**2 - 1)
t = np.linspace(0, theta_max, n)
right = np.empty((n, 2), np.float32)
for i in range(n):
ti = t[i]
right[i, 0] = rb * (math.cos(ti) + ti * math.sin(ti))
right[i, 1] = rb * (math.sin(ti) - ti * math.cos(ti))
left = np.empty((n, 2), np.float32)
for i in range(n):
ti = t[i]
left[i, 0] = rb * (math.cos(-ti) - ti * math.sin(-ti))
left[i, 1] = rb * (math.sin(-ti) + ti * math.cos(-ti))
left = left[::-1]
n_root = p.tooth_res // 3
angles = np.linspace(-np.pi/2 + p.pressure_angle, np.pi/2 - p.pressure_angle, n_root)
root = np.empty((n_root, 2), np.float32)
for i in range(n_root):
ang = angles[i]
root[i, 0], root[i, 1] = rf * math.cos(ang), rf * math.sin(ang)
profile = np.vstack((right, root, left))
if p.backlash > 0:
off = np.array([math.cos(p.pressure_angle), math.sin(p.pressure_angle)]) * p.backlash / 2
profile[:len(right)] -= off
profile[-len(left):] += off
return profile
@nb.njit(cache=True, fastmath=True, parallel=True)
def build_mesh_90_norm_color(p: WormParams, profile: np.ndarray, err: MachiningError) -> Tuple:
n_prof = len(profile)
z = np.linspace(0, p.length, p.axial_res, dtype=np.float32)
phase = (2 * np.pi * z / p.lead).astype(np.float32)
offset = (2 * np.pi * np.arange(p.teeth) / p.teeth).astype(np.float32)
rx, ry, rz = math.radians(err.tool_rx_err_deg), math.radians(err.tool_ry_err_deg), math.radians(err.tool_rz_err_deg)
cos_rz, sin_rz = math.cos(rz), math.sin(rz)
max_vertices = p.axial_res * n_prof * p.teeth
vertices = np.empty((max_vertices, 3), np.float32)
normals = np.empty((max_vertices, 3), np.float32)
colors = np.empty((max_vertices, 3), np.float32)
vertex_count = 0
faces = []
for i in nb.prange(p.axial_res):
cos_p, sin_p = math.cos(phase[i]), math.sin(phase[i])
z_val = z[i]
for j in range(n_prof):
x0, y0 = profile[j]
x0 += err.pitch_err_um * 1e-3 * math.cos(phase[i])
y0 += err.pitch_err_um * 1e-3 * math.sin(phase[i])
x0 += err.tool_dx_um * 1e-3; y0 += err.tool_dy_um * 1e-3
new_x = x0 * cos_rz - y0 * sin_rz
new_y = x0 * sin_rz + y0 * cos_rz
x0, y0 = new_x, new_y
for k in range(p.teeth):
cos_t, sin_t = math.cos(offset[k]), math.sin(offset[k])
x = x0*(cos_p*cos_t - sin_p*sin_t) - y0*(cos_p*sin_t + sin_p*cos_t)
y = x0*(sin_p*cos_t + cos_p*sin_t) + y0*(cos_p*cos_t - sin_p*sin_t)
ang = math.atan2(y, x)
if ang < 0: ang += 2*math.pi
if ang <= 3*math.pi/2:
idx = vertex_count
vertices[idx] = (x, y, z_val)
norm_len = math.sqrt(x*x + y*y)
normals[idx] = (x/norm_len, y/norm_len, 0.0) if norm_len > 1e-6 else (0,0,1)
err_val = math.sqrt((x-profile[j,0])**2 + (y-profile[j,1])**2)*1e3
colors[idx] = (min(1.0, err_val/50.0), 0.0, max(0.0, 1.0 - err_val/50.0))
vertex_count += 1
vertices = vertices[:vertex_count]; normals = normals[:vertex_count]; colors = colors[:vertex_count]
n_per_layer = vertex_count // p.axial_res
for i in range(p.axial_res - 1):
start = i*n_per_layer; next_start = (i+1)*n_per_layer
if next_start + n_per_layer > vertex_count: break
for j in range(n_per_layer - 1):
a, b, c, d = start+j, start+j+1, next_start+j, next_start+j+1
if d < vertex_count:
faces.append((a, b, c)); faces.append((b, d, c))
return vertices, normals, colors, np.array(faces, dtype=np.uint32)
def build_overlay_lines(p: WormParams, err: MachiningError) -> Tuple[np.ndarray, np.ndarray]:
prof_perf = theory_profile(p)
prof_err = apply_machining_errors(prof_perf, p, err)
z0 = 0.0
line_perf = np.column_stack([prof_perf[:, 0], prof_perf[:, 1], np.full(len(prof_perf), z0)])
line_err = np.column_stack([prof_err[:, 0], prof_err[:, 1], np.full(len(prof_err), z0)])
return line_perf.astype(np.float32), line_err.astype(np.float32)
@nb.njit(cache=True, fastmath=True)
def apply_machining_errors(profile: np.ndarray, p: WormParams, err: MachiningError) -> np.ndarray:
result = np.empty_like(profile)
rx_rad, ry_rad, rz_rad = math.radians(err.tool_rx_err_deg), math.radians(err.tool_ry_err_deg), math.radians(err.tool_rz_err_deg)
cos_rz, sin_rz = math.cos(rz_rad), math.sin(rz_rad)
for i in range(len(profile)):
x, y = profile[i]
x += err.tool_dx_um * 1e-3; y += err.tool_dy_um * 1e-3
new_x = x * cos_rz - y * sin_rz
new_y = x * sin_rz + y * cos_rz
result[i] = (new_x, new_y)
return result
# ---------------- GLRenderer ----------------
class GLRenderer:
def __init__(self):
self.vao = self.ebo = self.shader = self.line_shader = None
self.face_count = 0
self.line_vao_perf = self.line_vao_err = None
self.line_count_perf = self.line_count_err = 0
def setup_shaders(self):
vert = """#version 330 core
layout(location = 0) in vec3 aPos;
layout(location = 1) in vec3 aNorm;
layout(location = 2) in vec3 aColor;
uniform mat4 MVP; uniform vec3 lightDir;
out vec3 color;
void main(){
gl_Position = MVP * vec4(aPos, 1.0);
vec3 norm = normalize(aNorm);
float diff = max(dot(norm, normalize(lightDir)), 0.0);
color = aColor * (0.3 + 0.7 * diff);
}"""
frag = """#version 330 core
in vec3 color; out vec4 FragColor;
void main(){ FragColor = vec4(color, 1.0); }"""
self.shader = shaders.compileProgram(shaders.compileShader(vert, gl.GL_VERTEX_SHADER),
shaders.compileShader(frag, gl.GL_FRAGMENT_SHADER))
line_vert = """#version 330 core
layout(location = 0) in vec3 aPos; uniform mat4 MVP;
void main(){ gl_Position = MVP * vec4(aPos, 1.0); }"""
line_frag = """#version 330 core
uniform vec3 lineColor; out vec4 FragColor;
void main(){ FragColor = vec4(lineColor, 1.0); }"""
self.line_shader = shaders.compileProgram(shaders.compileShader(line_vert, gl.GL_VERTEX_SHADER),
shaders.compileShader(line_frag, gl.GL_FRAGMENT_SHADER))
def upload_mesh(self, vertices, normals, colors, faces):
if self.vao is None:
self.vao = gl.glGenVertexArrays(1)
gl.glBindVertexArray(self.vao)
# 顶点
if gl.glIsBuffer(gl.GLuint(0)) == gl.GL_FALSE: self.vbo = gl.glGenBuffers(1)
gl.glBindBuffer(gl.GL_ARRAY_BUFFER, self.vbo)
gl.glBufferData(gl.GL_ARRAY_BUFFER, vertices.nbytes, vertices, gl.GL_DYNAMIC_DRAW) # <<< NEW 5
gl.glVertexAttribPointer(0, 3, gl.GL_FLOAT, gl.GL_FALSE, 0, None); gl.glEnableVertexAttribArray(0)
# 法线
if gl.glIsBuffer(gl.GLuint(0)) == gl.GL_FALSE: self.nbo = gl.glGenBuffers(1)
gl.glBindBuffer(gl.GL_ARRAY_BUFFER, self.nbo)
gl.glBufferData(gl.GL_ARRAY_BUFFER, normals.nbytes, normals, gl.GL_DYNAMIC_DRAW)
gl.glVertexAttribPointer(1, 3, gl.GL_FLOAT, gl.GL_FALSE, 0, None); gl.glEnableVertexAttribArray(1)
# 颜色
if gl.glIsBuffer(gl.GLuint(0)) == gl.GL_FALSE: self.cbo = gl.glGenBuffers(1)
gl.glBindBuffer(gl.GL_ARRAY_BUFFER, self.cbo)
gl.glBufferData(gl.GL_ARRAY_BUFFER, colors.nbytes, colors, gl.GL_DYNAMIC_DRAW)
gl.glVertexAttribPointer(2, 3, gl.GL_FLOAT, gl.GL_FALSE, 0, None); gl.glEnableVertexAttribArray(2)
# 索引
if self.ebo is None: self.ebo = gl.glGenBuffers(1)
gl.glBindBuffer(gl.GL_ELEMENT_ARRAY_BUFFER, self.ebo)
gl.glBufferData(gl.GL_ELEMENT_ARRAY_BUFFER, faces.nbytes, faces, gl.GL_DYNAMIC_DRAW)
self.face_count = len(faces)
gl.glBindVertexArray(0)
def upload_lines(self, line_perf, line_err):
# 理论线
if self.line_vao_perf is None: self.line_vao_perf = gl.glGenVertexArrays(1)
gl.glBindVertexArray(self.line_vao_perf)
if gl.glIsBuffer(gl.GLuint(0)) == gl.GL_FALSE: self.line_vbo_perf = gl.glGenBuffers(1)
gl.glBindBuffer(gl.GL_ARRAY_BUFFER, self.line_vbo_perf)
gl.glBufferData(gl.GL_ARRAY_BUFFER, line_perf.nbytes, line_perf, gl.GL_DYNAMIC_DRAW)
gl.glVertexAttribPointer(0, 3, gl.GL_FLOAT, gl.GL_FALSE, 0, None); gl.glEnableVertexAttribArray(0)
self.line_count_perf = len(line_perf)
# 误差线
if self.line_vao_err is None: self.line_vao_err = gl.glGenVertexArrays(1)
gl.glBindVertexArray(self.line_vao_err)
if gl.glIsBuffer(gl.GLuint(0)) == gl.GL_FALSE: self.line_vbo_err = gl.glGenBuffers(1)
gl.glBindBuffer(gl.GL_ARRAY_BUFFER, self.line_vbo_err)
gl.glBufferData(gl.GL_ARRAY_BUFFER, line_err.nbytes, line_err, gl.GL_DYNAMIC_DRAW)
gl.glVertexAttribPointer(0, 3, gl.GL_FLOAT, gl.GL_FALSE, 0, None); gl.glEnableVertexAttribArray(0)
self.line_count_err = len(line_err)
gl.glBindVertexArray(0)
def render(self, mvp, light_dir):
gl.glUseProgram(self.shader)
gl.glUniformMatrix4fv(gl.glGetUniformLocation(self.shader, "MVP"), 1, gl.GL_FALSE, glm.value_ptr(mvp))
gl.glUniform3f(gl.glGetUniformLocation(self.shader, "lightDir"), *light_dir)
gl.glBindVertexArray(self.vao)
gl.glDrawElements(gl.GL_TRIANGLES, self.face_count, gl.GL_UNSIGNED_INT, None)
# 线条
gl.glUseProgram(self.line_shader)
gl.glUniformMatrix4fv(gl.glGetUniformLocation(self.line_shader, "MVP"), 1, gl.GL_FALSE, glm.value_ptr(mvp))
gl.glUniform3f(gl.glGetUniformLocation(self.line_shader, "lineColor"), 1.0, 1.0, 1.0)
gl.glBindVertexArray(self.line_vao_perf); gl.glDrawArrays(gl.GL_LINE_LOOP, 0, self.line_count_perf)
gl.glUniform3f(gl.glGetUniformLocation(self.line_shader, "lineColor"), 1.0, 0.0, 0.0)
gl.glBindVertexArray(self.line_vao_err); gl.glDrawArrays(gl.GL_LINE_LOOP, 0, self.line_count_err)
gl.glBindVertexArray(0)
# ---------------- GLWindow ----------------
class GLWindow:
def __init__(self, params: WormParams, errors: MachiningError):
self.params, self.errors = params, errors
self.renderer = GLRenderer()
self.camera_distance = 80.0
self.camera_rotation = glm.vec2(0.0, 0.0)
self.last_mouse_pos = None
self.window = None
self._init_glfw()
self._init_scene()
def _init_glfw(self):
if not glfw.init():
msg.showerror("初始化失败", "GLFW 初始化失败,请检查显卡驱动")
raise RuntimeError("GLFW init failed") # <<< NEW 3
glfw.window_hint(glfw.SAMPLES, 8)
glfw.window_hint(glfw.CONTEXT_VERSION_MAJOR, 3)
glfw.window_hint(glfw.CONTEXT_VERSION_MINOR, 3)
glfw.window_hint(glfw.OPENGL_PROFILE, glfw.OPENGL_CORE_PROFILE)
self.window = glfw.create_window(1400, 900, "UG级蜗杆3D - 90°切除+误差热图+叠加线", None, None)
if not self.window:
glfw.terminate() # <<< NEW 3
raise RuntimeError("窗口创建失败")
glfw.make_context_current(self.window)
glfw.set_cursor_pos_callback(self.window, self._on_mouse_move)
glfw.set_scroll_callback(self.window, self._on_scroll)
glfw.set_key_callback(self.window, self._on_key)
gl.glEnable(gl.GL_DEPTH_TEST)
gl.glEnable(gl.GL_MULTISAMPLE)
gl.glEnable(gl.GL_CULL_FACE); gl.glCullFace(gl.GL_BACK) # <<< NEW 4
gl.glClearColor(0.05, 0.05, 0.05, 1.0)
def _init_scene(self):
self.renderer.setup_shaders()
self.update_mesh() # <<< NEW 5
def update_mesh(self): # <<< NEW 5
profile = theory_profile(self.params)
vertices, normals, colors, faces = build_mesh_90_norm_color(self.params, profile, self.errors)
self.renderer.upload_mesh(vertices, normals, colors, faces)
line_perf, line_err = build_overlay_lines(self.params, self.errors)
self.renderer.upload_lines(line_perf, line_err)
def _get_mvp_matrix(self):
proj = glm.perspective(glm.radians(45.0), 1400/900, 0.1, 1000.0)
eye = glm.vec3(0, 0, self.camera_distance)
rot_x = glm.rotate(glm.mat4(1.0), self.camera_rotation.x, glm.vec3(1, 0, 0))
rot_y = glm.rotate(rot_x, self.camera_rotation.y, glm.vec3(0, 1, 0))
eye = glm.vec3(rot_y * glm.vec4(eye, 1.0))
return proj * glm.lookAt(eye, glm.vec3(0), glm.vec3(0, 1, 0))
def _on_mouse_move(self, window, x, y):
if self.last_mouse_pos is not None and glfw.get_mouse_button(window, glfw.MOUSE_BUTTON_LEFT) == glfw.PRESS:
dx, dy = x - self.last_mouse_pos[0], y - self.last_mouse_pos[1]
self.camera_rotation.x += dy * 0.01
self.camera_rotation.y += dx * 0.01
self.last_mouse_pos = (x, y)
def _on_scroll(self, window, dx, dy):
self.camera_distance = max(10.0, min(200.0, self.camera_distance - dy))
def _on_key(self, window, key, scancode, action, mods):
if action == glfw.PRESS and key == glfw.KEY_R:
self.camera_distance = 80.0; self.camera_rotation = glm.vec2(0.0, 0.0)
elif action == glfw.PRESS and key == glfw.KEY_S:
self._screenshot()
def _screenshot(self):
width, height = glfw.get_window_size(self.window)
gl.glPixelStorei(gl.GL_PACK_ALIGNMENT, 1)
data = gl.glReadPixels(0, 0, width, height, gl.GL_RGB, gl.GL_UNSIGNED_BYTE)
os.makedirs("screenshots", exist_ok=True)
import imageio, time
filename = f"screenshots/worm_gear_{time.strftime('%Y%m%d_%H%M%S')}.png"
image = np.flipud(np.frombuffer(data, dtype=np.uint8).reshape(height, width, 3))
imageio.imwrite(filename, image)
logging.info(f"截图已保存: {filename}")
def main_loop(self):
while not glfw.window_should_close(self.window):
glfw.poll_events()
gl.glClear(gl.GL_COLOR_BUFFER_BIT | gl.GL_DEPTH_BUFFER_BIT)
mvp = self._get_mvp_matrix()
light_dir = glm.normalize(glm.vec3(0.5, 0.5, 1.0))
self.renderer.render(mvp, light_dir)
glfw.swap_buffers(self.window)
glfw.terminate()
def shutdown(self): # <<< NEW 5
if self.window:
glfw.set_window_should_close(self.window, True)
# ---------------- Tkinter 控制 ----------------
class ControlPanel:
def __init__(self, root):
self.root = root
self.root.title("蜗杆参数控制")
self.params, self.errors = WormParams(), MachiningError()
self.gl_window = None # <<< NEW 5
self._build_ui()
def _build_ui(self):
# 参数区
pf = tk.LabelFrame(self.root, text="蜗杆参数"); pf.pack(fill=tk.X, padx=10, pady=5)
self.param_vars = {}
for i, (lab, attr, minv, maxv, step) in enumerate([
("模数 (mm)", "module", 1.0, 10.0, 0.1),
("齿数", "teeth", 1, 50, 1),
("长度 (mm)", "length", 10, 200, 1),
("压力角 (°)", "pressure_angle", 10, 30, 1),
("螺旋角 (°)", "helix_angle", 1, 20, 1),
("齿侧间隙 (mm)", "backlash", 0.0, 0.5, 0.01),
("齿根系数", "root_factor", 0.8, 1.4, 0.05)]):
fr = tk.Frame(pf); fr.grid(row=i//2, column=i%2, sticky="ew", padx=5, pady=2)
tk.Label(fr, text=lab, width=12).pack(side=tk.LEFT)
var = tk.DoubleVar(value=getattr(self.params, attr))
tk.Spinbox(fr, from_=minv, to=maxv, increment=step, textvariable=var, width=8).pack(side=tk.RIGHT)
self.param_vars[attr] = var
# 误差区
ef = tk.LabelFrame(self.root, text="加工误差"); ef.pack(fill=tk.X, padx=10, pady=5)
self.error_vars = {}
for i, (lab, attr, minv, maxv, step) in enumerate([
("X向误差 (μm)", "tool_dx_um", -50, 50, 1),
("Y向误差 (μm)", "tool_dy_um", -50, 50, 1),
("RX误差 (°)", "tool_rx_err_deg", -5, 5, 0.1),
("RY误差 (°)", "tool_ry_err_deg", -5, 5, 0.1),
("RZ误差 (°)", "tool_rz_err_deg", -5, 5, 0.1),
("螺距误差 (μm)", "pitch_err_um", -50, 50, 1)]):
fr = tk.Frame(ef); fr.grid(row=i//2, column=i%2, sticky="ew", padx=5, pady=2)
tk.Label(fr, text=lab, width=12).pack(side=tk.LEFT)
var = tk.DoubleVar(value=getattr(self.errors, attr))
tk.Spinbox(fr, from_=minv, to=maxv, increment=step, textvariable=var, width=8).pack(side=tk.RIGHT)
self.error_vars[attr] = var
# 按钮区
bf = tk.Frame(self.root); bf.pack(fill=tk.X, padx=10, pady=10)
tk.Button(bf, text="更新模型", command=self.update_model).pack(side=tk.LEFT, padx=5)
tk.Button(bf, text="重置参数", command=self.reset_params).pack(side=tk.LEFT, padx=5)
tk.Button(bf, text="导出STL", command=self.export_stl).pack(side=tk.LEFT, padx=5)
tk.Button(bf, text="退出", command=self.root.quit).pack(side=tk.RIGHT, padx=5)
def update_model(self):
for attr, var in self.param_vars.items():
setattr(self.params, attr, var.get())
for attr, var in self.error_vars.items():
setattr(self.errors, attr, var.get())
self.params._derive()
if self.gl_window is None: # 首次创建
self.gl_window = GLWindow(self.params, self.errors)
self.root.after(100, self._tk_poll_gl) # 非阻塞轮询
else: # 热更新
self.gl_window.params = self.params
self.gl_window.errors = self.errors
self.gl_window.update_mesh()
def _tk_poll_gl(self): # <<< NEW 2
if self.gl_window and glfw.window_should_close(self.gl_window.window):
self.gl_window.shutdown()
self.gl_window = None
if self.gl_window:
glfw.poll_events()
self.root.after(100, self._tk_poll_gl)
def reset_params(self):
self.params, self.errors = WormParams(), MachiningError()
for attr, var in {**self.param_vars, **self.error_vars}.items():
var.set(getattr(self.params if attr in self.param_vars else self.errors, attr))
def export_stl(self):
profile = theory_profile(self.params)
vertices, normals, colors, faces = build_mesh_90_norm_color(self.params, profile, self.errors)
stl_mesh = mesh.Mesh(np.zeros(faces.shape[0], dtype=mesh.Mesh.dtype))
for i, face in enumerate(faces):
stl_mesh.vectors[i] = vertices[face]
filename = fd.asksaveasfilename(defaultextension=".stl", filetypes=[("STL文件", "*.stl")], title="保存蜗杆STL文件")
if filename:
stl_mesh.save(filename)
msg.showinfo("导出成功", f"STL文件已保存到: {filename}")
# ---------------- main ----------------
def main():
root = tk.Tk()
ControlPanel(root)
root.mainloop()
if __name__ == "__main__":
main()