Add UPP STL asset indexing workflow

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

@@ -0,0 +1,157 @@
+#!/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()