PACS/UPP_数据库构建/03_统计STL名称.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""统计已处理STL名称，并按可解释的医学/结构类别归类。"""

from __future__ import annotations

import csv
import json
import re
from collections import Counter, defaultdict
from pathlib import Path
from typing import Any


BASE_DIR = Path(__file__).resolve().parents[1]
DEFAULT_PROCESSED_ROOT = BASE_DIR / "UPP_STL处理" / "已处理STL数据"
DEFAULT_JSON = BASE_DIR / "UPP_数据库构建" / "UPP_STL名称统计.json"
DEFAULT_CSV = BASE_DIR / "UPP_数据库构建" / "UPP_STL名称统计.csv"
DEFAULT_FAMILY_CSV = BASE_DIR / "UPP_数据库构建" / "UPP_STL_family统计.csv"
DEFAULT_ORDERED_CSV = BASE_DIR / "UPP_数据库构建" / "UPP_STL文件family顺序明细.csv"


def classify(segment_name: str) -> tuple[str, str]:
    if segment_name in {"liver", "liver_left", "liver_right"}:
        return "肝脏主体", segment_name
    if re.fullmatch(r"liver_segment_S[1-8]", segment_name):
        return "肝段", segment_name
    if segment_name in {"liver_artery", "liver_vein", "portal_vein", "bile_duct"}:
        return "血管胆管", segment_name
    if segment_name in {"pancreas", "spleen", "cholecyst"}:
        return "腹部脏器", segment_name
    if segment_name in {"skin", "rib", "vertebrae", "sternum", "hipbone", "sacrum"}:
        return "体表骨骼", segment_name
    if re.fullmatch(r"liver_tumor_\d+", segment_name):
        return "肝脏肿瘤", "liver_tumor_*"
    if re.fullmatch(r"liver_cyst_\d+", segment_name):
        return "肝囊肿", "liver_cyst_*"
    if re.fullmatch(r"liver_hemangioma_\d+", segment_name):
        return "肝血管瘤", "liver_hemangioma_*"
    if re.fullmatch(r"pancreas_tumor_\d+", segment_name):
        return "胰腺肿瘤", "pancreas_tumor_*"
    if re.fullmatch(r"Segment_\d+", segment_name):
        return "未命名分割", "Segment_*"
    return "其他", segment_name


def collect(processed_root: Path) -> dict[str, Any]:
    segment_counter: Counter[str] = Counter()
    category_counter: Counter[str] = Counter()
    family_counter: Counter[str] = Counter()
    family_ct: dict[str, set[str]] = defaultdict(set)
    category_ct: dict[str, set[str]] = defaultdict(set)
    name_ct: dict[str, set[str]] = defaultdict(set)
    family_members: dict[str, set[str]] = defaultdict(set)
    ordered_files: list[dict[str, Any]] = []
    ct_dirs = [item for item in processed_root.iterdir() if item.is_dir()]

    for ct_dir in sorted(ct_dirs):
        ct_number = ct_dir.name
        for order_no, stl_file in enumerate(sorted(ct_dir.glob("*.stl")), start=1):
            segment_name = stl_file.stem
            category, family = classify(segment_name)
            segment_counter[segment_name] += 1
            category_counter[category] += 1
            family_counter[family] += 1
            name_ct[segment_name].add(ct_number)
            family_ct[family].add(ct_number)
            category_ct[category].add(ct_number)
            family_members[family].add(segment_name)
            ordered_files.append(
                {
                    "ct_number": ct_number,
                    "order_no": order_no,
                    "segment_name": segment_name,
                    "family": family,
                    "category": category,
                    "file_name": stl_file.name,
                    "file_path": str(stl_file),
                }
            )

    by_name = [
        {
            "segment_name": name,
            "category": classify(name)[0],
            "family": classify(name)[1],
            "file_count": count,
            "ct_count": len(name_ct[name]),
        }
        for name, count in segment_counter.most_common()
    ]
    by_category = [
        {
            "category": category,
            "file_count": count,
            "ct_count": len(category_ct[category]),
        }
        for category, count in category_counter.most_common()
    ]
    by_family = [
        {
            "family": family,
            "category": classify(family.replace("*", "1"))[0] if "*" in family else classify(family)[0],
            "file_count": count,
            "ct_count": len(family_ct[family]),
            "segment_name_count": len(family_members[family]),
            "segment_names": "|".join(sorted(family_members[family])),
        }
        for family, count in family_counter.most_common()
    ]
    return {
        "processed_root": str(processed_root.relative_to(BASE_DIR)),
        "ct_count": len(ct_dirs),
        "stl_file_count": sum(segment_counter.values()),
        "unique_segment_name_count": len(segment_counter),
        "by_category": by_category,
        "by_family": by_family,
        "by_name": by_name,
        "ordered_files": ordered_files,
    }


def write_csv(path: Path, rows: list[dict[str, Any]], fieldnames: list[str]) -> None:
    with path.open("w", encoding="utf-8", newline="") as file:
        writer = csv.DictWriter(file, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)


def main() -> None:
    result = collect(DEFAULT_PROCESSED_ROOT)
    ordered_files = result.pop("ordered_files")
    DEFAULT_JSON.write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8")
    write_csv(DEFAULT_CSV, result["by_name"], ["segment_name", "category", "family", "file_count", "ct_count"])
    write_csv(
        DEFAULT_FAMILY_CSV,
        result["by_family"],
        ["family", "category", "file_count", "ct_count", "segment_name_count", "segment_names"],
    )
    write_csv(
        DEFAULT_ORDERED_CSV,
        ordered_files,
        ["ct_number", "order_no", "segment_name", "family", "category", "file_name", "file_path"],
    )
    print(json.dumps({
        "ct_count": result["ct_count"],
        "stl_file_count": result["stl_file_count"],
        "unique_segment_name_count": result["unique_segment_name_count"],
        "json": str(DEFAULT_JSON),
        "csv": str(DEFAULT_CSV),
        "family_csv": str(DEFAULT_FAMILY_CSV),
        "ordered_csv": str(DEFAULT_ORDERED_CSV),
    }, ensure_ascii=False))


if __name__ == "__main__":
    main()