体素化

from OCC.Core.STEPControl import STEPControl_Reader
from OCC.Core.IFSelect import IFSelect_RetDone
from OCC.Core.BRepPrimAPI import BRepPrimAPI_MakeBox, BRepPrimAPI_MakeCylinder
from OCC.Core._BRepBndLib import brepbndlib_Add
from OCC.Core.gp import gp_Pnt, gp_Vec, gp_Trsf
from OCC.Core.BRepBuilderAPI import BRepBuilderAPI_Transform
from OCC.Core.BRepClass3d import BRepClass3d_SolidClassifier
from OCC.Core.TopAbs import TopAbs_IN
from OCC.Core.Bnd import Bnd_Box
from OCC.Display.SimpleGui import init_display
from OCC.Core.Quantity import Quantity_Color, Quantity_TOC_RGB
import math
import random
import time


# 导入STEP文件
def import_step(file_path):
    step_reader = STEPControl_Reader()
    status = step_reader.ReadFile(file_path)
    if status != IFSelect_RetDone:
        raise Exception(f"STEP 文件读取失败, code: {status}")
    step_reader.TransferRoots()
    shape = step_reader.Shape()
    return shape


# 获取shape的bounding box的min和max坐标
def get_bounding_box(shape):
    bbox = Bnd_Box()
    brepbndlib_Add(shape, bbox)
    xmin, ymin, zmin, xmax, ymax, zmax = bbox.Get()
    return xmin, ymin, zmin, xmax, ymax, zmax


# 生成立方体毛坯
def make_box(pnt, dx, dy, dz):
    box = BRepPrimAPI_MakeBox(pnt, dx, dy, dz).Shape()
    return box


# 生成圆柱毛坯
def make_cylinder(radius, height):
    cyl = BRepPrimAPI_MakeCylinder(radius, height).Shape()
    return cyl


# 生成体素小立方体用于可视化
def make_voxel(p, size):
    half_size = size / 2.0
    voxel_box = BRepPrimAPI_MakeBox(
        gp_Pnt(p.X() - half_size, p.Y() - half_size, p.Z() - half_size),
        size, size, size
    ).Shape()
    return voxel_box


# 判断点是否在形状内部
def is_point_inside_shape(shape, point):
    classifier = BRepClass3d_SolidClassifier(shape, point, 1e-6)
    return classifier.State() == TopAbs_IN


