|
@@ -0,0 +1,104 @@
|
|
|
|
|
+import os
|
|
|
|
|
+import vtk
|
|
|
|
|
+import pyvista as pv
|
|
|
|
|
+
|
|
|
|
|
+# -----------------------------------------------------------------------------
|
|
|
|
|
+# 1. 配置:要展示的文件夹路径(改成你自己的)
|
|
|
|
|
+# -----------------------------------------------------------------------------
|
|
|
|
|
+ROOT_FOLDER = r"./" # 示例:当前目录
|
|
|
|
|
+MAX_DEPTH = 3 # 限制深度,防止太密
|
|
|
|
|
+NODE_SPACING_X = 2.0
|
|
|
|
|
+NODE_SPACING_Y = 1.5
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+# -----------------------------------------------------------------------------
|
|
|
|
|
+# 2. 递归遍历文件夹,生成树形结构数据
|
|
|
|
|
+# -----------------------------------------------------------------------------
|
|
|
|
|
+class TreeNode:
|
|
|
|
|
+ def __init__(self, name, path, depth=0):
|
|
|
|
|
+ self.name = name
|
|
|
|
|
+ self.path = path
|
|
|
|
|
+ self.depth = depth
|
|
|
|
|
+ self.children = []
|
|
|
|
|
+ self.x = 0.0
|
|
|
|
|
+ self.y = 0.0
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+def build_tree(folder, depth=0, max_depth=MAX_DEPTH):
|
|
|
|
|
+ node = TreeNode(os.path.basename(folder) or folder, folder, depth)
|
|
|
|
|
+ if depth >= max_depth:
|
|
|
|
|
+ return node
|
|
|
|
|
+
|
|
|
|
|
+ try:
|
|
|
|
|
+ entries = sorted(os.listdir(folder))
|
|
|
|
|
+ except PermissionError:
|
|
|
|
|
+ return node
|
|
|
|
|
+
|
|
|
|
|
+ for entry in entries:
|
|
|
|
|
+ full = os.path.join(folder, entry)
|
|
|
|
|
+ if os.path.isdir(full):
|
|
|
|
|
+ child = build_tree(full, depth + 1, max_depth)
|
|
|
|
|
+ node.children.append(child)
|
|
|
|
|
+ return node
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+# -----------------------------------------------------------------------------
|
|
|
|
|
+# 3. 树布局算法(简单层级横向布局)
|
|
|
|
|
+# -----------------------------------------------------------------------------
|
|
|
|
|
+def layout_tree(node: TreeNode, x=0.0, y=0.0):
|
|
|
|
|
+ node.y = y
|
|
|
|
|
+ node.x = x
|
|
|
|
|
+
|
|
|
|
|
+ child_x = x - (len(node.children) - 1) * NODE_SPACING_X / 2
|
|
|
|
|
+ for child in node.children:
|
|
|
|
|
+ layout_tree(child, child_x, y - NODE_SPACING_Y)
|
|
|
|
|
+ child_x += NODE_SPACING_X
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+# -----------------------------------------------------------------------------
|
|
|
|
|
+# 4. 收集所有节点、边用于 VTK 绘制
|
|
|
|
|
+# -----------------------------------------------------------------------------
|
|
|
|
|
+def collect_geometry(node):
|
|
|
|
|
+ points = []
|
|
|
|
|
+ lines = []
|
|
|
|
|
+ labels = []
|
|
|
|
|
+ stack = [node]
|
|
|
|
|
+
|
|
|
|
|
+ while stack:
|
|
|
|
|
+ n = stack.pop()
|
|
|
|
|
+ points.append([n.x, n.y, 0.0])
|
|
|
|
|
+ labels.append(n.name)
|
|
|
|
|
+ for child in n.children:
|
|
|
|
|
+ points.append([child.x, child.y, 0.0])
|
|
|
|
|
+ labels.append(child.name)
|
|
|
|
|
+ lines.append(2) # 每个线段有两个点
|
|
|
|
|
+ lines.append(len(points)-2) # 父节点的索引
|
|
|
|
|
+ lines.append(len(points)-1) # 子节点的索引
|
|
|
|
|
+ stack.append(child)
|
|
|
|
|
+ return points, lines, labels
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+# -----------------------------------------------------------------------------
|
|
|
|
|
+# 5. PyVista + VTK 绘制树
|
|
|
|
|
+# -----------------------------------------------------------------------------
|
|
|
|
|
+if __name__ == "__main__":
|
|
|
|
|
+ tree = build_tree(ROOT_FOLDER)
|
|
|
|
|
+ layout_tree(tree)
|
|
|
|
|
+ points, lines, labels = collect_geometry(tree)
|
|
|
|
|
+
|
|
|
|
|
+ # 构建点线云
|
|
|
|
|
+ cloud = pv.PolyData(points)
|
|
|
|
|
+ cloud.lines = lines
|
|
|
|
|
+ cloud["label"] = labels
|
|
|
|
|
+
|
|
|
|
|
+ # 绘图
|
|
|
|
|
+ plotter = pv.Plotter(window_size=(1200, 800))
|
|
|
|
|
+ plotter.add_mesh(cloud, color="#66ccff", point_size=12, render_lines_as_tubes=True, line_width=3)
|
|
|
|
|
+ plotter.add_point_labels(cloud, "label", font_size=8, shape_color="white")
|
|
|
|
|
+
|
|
|
|
|
+ plotter.add_text(
|
|
|
|
|
+ f"Folder Tree: {ROOT_FOLDER}\nMax Depth: {MAX_DEPTH}",
|
|
|
|
|
+ font_size=10, color="blue"
|
|
|
|
|
+ )
|
|
|
|
|
+ plotter.camera_position = "xy"
|
|
|
|
|
+ plotter.show()
|