Browse Source

refactor: config

main
mckay 3 months ago
parent
commit
130735bae1
  1. 9
      brep2sdf/config/default_config.py
  2. 6
      brep2sdf/scripts/process_brep.py
  3. 5
      brep2sdf/scripts/process_furniture.py

9
brep2sdf/config/default_config.py

@ -6,7 +6,7 @@ class ModelConfig:
"""模型相关配置"""
brep_feature_dim: int = 16
use_cf: bool = True
embed_dim: int = 192 # 3 的 倍数
embed_dim: int = 384 # 3 的 倍数
latent_dim: int = 16
# 点云采样配置
@ -27,10 +27,13 @@ class ModelConfig:
@dataclass
class DataConfig:
"""数据相关配置"""
max_face: int = 8
max_edge: int = 16
max_face: int = 32
max_edge: int = 128
num_query_points: int = 4096 # 限制查询点数量,sdf 采样点数 本来是 128*128*128 ,在data load时随机采样
bbox_scaled: float = 1.0
# pre_process
origin_brep_dir: str = '/mnt/mynewdisk/dataset/furniture/step/furniture_dataset_step/'
# 数据路径
brep_dir: str = '/home/wch/brep2sdf/test_data/pkl'

6
brep2sdf/scripts/process_brep.py

@ -492,9 +492,9 @@ def process_furniture_step(data_path):
def main():
"""主函数:处理多个STEP文件"""
# 定义路径常量
INPUT = '/mnt/disk2/dataset/furniture/step/furniture_dataset_step/'
OUTPUT = '../test_data/pkl/'
RESULT = '../test_data/result/pkl/' # 用于存储成功/失败文件记录
INPUT = config.data.origin_brep_dir
OUTPUT = config.data.brep_dir
RESULT = config.data.valid_data_dir # 用于存储成功/失败文件记录
# 清理输出目录
def clean_directory(directory):

5
brep2sdf/scripts/process_furniture.py

@ -16,6 +16,9 @@ from OCC.Core.BRepMesh import BRepMesh_IncrementalMesh
from OCC.Core.StlAPI import StlAPI_Writer
from OCC.Core.IFSelect import IFSelect_RetDone
from brep2sdf.config.default_config import get_default_config
config = get_default_config()
# 配置日志记录
def setup_logger():
logger = logging.getLogger('furniture_processing')
@ -244,7 +247,7 @@ def process(step_file: str, set_name:str) -> bool:
def main():
"""主函数:并行处理所有STEP文件"""
INPUT = '/mnt/disk2/dataset/furniture/step/furniture_dataset_step'
INPUT = config.data.origin_brep_dir
if not os.path.exists(INPUT):
logger.error(f"输入目录不存在: {INPUT}")
return

Loading…
Cancel
Save