206 lines
7.8 KiB
Python
206 lines
7.8 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
Kimi 视觉兜底识别：用于 pdftotext/Mineru 对局部字段解析不稳定时，人工指定 PDF 图片页进行辅助抽取。
|
||
|
||
默认只输出 JSON 文件，不直接覆盖 PostgreSQL。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import base64
|
||
import json
|
||
import mimetypes
|
||
import os
|
||
import subprocess
|
||
import sys
|
||
import tempfile
|
||
import urllib.error
|
||
import urllib.request
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
|
||
PROJECT_ROOT = Path(__file__).resolve().parents[2]
|
||
DEFAULT_IMAGE_ROOT = PROJECT_ROOT / "数据处理结果区" / "06_PDF图片对照" / "图片"
|
||
DEFAULT_PDF_ROOT = PROJECT_ROOT / "已处理-患者首页PDF" / "2026_5_25_处理"
|
||
DEFAULT_OUTPUT_DIR = PROJECT_ROOT / "数据处理结果区" / "07_Kimi视觉识别"
|
||
DEFAULT_API_BASE = "https://api.moonshot.cn/v1"
|
||
DEFAULT_MODEL = os.environ.get("KIMI_MODEL", "kimi-k2.6")
|
||
|
||
|
||
PROMPT = """请读取这张住院病案首页图片，仅返回 JSON。
|
||
优先核对这些区域：基本信息、地址联系人、入出院信息、诊断、手术及操作、质控签名、离院方式、费用。
|
||
如果是手术及操作区域，请按以下列拆分：
|
||
手术操作编码、手术操作日期、手术级别、手术操作名称、术者、I助、II助、切口愈合等级、麻醉方式、麻醉医师。
|
||
无法确认的字段填空字符串，不要编造。输出结构：
|
||
{
|
||
"page_summary": "",
|
||
"suspected_missing_fields": [],
|
||
"fields": {},
|
||
"discharge_diagnoses": [],
|
||
"operations": [],
|
||
"fees": {},
|
||
"notes": []
|
||
}
|
||
"""
|
||
|
||
|
||
def parse_args() -> argparse.Namespace:
|
||
parser = argparse.ArgumentParser(description="使用 Kimi 视觉模型辅助识别患者首页图片")
|
||
parser.add_argument("--image-root", type=Path, default=DEFAULT_IMAGE_ROOT, help="PDF 转图片根目录")
|
||
parser.add_argument("--pdf-root", type=Path, default=DEFAULT_PDF_ROOT, help="没有图片目录时，从该 PDF 目录临时渲染页面")
|
||
parser.add_argument("-o", "--output-dir", type=Path, default=DEFAULT_OUTPUT_DIR, help="Kimi 识别结果输出目录")
|
||
parser.add_argument("--case", help="只处理指定病例目录名，例如 ZY010001672803_ori")
|
||
parser.add_argument("--page", help="只处理指定图片名，例如 page-001.png")
|
||
parser.add_argument("--pdf-page", type=int, help="从 PDF 临时渲染时，只处理指定页码，从 1 开始")
|
||
parser.add_argument("--dpi", type=int, default=160, help="PDF 临时渲染分辨率")
|
||
parser.add_argument("--model", default=DEFAULT_MODEL, help="Kimi 模型名")
|
||
parser.add_argument("--api-base", default=os.environ.get("MOONSHOT_API_BASE", DEFAULT_API_BASE), help="Moonshot API base URL")
|
||
parser.add_argument("--api-key", default=os.environ.get("MOONSHOT_API_KEY") or os.environ.get("KIMI_API_KEY"), help="API Key，也可用 MOONSHOT_API_KEY")
|
||
return parser.parse_args()
|
||
|
||
|
||
def image_to_data_url(path: Path) -> str:
|
||
mime_type, _ = mimetypes.guess_type(path.name)
|
||
if mime_type not in {"image/png", "image/jpeg", "image/webp", "image/gif"}:
|
||
mime_type = "image/png"
|
||
encoded = base64.b64encode(path.read_bytes()).decode("utf-8")
|
||
return f"data:{mime_type};base64,{encoded}"
|
||
|
||
|
||
def call_kimi(image_path: Path, args: argparse.Namespace) -> dict[str, Any]:
|
||
if not args.api_key:
|
||
raise RuntimeError("未设置 MOONSHOT_API_KEY/KIMI_API_KEY，无法调用 Kimi API。")
|
||
payload = {
|
||
"model": args.model,
|
||
"messages": [
|
||
{"role": "system", "content": "你是严谨的病案首页结构化抽取助手，只输出 JSON。"},
|
||
{
|
||
"role": "user",
|
||
"content": [
|
||
{"type": "image_url", "image_url": {"url": image_to_data_url(image_path)}},
|
||
{"type": "text", "text": PROMPT},
|
||
],
|
||
},
|
||
],
|
||
}
|
||
request = urllib.request.Request(
|
||
f"{args.api_base.rstrip('/')}/chat/completions",
|
||
data=json.dumps(payload, ensure_ascii=False).encode("utf-8"),
|
||
headers={
|
||
"Content-Type": "application/json",
|
||
"Authorization": f"Bearer {args.api_key}",
|
||
},
|
||
method="POST",
|
||
)
|
||
try:
|
||
with urllib.request.urlopen(request, timeout=180) as response:
|
||
data = json.loads(response.read().decode("utf-8"))
|
||
except urllib.error.HTTPError as exc:
|
||
detail = exc.read().decode("utf-8", errors="replace")
|
||
raise RuntimeError(f"Kimi API 返回错误 {exc.code}: {detail}") from exc
|
||
|
||
content = data["choices"][0]["message"]["content"]
|
||
parsed = parse_json_content(content)
|
||
return {
|
||
"image": str(image_path),
|
||
"model": args.model,
|
||
"recognized_at": datetime.now().isoformat(timespec="seconds"),
|
||
"raw_response": content,
|
||
"parsed": parsed,
|
||
}
|
||
|
||
|
||
def parse_json_content(content: str) -> Any:
|
||
text = content.strip()
|
||
if text.startswith("```"):
|
||
text = text.strip("`")
|
||
if text.startswith("json"):
|
||
text = text[4:].strip()
|
||
try:
|
||
return json.loads(text)
|
||
except json.JSONDecodeError:
|
||
start = text.find("{")
|
||
end = text.rfind("}")
|
||
if start >= 0 and end > start:
|
||
return json.loads(text[start : end + 1])
|
||
return {"notes": ["Kimi 返回内容不是合法 JSON"], "text": content}
|
||
|
||
|
||
def iter_images(args: argparse.Namespace) -> list[Path]:
|
||
root = args.image_root
|
||
if args.case:
|
||
case_dir = root / args.case
|
||
if args.page:
|
||
return [case_dir / args.page]
|
||
return sorted(case_dir.glob("page-*.*"))
|
||
images: list[Path] = []
|
||
for case_dir in sorted(path for path in root.iterdir() if path.is_dir()):
|
||
images.extend(sorted(case_dir.glob("page-*.*")))
|
||
return images
|
||
|
||
|
||
def find_case_pdf(args: argparse.Namespace) -> Path | None:
|
||
if not args.case:
|
||
return None
|
||
stem = Path(args.case).stem
|
||
candidates = [
|
||
args.pdf_root / f"{stem}.pdf",
|
||
args.pdf_root / args.case,
|
||
]
|
||
return next((path for path in candidates if path.exists() and path.is_file()), None)
|
||
|
||
|
||
def render_pdf_pages(pdf_path: Path, output_root: Path, args: argparse.Namespace) -> list[Path]:
|
||
case_dir = output_root / pdf_path.stem
|
||
case_dir.mkdir(parents=True, exist_ok=True)
|
||
if args.pdf_page:
|
||
prefix = case_dir / f"page-{args.pdf_page:03d}"
|
||
command = [
|
||
"pdftoppm",
|
||
"-png",
|
||
"-r",
|
||
str(args.dpi),
|
||
"-f",
|
||
str(args.pdf_page),
|
||
"-singlefile",
|
||
str(pdf_path),
|
||
str(prefix),
|
||
]
|
||
else:
|
||
command = ["pdftoppm", "-png", "-r", str(args.dpi), str(pdf_path), str(case_dir / "page")]
|
||
subprocess.run(command, check=True)
|
||
return sorted(case_dir.glob("page-*.png")) or sorted(case_dir.glob("page.png"))
|
||
|
||
|
||
def main() -> int:
|
||
args = parse_args()
|
||
args.output_dir.mkdir(parents=True, exist_ok=True)
|
||
images = [path for path in iter_images(args) if path.exists() and path.is_file()]
|
||
with tempfile.TemporaryDirectory(prefix="frontpage_kimi_") as temp_dir:
|
||
if not images:
|
||
pdf_path = find_case_pdf(args)
|
||
if pdf_path:
|
||
images = render_pdf_pages(pdf_path, Path(temp_dir), args)
|
||
if not images:
|
||
raise SystemExit("未找到待识别图片，也未找到可临时渲染的 PDF。")
|
||
|
||
for image_path in images:
|
||
print(f"识别：{image_path}")
|
||
result = call_kimi(image_path, args)
|
||
out_path = args.output_dir / f"{image_path.parent.name}_{image_path.stem}.json"
|
||
out_path.write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8")
|
||
print(f"输出：{out_path}")
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
try:
|
||
raise SystemExit(main())
|
||
except Exception as exc: # noqa: BLE001
|
||
print(f"错误：{exc}", file=sys.stderr)
|
||
raise SystemExit(1)
|