143 lines
4.9 KiB
Python
143 lines
4.9 KiB
Python
"""
|
||
OCR service: Use vision model to extract medical report data from images.
|
||
Supports OpenAI-compatible APIs (e.g. SiliconFlow).
|
||
"""
|
||
|
||
import base64
|
||
import json
|
||
import logging
|
||
from pathlib import Path
|
||
|
||
import aiofiles
|
||
from openai import AsyncOpenAI
|
||
|
||
from app.config import settings
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
OCR_PROMPT = """请分析这张图片，首先判断它是否是孕期/产检相关的医疗检查报告。
|
||
|
||
**孕期相关报告包括但不限于：**
|
||
- 产检报告（B超/彩超、NT检查、四维彩超、大排畸等）
|
||
- 血液检查（血常规、血糖、血型、唐氏筛查、TORCH、甲功等）
|
||
- 尿液检查（尿常规等）
|
||
- 肝肾功能检查
|
||
- 凝血功能检查
|
||
- 心电图报告
|
||
- 白带常规
|
||
- 胎心监护报告
|
||
- 其他与孕期产检相关的检查报告
|
||
|
||
**不属于孕期报告的情况：**
|
||
- 非医疗报告（发票、收据、菜单、截图、照片等）
|
||
- 与孕期无关的医疗报告（如骨科、眼科、口腔科等非产检相关）
|
||
- 无法识别内容的模糊图片
|
||
|
||
**请严格按以下JSON格式返回（不要输出其他内容）：**
|
||
|
||
如果**不是**孕期相关报告，返回：
|
||
{
|
||
"is_pregnancy_report": false,
|
||
"reject_reason": "简短说明为什么不是孕期报告（如：这是一张超市购物小票）"
|
||
}
|
||
|
||
如果**是**孕期相关报告，提取所有检查指标并返回：
|
||
{
|
||
"is_pregnancy_report": true,
|
||
"patient_name": "患者姓名（如有）",
|
||
"report_type": "报告类型（如血常规、尿常规、B超等）",
|
||
"report_date": "报告日期（如有，格式 YYYY-MM-DD）",
|
||
"data_items": [
|
||
{
|
||
"name": "测定项目",
|
||
"value": "测定结果（只提取数字，不要单位）",
|
||
"unit": "单位",
|
||
"reference_range": "参考范围（如有）"
|
||
}
|
||
]
|
||
}
|
||
|
||
**重要提示：**
|
||
1. 必须先判断是否为孕期报告，is_pregnancy_report 字段必填
|
||
2. value字段只填写纯数字，不要包含单位（如"23.5"而不是"23.5mm"）
|
||
3. 单位单独放在unit字段中
|
||
4. 如果是B超报告，务必确保双顶径(BPD)、腹围(AC)、股骨长(FL)这三个指标被正确提取"""
|
||
|
||
|
||
def _get_client() -> AsyncOpenAI:
|
||
return AsyncOpenAI(api_key=settings.ai_api_key, base_url=settings.ai_base_url)
|
||
|
||
|
||
async def ocr_recognize(image_path: str) -> dict:
|
||
"""
|
||
Use vision model to extract medical report data from an image.
|
||
Returns structured OCR results.
|
||
"""
|
||
# Resolve path relative to upload_dir
|
||
if image_path.startswith("/uploads/"):
|
||
relative = image_path[len("/uploads/"):]
|
||
full_path = Path(settings.upload_dir) / relative
|
||
else:
|
||
full_path = Path(settings.upload_dir) / image_path.lstrip("/")
|
||
|
||
if not full_path.exists():
|
||
logger.error("OCR 文件不存在: %s", full_path)
|
||
raise FileNotFoundError(f"报告图片不存在: {image_path}")
|
||
|
||
# Async file read
|
||
async with aiofiles.open(full_path, "rb") as f:
|
||
raw = await f.read()
|
||
image_data = base64.standard_b64encode(raw).decode("utf-8")
|
||
|
||
# Determine media type
|
||
suffix = full_path.suffix.lower()
|
||
media_type_map = {
|
||
".jpg": "image/jpeg",
|
||
".jpeg": "image/jpeg",
|
||
".png": "image/png",
|
||
".webp": "image/webp",
|
||
}
|
||
media_type = media_type_map.get(suffix, "image/jpeg")
|
||
|
||
client = _get_client()
|
||
try:
|
||
response = await client.chat.completions.create(
|
||
model=settings.ai_vision_model,
|
||
max_tokens=2048,
|
||
messages=[
|
||
{
|
||
"role": "user",
|
||
"content": [
|
||
{
|
||
"type": "image_url",
|
||
"image_url": {
|
||
"url": f"data:{media_type};base64,{image_data}",
|
||
},
|
||
},
|
||
{
|
||
"type": "text",
|
||
"text": OCR_PROMPT,
|
||
},
|
||
],
|
||
}
|
||
],
|
||
)
|
||
except Exception as e:
|
||
err_msg = str(e)
|
||
logger.error("AI OCR 调用失败: %s: %s", type(e).__name__, err_msg[:200])
|
||
if "auth" in err_msg.lower() or "api key" in err_msg.lower():
|
||
raise RuntimeError("AI 服务认证失败，请检查 API Key 是否有效") from e
|
||
if "<html" in err_msg.lower() or "cloudflare" in err_msg.lower():
|
||
raise RuntimeError("AI 服务请求被拦截，请检查 API 配置") from e
|
||
raise RuntimeError("AI 识别服务暂时不可用，请稍后重试") from e
|
||
|
||
response_text = response.choices[0].message.content
|
||
try:
|
||
if "```json" in response_text:
|
||
response_text = response_text.split("```json")[1].split("```")[0]
|
||
elif "```" in response_text:
|
||
response_text = response_text.split("```")[1].split("```")[0]
|
||
return json.loads(response_text.strip())
|
||
except json.JSONDecodeError:
|
||
return {"raw_text": response_text, "data_items": []}
|