#!/usr/bin/env python3 # -*- coding: utf-8 -*- """整理UPP STL目录，并把CT号唯一资产索引同步到PostgreSQL。""" from __future__ import annotations import argparse import csv import json import os import re import shutil import subprocess import tempfile import unicodedata from dataclasses import dataclass from datetime import datetime from pathlib import Path from typing import Any BASE_DIR = Path(__file__).resolve().parents[1] DEFAULT_STL_ROOT = BASE_DIR / "UPP_STL处理" / "待处理STL数据" DEFAULT_PROCESSED_ROOT = BASE_DIR / "UPP_STL处理" / "已处理STL数据" DEFAULT_LIST_JSON = BASE_DIR / "UPP列表处理" / "数据处理结果区" / "全量分片结果" / "合并_图片表格_结构化.json" DEFAULT_SCHEMA = BASE_DIR / "UPP_数据库构建" / "01_UPP资产索引建表.sql" DEFAULT_REPORT = BASE_DIR / "UPP_数据库构建" / "UPP_STL资产同步报告.json" CT_PATTERN = re.compile(r"(D?CT\d{8,})", re.IGNORECASE) VALID_CT_PATTERN = re.compile(r"^D?CT\d{8,}$") ASSET_FIELDS = [ "ct_number", "list_present", "stl_present", "patient_name", "patient_sex", "patient_age", "patient_id_masked", "exam_date", "task_created_at", "exam_description", "exam_device", "algorithm_model", "upp_status", "list_record_count", "selected_list_record", "list_records", "selected_source_case_dir", "selected_source_stl_dir", "processed_stl_dir", "stl_case_name", "stl_sequence_no", "stl_file_count", "stl_total_bytes", "stl_files", "stl_candidates", ] STL_FIELDS = [ "ct_number", "file_count", "total_bytes", "segment_names", "segment_families", "segment_categories", "file_names", "source_file_paths", "processed_file_paths", "files", ] MODEL_VALUES = {"肝胆模型", "泌尿模型", "胸外模型"} DEVICE_VALUES = {"CT", "MR", "DR", "CR", "DX", "US", "XA", "NM", "PT"} def normalize_text(value: Any) -> str: if value is None: return "" text = unicodedata.normalize("NFKC", str(value)).replace("\u3000", " ") return re.sub(r"\s+", " ", text).strip() def normalize_ct(value: Any) -> str: return re.sub(r"\s+", "", normalize_text(value)).upper() def valid_ct(value: str) -> bool: return bool(VALID_CT_PATTERN.fullmatch(value)) def extract_ct(text: str) -> str: match = CT_PATTERN.search(text) if not match: return "" ct_number = normalize_ct(match.group(1)) return ct_number if valid_ct(ct_number) else "" def json_dump(value: Any) -> str: return json.dumps(value, ensure_ascii=False, separators=(",", ":")) def path_text(path: Path) -> str: return str(path.expanduser().absolute()) def parse_time(value: Any) -> datetime | None: text = normalize_text(value) if not text: return None for fmt in ("%Y-%m-%d %H:%M:%S", "%Y/%m/%d %H:%M:%S", "%Y-%m-%d", "%Y/%m/%d"): try: return datetime.strptime(text, fmt) except ValueError: pass return None def normalize_status(value: Any) -> str: text = normalize_text(value).replace("\ufe0f", "") compact = re.sub(r"\s+", "", text) if not compact: return "" if "部分重建成功" in compact: return "√ 部分重建成功" if "重建失败" in compact or compact.startswith(("×", "✗")): return "× 重建失败" if "重建成功" in compact: return "√ 重建成功" if "已报告" in compact: return "√ 已报告" return text def status_like(value: Any) -> bool: text = normalize_text(value) return bool(re.search(r"(重建成功|重建失败|部分重建成功|已报告)", text)) def classify_segment(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 clean_record_info(info: dict[str, Any]) -> dict[str, Any]: cleaned = dict(info) exam_description = normalize_text(cleaned.get("检查描述")) exam_device = normalize_text(cleaned.get("检查设备")) algorithm_model = normalize_text(cleaned.get("算法模型")) if exam_description in DEVICE_VALUES and exam_device in MODEL_VALUES and status_like(algorithm_model): cleaned["检查描述"] = "" cleaned["检查设备"] = exam_description cleaned["算法模型"] = exam_device cleaned["状态"] = normalize_status(algorithm_model) elif exam_device in MODEL_VALUES and status_like(algorithm_model) and normalize_text(cleaned.get("状态")) in {"土", "士", "↓"}: cleaned["检查设备"] = "CT" cleaned["算法模型"] = exam_device cleaned["状态"] = normalize_status(algorithm_model) else: cleaned["状态"] = normalize_status(cleaned.get("状态")) return cleaned def clean_record(record: dict[str, Any]) -> dict[str, Any]: cleaned = dict(record) info = cleaned.get("记录信息") if isinstance(info, dict): cleaned["记录信息"] = clean_record_info(info) return cleaned def sql_quote_path(path: Path) -> str: return "'" + str(path.expanduser().absolute()).replace("'", "''") + "'" @dataclass(frozen=True) class StlCandidate: ct_number: str source_case_dir: Path source_stl_dir: Path sequence_no: int | None files: tuple[Path, ...] @property def file_count(self) -> int: return len(self.files) @property def total_bytes(self) -> int: return sum(file.stat().st_size for file in self.files) def score(self) -> tuple[int, int, str, str]: return ( self.file_count, self.sequence_no if self.sequence_no is not None else -1, self.source_case_dir.name, self.source_stl_dir.name, ) def summary(self) -> dict[str, Any]: return { "source_case_dir": path_text(self.source_case_dir), "source_stl_dir": path_text(self.source_stl_dir), "case_name": self.source_case_dir.name, "sequence_no": self.sequence_no, "file_count": self.file_count, "total_bytes": self.total_bytes, } def sequence_no(case_dir: Path, stl_dir: Path) -> int | None: values: list[int] = [] for pattern in (r"-(\d{3,6})_D?CT", r"^(\d{3,6})-STL$"): for text in (case_dir.name, stl_dir.name): values.extend(int(item) for item in re.findall(pattern, text, re.IGNORECASE)) return max(values) if values else None def scan_stl_candidates(stl_root: Path) -> tuple[dict[str, list[StlCandidate]], list[str]]: candidates: dict[str, list[StlCandidate]] = {} no_ct_dirs: list[str] = [] for directory in sorted(stl_root.rglob("*")): if not directory.is_dir(): continue files = tuple(sorted([item for item in directory.iterdir() if item.is_file() and item.suffix.lower() == ".stl"])) if not files: continue ct_number = extract_ct("/".join(directory.relative_to(stl_root).parts)) if not ct_number: no_ct_dirs.append(path_text(directory)) continue candidate = StlCandidate( ct_number=ct_number, source_case_dir=directory.parent, source_stl_dir=directory, sequence_no=sequence_no(directory.parent, directory), files=files, ) candidates.setdefault(ct_number, []).append(candidate) return candidates, no_ct_dirs def choose_candidates(candidates: dict[str, list[StlCandidate]]) -> dict[str, StlCandidate]: return {ct_number: max(items, key=lambda item: item.score()) for ct_number, items in candidates.items()} def clear_directory(directory: Path) -> None: directory.mkdir(parents=True, exist_ok=True) for child in directory.iterdir(): if child.is_dir() and not child.is_symlink(): shutil.rmtree(child) else: child.unlink() def place_file(source: Path, destination: Path, mode: str) -> str: destination.parent.mkdir(parents=True, exist_ok=True) if mode == "symlink": destination.symlink_to(source.expanduser().absolute()) return "symlink" if mode == "copy": shutil.copy2(source, destination) return "copy" try: os.link(source, destination) return "hardlink" except OSError: shutil.copy2(source, destination) return "copy" def materialize_selected_stl( selected: dict[str, StlCandidate], processed_root: Path, link_mode: str, refresh_files: bool, ) -> tuple[dict[str, list[dict[str, Any]]], dict[str, str]]: processed_files: dict[str, list[dict[str, Any]]] = {} link_methods: dict[str, str] = {} processed_root.mkdir(parents=True, exist_ok=True) for ct_number, candidate in sorted(selected.items()): ct_dir = processed_root / ct_number if refresh_files: clear_directory(ct_dir) else: ct_dir.mkdir(parents=True, exist_ok=True) file_rows: list[dict[str, Any]] = [] used_methods: set[str] = set() for source in candidate.files: destination = ct_dir / source.name if destination.exists() or destination.is_symlink(): destination.unlink() method = place_file(source, destination, link_mode) used_methods.add(method) category, family = classify_segment(source.stem) file_rows.append( { "segment_name": source.stem, "family": family, "category": category, "file_name": source.name, "source_file_path": path_text(source), "processed_file_path": path_text(destination), "size_bytes": source.stat().st_size, } ) manifest = { "ct_number": ct_number, "selected_source_case_dir": path_text(candidate.source_case_dir), "selected_source_stl_dir": path_text(candidate.source_stl_dir), "processed_stl_dir": path_text(ct_dir), "selection_rule": "file_count_desc_then_sequence_no_desc", "stl_file_count": candidate.file_count, "stl_total_bytes": candidate.total_bytes, "files": file_rows, } (ct_dir / "manifest.json").write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8") processed_files[ct_number] = file_rows link_methods[ct_number] = "+".join(sorted(used_methods)) return processed_files, link_methods def load_list_records(list_json: Path) -> tuple[dict[str, list[dict[str, Any]]], list[dict[str, Any]]]: data = json.loads(list_json.read_text(encoding="utf-8")) records_by_ct: dict[str, list[dict[str, Any]]] = {} invalid: list[dict[str, Any]] = [] for index, record in enumerate(data.get("图片表格记录", []), start=1): record = clean_record(record) info = record.get("记录信息") or {} ct_number = normalize_ct(info.get("检查号", "")) if not valid_ct(ct_number): invalid.append({"row_index": index, "ct_number": ct_number, "record": record}) continue records_by_ct.setdefault(ct_number, []).append(record) return records_by_ct, invalid def select_list_record(records: list[dict[str, Any]]) -> dict[str, Any]: def key(record: dict[str, Any]) -> tuple[datetime, datetime, int, int]: info = record.get("记录信息") or {} image = record.get("图片信息") or {} task_time = parse_time(info.get("任务创建时间")) or datetime.min exam_time = parse_time(info.get("检查日期")) or datetime.min page_no = -1 seq = image.get("图片序号") if isinstance(seq, list): numbers = [item for item in seq if isinstance(item, int)] if numbers: page_no = numbers[0] row_no = int(image.get("图片内行号") or -1) return (task_time, exam_time, page_no, row_no) return max(records, key=key) def build_rows( selected: dict[str, StlCandidate], all_candidates: dict[str, list[StlCandidate]], processed_root: Path, processed_files: dict[str, list[dict[str, Any]]], list_records: dict[str, list[dict[str, Any]]], ) -> tuple[list[dict[str, Any]], list[dict[str, Any]], int]: asset_rows: list[dict[str, Any]] = [] stl_rows: list[dict[str, Any]] = [] stl_file_count = 0 for ct_number in sorted(set(selected) | set(list_records)): candidate = selected.get(ct_number) records = list_records.get(ct_number, []) chosen_record = select_list_record(records) if records else None info = (chosen_record or {}).get("记录信息") or {} files = processed_files.get(ct_number, []) processed_dir = processed_root / ct_number asset_rows.append( { "ct_number": ct_number, "list_present": bool(records), "stl_present": candidate is not None, "patient_name": normalize_text(info.get("姓名")), "patient_sex": normalize_text(info.get("性别")), "patient_age": normalize_text(info.get("年龄")), "patient_id_masked": normalize_text(info.get("患者号")), "exam_date": normalize_text(info.get("检查日期")), "task_created_at": normalize_text(info.get("任务创建时间")), "exam_description": normalize_text(info.get("检查描述")), "exam_device": normalize_text(info.get("检查设备")), "algorithm_model": normalize_text(info.get("算法模型")), "upp_status": normalize_text(info.get("状态")), "list_record_count": len(records), "selected_list_record": json_dump(chosen_record) if chosen_record else "", "list_records": json_dump(records), "selected_source_case_dir": path_text(candidate.source_case_dir) if candidate else "", "selected_source_stl_dir": path_text(candidate.source_stl_dir) if candidate else "", "processed_stl_dir": path_text(processed_dir) if candidate else "", "stl_case_name": candidate.source_case_dir.name if candidate else "", "stl_sequence_no": candidate.sequence_no if candidate and candidate.sequence_no is not None else "", "stl_file_count": candidate.file_count if candidate else 0, "stl_total_bytes": candidate.total_bytes if candidate else 0, "stl_files": json_dump(files), "stl_candidates": json_dump([item.summary() for item in sorted(all_candidates.get(ct_number, []), key=lambda item: item.score(), reverse=True)]), } ) if files: stl_file_count += len(files) stl_rows.append( { "ct_number": ct_number, "file_count": len(files), "total_bytes": sum(int(file_info.get("size_bytes") or 0) for file_info in files), "segment_names": json_dump([file_info["segment_name"] for file_info in files]), "segment_families": json_dump([file_info["family"] for file_info in files]), "segment_categories": json_dump([file_info["category"] for file_info in files]), "file_names": json_dump([file_info["file_name"] for file_info in files]), "source_file_paths": json_dump([file_info["source_file_path"] for file_info in files]), "processed_file_paths": json_dump([file_info["processed_file_path"] for file_info in files]), "files": json_dump(files), } ) return asset_rows, stl_rows, stl_file_count def write_csv(rows: list[dict[str, Any]], fields: list[str], suffix: str) -> Path: temp_file = tempfile.NamedTemporaryFile("w", encoding="utf-8", newline="", suffix=suffix, delete=False) with temp_file: writer = csv.DictWriter(temp_file, fieldnames=fields) writer.writeheader() for row in rows: writer.writerow({field: row.get(field, "") for field in fields}) return Path(temp_file.name) def run_psql(args: argparse.Namespace, asset_csv: Path, stl_csv: Path) -> None: sql = f""" \\set ON_ERROR_STOP on \\i {sql_quote_path(Path(args.schema))} CREATE TEMP TABLE stg_upp_exam_assets ( ct_number text, list_present boolean, stl_present boolean, patient_name text, patient_sex text, patient_age text, patient_id_masked text, exam_date timestamptz, task_created_at timestamptz, exam_description text, exam_device text, algorithm_model text, upp_status text, list_record_count integer, selected_list_record text, list_records text, selected_source_case_dir text, selected_source_stl_dir text, processed_stl_dir text, stl_case_name text, stl_sequence_no integer, stl_file_count integer, stl_total_bytes bigint, stl_files text, stl_candidates text ); CREATE TEMP TABLE stg_upp_stl_files ( ct_number text, file_count integer, total_bytes bigint, segment_names text, segment_families text, segment_categories text, file_names text, source_file_paths text, processed_file_paths text, files text ); \\copy stg_upp_exam_assets({",".join(ASSET_FIELDS)}) FROM {sql_quote_path(asset_csv)} WITH (FORMAT csv, HEADER true, NULL '') \\copy stg_upp_stl_files({",".join(STL_FIELDS)}) FROM {sql_quote_path(stl_csv)} WITH (FORMAT csv, HEADER true, NULL '') INSERT INTO upp_exam_assets ( ct_number, list_present, stl_present, patient_name, patient_sex, patient_age, patient_id_masked, exam_date, task_created_at, exam_description, exam_device, algorithm_model, upp_status, list_record_count, selected_list_record, list_records, selected_source_case_dir, selected_source_stl_dir, processed_stl_dir, stl_case_name, stl_sequence_no, stl_file_count, stl_total_bytes, stl_files, stl_candidates, updated_at ) SELECT ct_number, list_present, stl_present, patient_name, patient_sex, patient_age, patient_id_masked, exam_date, task_created_at, exam_description, exam_device, algorithm_model, upp_status, COALESCE(list_record_count, 0), NULLIF(selected_list_record, '')::jsonb, COALESCE(NULLIF(list_records, ''), '[]')::jsonb, selected_source_case_dir, selected_source_stl_dir, processed_stl_dir, stl_case_name, stl_sequence_no, COALESCE(stl_file_count, 0), COALESCE(stl_total_bytes, 0), COALESCE(NULLIF(stl_files, ''), '[]')::jsonb, COALESCE(NULLIF(stl_candidates, ''), '[]')::jsonb, now() FROM stg_upp_exam_assets ON CONFLICT (ct_number) DO UPDATE SET list_present = EXCLUDED.list_present, stl_present = EXCLUDED.stl_present, patient_name = EXCLUDED.patient_name, patient_sex = EXCLUDED.patient_sex, patient_age = EXCLUDED.patient_age, patient_id_masked = EXCLUDED.patient_id_masked, exam_date = EXCLUDED.exam_date, task_created_at = EXCLUDED.task_created_at, exam_description = EXCLUDED.exam_description, exam_device = EXCLUDED.exam_device, algorithm_model = EXCLUDED.algorithm_model, upp_status = EXCLUDED.upp_status, list_record_count = EXCLUDED.list_record_count, selected_list_record = EXCLUDED.selected_list_record, list_records = EXCLUDED.list_records, selected_source_case_dir = EXCLUDED.selected_source_case_dir, selected_source_stl_dir = EXCLUDED.selected_source_stl_dir, processed_stl_dir = EXCLUDED.processed_stl_dir, stl_case_name = EXCLUDED.stl_case_name, stl_sequence_no = EXCLUDED.stl_sequence_no, stl_file_count = EXCLUDED.stl_file_count, stl_total_bytes = EXCLUDED.stl_total_bytes, stl_files = EXCLUDED.stl_files, stl_candidates = EXCLUDED.stl_candidates, updated_at = now(); DELETE FROM upp_stl_files WHERE ct_number IN (SELECT ct_number FROM stg_upp_exam_assets); INSERT INTO upp_stl_files ( ct_number, file_count, total_bytes, segment_names, segment_families, segment_categories, file_names, source_file_paths, processed_file_paths, files, updated_at ) SELECT ct_number, COALESCE(file_count, 0), COALESCE(total_bytes, 0), COALESCE(NULLIF(segment_names, ''), '[]')::jsonb, COALESCE(NULLIF(segment_families, ''), '[]')::jsonb, COALESCE(NULLIF(segment_categories, ''), '[]')::jsonb, COALESCE(NULLIF(file_names, ''), '[]')::jsonb, COALESCE(NULLIF(source_file_paths, ''), '[]')::jsonb, COALESCE(NULLIF(processed_file_paths, ''), '[]')::jsonb, COALESCE(NULLIF(files, ''), '[]')::jsonb, now() FROM stg_upp_stl_files; """ env = os.environ.copy() password = args.password or env.get("PGPASSWORD") if password: env["PGPASSWORD"] = password command = ["psql"] if args.host: command.extend(["-h", args.host]) if args.port: command.extend(["-p", str(args.port)]) if args.user: command.extend(["-U", args.user]) if args.dbname: command.extend(["-d", args.dbname]) command.extend(["-v", "ON_ERROR_STOP=1"]) completed = subprocess.run(command, input=sql, text=True, env=env, capture_output=True) print(completed.stdout, end="") if completed.returncode != 0: print(completed.stderr, end="") completed.check_returncode() def write_report(path: Path, report: dict[str, Any]) -> None: path.parent.mkdir(parents=True, exist_ok=True) path.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8") def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--stl-root", default=str(DEFAULT_STL_ROOT), help="待处理STL数据根目录") parser.add_argument("--processed-root", default=str(DEFAULT_PROCESSED_ROOT), help="规范化后的STL输出目录") parser.add_argument("--list-json", default=str(DEFAULT_LIST_JSON), help="UPP列表合并结构化JSON") parser.add_argument("--schema", default=str(DEFAULT_SCHEMA), help="PostgreSQL建表SQL") parser.add_argument("--report", default=str(DEFAULT_REPORT), help="同步报告JSON") parser.add_argument("--link-mode", choices=["hardlink", "copy", "symlink"], default="hardlink", help="输出STL文件方式") parser.add_argument("--no-refresh-files", action="store_true", help="不清空已处理CT目录中的旧文件") parser.add_argument("--dry-run", action="store_true", help="只整理文件和生成报告，不写数据库") parser.add_argument("--host", default=os.getenv("PGHOST", "")) parser.add_argument("--port", default=os.getenv("PGPORT", "")) parser.add_argument("--dbname", default=os.getenv("PGDATABASE", "")) parser.add_argument("--user", default=os.getenv("PGUSER", "")) parser.add_argument("--password", default=os.getenv("PGPASSWORD", ""), help="数据库密码；也可用PGPASSWORD环境变量") return parser.parse_args() def main() -> None: args = parse_args() stl_root = Path(args.stl_root) processed_root = Path(args.processed_root) all_candidates, no_ct_dirs = scan_stl_candidates(stl_root) selected = choose_candidates(all_candidates) processed_files, link_methods = materialize_selected_stl( selected=selected, processed_root=processed_root, link_mode=args.link_mode, refresh_files=not args.no_refresh_files, ) list_records, invalid_list_records = load_list_records(Path(args.list_json)) asset_rows, stl_rows, stl_file_count = build_rows( selected=selected, all_candidates=all_candidates, processed_root=processed_root, processed_files=processed_files, list_records=list_records, ) duplicate_ct_numbers = sorted(ct_number for ct_number, items in all_candidates.items() if len(items) > 1) report = { "stl_candidate_dirs": sum(len(items) for items in all_candidates.values()), "stl_unique_ct": len(selected), "stl_no_ct_dirs": len(no_ct_dirs), "stl_duplicate_ct": len(duplicate_ct_numbers), "list_unique_valid_ct": len(list_records), "list_invalid_ct_records": len(invalid_list_records), "asset_rows": len(asset_rows), "stl_table_rows": len(stl_rows), "stl_file_rows": stl_file_count, "matched_list_and_stl_ct": len(set(selected) & set(list_records)), "list_without_stl_ct": len(set(list_records) - set(selected)), "stl_without_list_ct": len(set(selected) - set(list_records)), "duplicate_ct_examples": { ct_number: [item.summary() for item in sorted(all_candidates[ct_number], key=lambda item: item.score(), reverse=True)] for ct_number in duplicate_ct_numbers[:20] }, "no_ct_dir_examples": no_ct_dirs[:20], "invalid_list_ct_examples": invalid_list_records[:20], "link_methods": dict(sorted(link_methods.items())[:20]), "processed_root": path_text(processed_root), "selection_rule": "同一CT号优先选择STL文件数更多；文件数相同选择目录/病例名中的较大序号", "dry_run": args.dry_run, } write_report(Path(args.report), report) asset_csv = write_csv(asset_rows, ASSET_FIELDS, "_upp_exam_assets.csv") stl_csv = write_csv(stl_rows, STL_FIELDS, "_upp_stl_files.csv") try: if not args.dry_run: run_psql(args, asset_csv, stl_csv) finally: asset_csv.unlink(missing_ok=True) stl_csv.unlink(missing_ok=True) print(json.dumps({k: report[k] for k in [ "stl_candidate_dirs", "stl_unique_ct", "stl_no_ct_dirs", "stl_duplicate_ct", "list_unique_valid_ct", "asset_rows", "stl_table_rows", "stl_file_rows", "matched_list_and_stl_ct", "list_without_stl_ct", "stl_without_list_ct", ]}, ensure_ascii=False)) if __name__ == "__main__": main()