#!/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)