ai-cad-quant/experiments/dxf-parser/main.py

"""DXF 解析工具 —— 读取 DXF 文件，输出图层分析和实体分析结果。"""

import json
from datetime import datetime
from pathlib import Path
import ezdxf

from layer_analysis import analyze_layers
from entity_analysis import analyze_entities
from element_detector import ElementDetector

# 路径配置
BASE_DIR = Path(__file__).parent.parent.parent
SAMPLE_DIR = BASE_DIR / "samples"
OUTPUT_DIR = Path(__file__).parent / "output"

# 要解析的 DXF 文件
DXF_FILE = SAMPLE_DIR / "洗浴中心C-48.dxf"


def build_file_info(doc, filename: str) -> dict:
    """提取 DXF 文件基本信息。"""
    info = {
        "filename": filename,
        "dxf_version": doc.dxfversion,
        "encoding": doc.encoding,
    }

    # 单位
    units = doc.header.get("$INSUNITS", None)
    info["units"] = units

    # 图形范围
    extmin = doc.header.get("$EXTMIN", None)
    extmax = doc.header.get("$EXTMAX", None)
    if extmin and extmax:
        info["extents"] = {"min": list(extmin), "max": list(extmax)}
    else:
        info["extents"] = None

    return info


def print_summary(file_info: dict, layers: list, entities: dict):
    """终端打印解析摘要。"""
    print(f"\n{'=' * 60}")
    print(f"  CAD 文件解析报告")
    print(f"{'=' * 60}")
    print(f"  文件     : {file_info['filename']}")
    print(f"  版本     : {file_info['dxf_version']}")
    print(f"  编码     : {file_info['encoding']}")
    print(f"  实体总数 : {entities['total_entities']}")
    if file_info.get("extents"):
        print(f"  图形范围 : {file_info['extents']}")

    # 图层摘要
    print(f"\n  {'─' * 56}")
    print(f"  图层分析（共 {len(layers)} 个图层）:")
    print(f"  {'名称':<25} {'实体数':<8} {'角色推断'}")
    print(f"  {'-' * 50}")
    for layer in layers[:20]:  # 最多显示前 20 个
        print(f"  {layer['name']:<25} {layer['entity_count']:<8} {layer['suggested_role']}")
    if len(layers) > 20:
        print(f"  ... 还有 {len(layers) - 20} 个图层")

    # 实体类型统计
    type_counts = entities["type_counts"]
    print(f"\n  {'─' * 56}")
    print(f"  实体类型统计:")
    print(f"  {'类型':<25} {'数量':<8}")
    print(f"  {'-' * 35}")
    for dtype, count in sorted(type_counts.items(), key=lambda x: -x[1]):
        print(f"  {dtype:<25} {count:<8}")

    # 块统计
    blocks = entities["blocks"]
    if blocks:
        print(f"\n  {'─' * 56}")
        print(f"  块引用统计（共 {len(blocks)} 种块）:")
        print(f"  {'块名':<35} {'引用次数':<10}")
        print(f"  {'-' * 47}")
        for blk in blocks[:15]:
            print(f"  {blk['block_name']:<35} {blk['count']:<10}")
        if len(blocks) > 15:
            print(f"  ... 还有 {len(blocks) - 15} 种块")

    # 文本
    texts = entities["texts"]
    if texts:
        print(f"\n  {'─' * 56}")
        print(f"  文本内容（共 {len(texts)} 条）:")
        for t in texts[:10]:
            print(f"  [{t['layer']}] {t['content'][:60]}")
        if len(texts) > 10:
            print(f"  ... 还有 {len(texts) - 10} 条文本")

    print(f"\n{'=' * 60}")
    print(f"  解析完成")
    print(f"{'=' * 60}\n")


def main():
    if not DXF_FILE.exists():
        print(f"错误：文件不存在 -> {DXF_FILE}")
        return

    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    parse_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

    print(f"解析文件: {DXF_FILE.name}")
    print(f"解析时间: {parse_time}")

    # 读取 DXF
    doc = ezdxf.readfile(str(DXF_FILE))

    # 文件基本信息
    file_info = build_file_info(doc, DXF_FILE.name)

    # 图层分析
    layers = analyze_layers(doc)

    # 实体分析
    entities = analyze_entities(doc)

    # 添加实体总数到文件信息
    file_info["total_entities"] = entities["total_entities"]

    # 终端打印摘要
    print_summary(file_info, layers, entities)

    # 组装并输出 JSON
    output = {
        "file_info": file_info,
        "layers": layers,
        "entities": entities["type_counts"],
        "blocks": entities["blocks"],
        "texts": entities["texts"],
    }

    output_json = OUTPUT_DIR / f"{DXF_FILE.stem}_cad_analysis.json"
    output_json.write_text(
        json.dumps(output, ensure_ascii=False, indent=2),
        encoding="utf-8",
    )
    print(f"JSON 已保存到: {output_json}")

    # ========== 第二阶段：元素识别 ==========
    print(f"\n{'─' * 56}")
    print(f"  开始元素识别...")
    detector = ElementDetector()
    detection_result = detector.detect_all(doc)

    # 输出元素检测 JSON
    detection_json = OUTPUT_DIR / f"{DXF_FILE.stem}_element_detection.json"
    detection_json.write_text(
        json.dumps(detection_result, ensure_ascii=False, indent=2),
        encoding="utf-8",
    )
    print(f"元素检测结果已保存到: {detection_json}")

    # 输出统计报告
    from element_detector import generate_report
    report = generate_report(detection_result)
    report_path = OUTPUT_DIR / f"{DXF_FILE.stem}_element_report.txt"
    report_path.write_text(report, encoding="utf-8")
    print(f"分析报告已保存到: {report_path}")
    print(report)


if __name__ == "__main__":
    main()