TellmeStpToGlb/core/converter.py

#!/usr/bin/env python3
"""
STP到GLB转换器类
封装原有转换逻辑，支持进度回调
"""

import os
import trimesh
from typing import Callable, Optional
from OCC.Extend.DataExchange import read_step_file, write_stl_file


class StpToGlbConverter:
    """STP到GLB转换器"""
    
    def __init__(self):
        self.progress_callback: Optional[Callable[[int, str], None]] = None
    
    def set_progress_callback(self, callback: Callable[[int, str], None]) -> None:
        """设置进度回调函数"""
        self.progress_callback = callback
    
    def _update_progress(self, progress: int, message: str) -> None:
        """更新进度"""
        if self.progress_callback:
            self.progress_callback(progress, message)
    
    def stp_to_stl(self, stp_path: str, stl_path: str) -> None:
        """pythonocc 读 STP 写二进制 STL"""
        self._update_progress(5, "检查输入文件...")
        
        if not os.path.isfile(stp_path):
            raise FileNotFoundError(f"找不到文件：{stp_path}")
        
        self._update_progress(20, "读取STP文件...")
        shape = read_step_file(stp_path)
        
        if shape is None:
            raise ValueError(f"无法读取STP文件：{stp_path}")
        
        self._update_progress(40, "转换为STL格式...")
        write_stl_file(shape, stl_path, mode="binary", linear_deflection=0.01, angular_deflection=0.1)
        
        self._update_progress(60, "STL文件生成完成")
    
    def stl_to_glb(self, stl_path: str, glb_path: str, auto_scale: bool = True, auto_center: bool = True) -> None:
        """使用 trimesh 转换 STL 到 GLB，并优化用于Blender"""
        self._update_progress(65, "检查STL文件...")
        
        if not os.path.isfile(stl_path):
            raise FileNotFoundError(f"找不到STL文件：{stl_path}")
        
        self._update_progress(70, "加载STL网格...")
        mesh = trimesh.load(stl_path)
        
        if mesh.is_empty:
            raise ValueError(f"STL文件为空或无效：{stl_path}")
        
        if len(mesh.vertices) == 0 or len(mesh.faces) == 0:
            raise ValueError("网格数据无效")
        
        self._update_progress(80, "优化网格...")
        
        # 自动缩放和居中
        if auto_scale or auto_center:
            bounds = mesh.bounds
            max_dimension = max(bounds[1] - bounds[0])
            
            if auto_scale and max_dimension > 1000:
                scale_factor = 50.0 / max_dimension
                mesh.apply_scale(scale_factor)
            
            if auto_center:
                mesh.apply_translation(-mesh.centroid)
        
        self._update_progress(90, "导出GLB文件...")
        
        try:
            mesh.export(glb_path)
            if os.path.getsize(glb_path) == 0:
                raise ValueError("生成的GLB文件为空")
        except Exception as e:
            raise ValueError(f"GLB导出失败：{e}")
        
        self._update_progress(100, "转换完成")
    
    def convert(self, stp_path: str, glb_path: str, auto_scale: bool = True, auto_center: bool = True) -> None:
        """完整的STP到GLB转换流程"""
        stl_tmp = os.path.splitext(glb_path)[0] + ".tmp.stl"
        
        try:
            self._update_progress(0, "开始转换...")
            self.stp_to_stl(stp_path, stl_tmp)
            self.stl_to_glb(stl_tmp, glb_path, auto_scale, auto_center)
        finally:
            # 清理临时文件
            if os.path.isfile(stl_tmp):
                os.remove(stl_tmp)
    
    def _split_mesh_components(self, stl_path: str, auto_scale: bool = True, 
                              auto_center: bool = True) -> trimesh.Scene:
        """基于trimesh分离连通组件来模拟层级结构"""
        self._update_progress(70, "加载STL文件...")
        
        if not os.path.isfile(stl_path):
            raise FileNotFoundError(f"找不到STL文件：{stl_path}")
        
        # 加载STL文件
        mesh = trimesh.load(stl_path)
        
        if mesh.is_empty:
            raise ValueError(f"STL文件为空或无效：{stl_path}")
        
        self._update_progress(75, "分析网格连通组件...")
        
        # 尝试分离连通组件
        try:
            # 分离不连通的网格部分
            components = mesh.split(only_watertight=False)
            
            if len(components) > 1:
                self._update_progress(80, f"发现 {len(components)} 个分离组件")
                
                # 创建场景，每个组件作为独立零件
                scene = trimesh.Scene()
                
                for i, component in enumerate(components):
                    if not component.is_empty and len(component.vertices) > 0:
                        part_name = f'Part_{i:03d}'
                        scene.add_geometry(component, node_name=part_name)
                        self._update_progress(80 + int(5 * i / len(components)), 
                                            f"添加组件 {part_name}")
                
                # 对整个场景进行优化
                if auto_scale or auto_center:
                    self._update_progress(87, "优化场景...")
                    bounds = scene.bounds
                    if bounds is not None:
                        max_dimension = max(bounds[1] - bounds[0])
                        
                        if auto_scale and max_dimension > 1000:
                            scale_factor = 50.0 / max_dimension
                            scene.apply_scale(scale_factor)
                        
                        if auto_center:
                            scene.apply_translation(-scene.centroid)
                
                return scene
            else:
                self._update_progress(80, "未发现分离组件，使用单一网格")
                return None
                
        except Exception as e:
            self._update_progress(80, f"组件分离失败: {e}")
            return None
    
    def convert_with_hierarchy(self, stp_path: str, glb_path: str, 
                              auto_scale: bool = True, auto_center: bool = True) -> None:
        """完整的STP到GLB转换流程，尝试保持组件分离结构"""
        stl_tmp = os.path.splitext(glb_path)[0] + ".tmp.stl"
        
        try:
            self._update_progress(0, "开始层级转换...")
            
            # 首先执行标准的STP到STL转换
            self.stp_to_stl(stp_path, stl_tmp)
            
            # 尝试分离网格组件
            self._update_progress(65, "尝试分离网格组件...")
            scene = self._split_mesh_components(stl_tmp, auto_scale, auto_center)
            
            if scene is not None:
                # 成功分离组件，导出场景
                self._update_progress(90, "导出分离组件GLB文件...")
                
                try:
                    scene.export(glb_path)
                    if os.path.getsize(glb_path) == 0:
                        raise ValueError("生成的GLB文件为空")
                    
                    # 计算组件数量
                    component_count = len([name for name in scene.geometry.keys()])
                    self._update_progress(100, f"层级转换完成，包含 {component_count} 个分离组件")
                    
                except Exception as e:
                    raise ValueError(f"GLB导出失败：{e}")
            else:
                # 没有发现分离组件，回退到标准转换
                self._update_progress(85, "未发现分离组件，使用标准转换...")
                self.stl_to_glb(stl_tmp, glb_path, auto_scale, auto_center)
                
        except Exception as e:
            # 如果层级转换失败，提供回退选项
            self._update_progress(50, f"层级转换失败: {e}，尝试标准转换...")
            try:
                self.convert(stp_path, glb_path, auto_scale, auto_center)
            except Exception as fallback_error:
                raise ValueError(f"转换失败: {fallback_error}")
        finally:
            # 清理临时文件
            if os.path.isfile(stl_tmp):
                os.remove(stl_tmp)