if __name__ == "__main__":
    # 参数设置
    file_path = r"D:\Desktop\Paper_TurnMilling\code\database\plate.step"
    offset = 2.0  # 体素毛坯轴半径偏移2mm
    voxel_size = 1.0  # 体素大小（mm）
    blank_type = "cube"  # 可选 "cube" 或 "cylinder"

    # 采样率 - 控制体素数量，值越大体素越少
    sampling_rate = 2  # 每个方向上的采样率

    # 显示选项
    show_part = True
    show_blank = True
    show_part_voxels = True
    show_blank_voxels = True

    # 初始化显示
    display, start_display, add_menu, add_function_to_menu = init_display()

    print("正在加载STEP文件...")
    # 1. 导入STEP文件
    part_shape = import_step(file_path)
    xmin, ymin, zmin, xmax, ymax, zmax = get_bounding_box(part_shape)
    print(f"零件边界盒: X[{xmin:.2f}, {xmax:.2f}], Y[{ymin:.2f}, {ymax:.2f}], Z[{zmin:.2f}, {zmax:.2f}]")

    # 2. 生成毛坯
    print("正在生成毛坯...")
    if blank_type == "cube":
        blank_dx = xmax - xmin + 2 * offset
        blank_dy = ymax - ymin + 2 * offset
        blank_dz = zmax - zmin + 2 * offset
        blank_shape = make_box(
            gp_Pnt(xmin - offset, ymin - offset, zmin - offset),
            blank_dx, blank_dy, blank_dz
        )


        def is_inside_blank(p):
            return ((xmin - offset) <= p.X() <= (xmax + offset) and
                    (ymin - offset) <= p.Y() <= (ymax + offset) and
                    (zmin - offset) <= p.Z() <= (zmax + offset))
    elif blank_type == "cylinder":
        cx = (xmin + xmax) / 2
        cy = (ymin + ymax) / 2
        blank_height = zmax - zmin + 2 * offset
        blank_radius = max((xmax - xmin) / 2, (ymax - ymin) / 2) + offset
        blank_cyl = make_cylinder(blank_radius, blank_height)
        trsf = gp_Trsf()
        trsf.SetTranslation(gp_Vec(cx, cy, zmin - offset))
        blank_shape = BRepBuilderAPI_Transform(blank_cyl, trsf).Shape()


        def is_inside_blank(p):
            dx = p.X() - cx
            dy = p.Y() - cy
            return (dx * dx + dy * dy <= blank_radius * blank_radius and
                    (zmin - offset) <= p.Z() <= (zmax + offset))
    else:
        raise ValueError("Invalid blank_type: choose 'cube' or 'cylinder'")

    # 3. 生成体素毛坯和体素零件
    print("正在生成体素模型...")
    # 定义扩展的边界框（考虑偏移）
    extended_min = [xmin - offset, ymin - offset, zmin - offset]
    extended_max = [xmax + offset, ymax + offset, zmax + offset]
    nx = int(math.ceil((extended_max[0] - extended_min[0]) / voxel_size))
    ny = int(math.ceil((extended_max[1] - extended_min[1]) / voxel_size))
    nz = int(math.ceil((extended_max[2] - extended_min[2]) / voxel_size))

    print(f"体素网格大小: {nx} x {ny} x {nz}")
    print(f"总体素数量: {nx * ny * nz}")
    print(f"采样率: {sampling_rate} (每个方向)")

    # 存储体素
    part_voxels = []
    blank_voxels = []

    # 计数器
    total_voxels = 0
    part_voxel_count = 0
    blank_voxel_count = 0

    start_time = time.time()

    # 生成体素并判断它们是否在零件或毛坯内部
    for i in range(0, nx, sampling_rate):
        for j in range(0, ny, sampling_rate):
            for k in range(0, nz, sampling_rate):
                # 计算体素中心点坐标
                x = extended_min[0] + (i + 0.5) * voxel_size
                y = extended_min[1] + (j + 0.5) * voxel_size
                z = extended_min[2] + (k + 0.5) * voxel_size
                p = gp_Pnt(x, y, z)

                total_voxels += 1

                # 检查是否在零件内部
                if is_point_inside_shape(part_shape, p):
                    part_voxels.append(p)
                    part_voxel_count += 1
                # 检查是否在毛坯内部但不在零件内部
                elif is_inside_blank(p):
                    blank_voxels.append(p)
                    blank_voxel_count += 1

    end_time = time.time()
    print(f"体素生成完成，耗时: {end_time - start_time:.2f}秒")
    print(f"零件体素数量: {part_voxel_count}")
    print(f"毛坯体素数量(不包括零件): {blank_voxel_count}")

    # 4. 显示模型
    print("正在显示模型...")

    # 显示原始零件
    if show_part:
        display.DisplayShape(part_shape, color="blue", transparency=0.7)

    # 显示毛坯
    if show_blank:
        display.DisplayShape(blank_shape, color="yellow", transparency=0.9)

    # 显示零件体素
    if show_part_voxels:
        part_color = Quantity_Color(0.2, 0.6, 1.0, Quantity_TOC_RGB)  # 蓝色
        for p in part_voxels:
            voxel = make_voxel(p, voxel_size * 0.95)  # 稍微缩小以便观察
            display.DisplayShape(voxel, color=part_color)

    # 显示毛坯体素（不包括零件体素）
    if show_blank_voxels:
        blank_color = Quantity_Color(1.0, 0.8, 0.0, Quantity_TOC_RGB)  # 黄色
        for p in blank_voxels:
            voxel = make_voxel(p, voxel_size * 0.95)  # 稍微缩小以便观察
            display.DisplayShape(voxel, color=blank_color)

    # 设置视图
    display.View_Iso()
    display.FitAll()

    print("显示完成，请旋转查看模型")
    start_display()