BBIT-Kai
153 lines
4.4 KiB
Python
153 lines
4.4 KiB
Python
from typing import TypedDict
|
||
|
||
from langchain_core.messages import HumanMessage
|
||
from langgraph.graph import StateGraph, END
|
||
|
||
from config.llm import llmVision
|
||
|
||
|
||
# -------- 定义状态 --------
|
||
class State(TypedDict):
|
||
image_url: str # 证件照
|
||
type: int # 证件类别
|
||
content: str # 最终内容
|
||
|
||
|
||
def analyze_card(state: State, prompt_text: str):
|
||
messages = [
|
||
HumanMessage(
|
||
content=[
|
||
{"type": "text", "text": prompt_text},
|
||
{"type": "image_url", "image_url": {"url": state["image_url"]}},
|
||
]
|
||
)
|
||
]
|
||
return llmVision.invoke(messages).content
|
||
|
||
|
||
# -------- 定义节点 --------
|
||
def decide_source(state: State, max_retry=3):
|
||
choice = analyze_card(
|
||
state,
|
||
"""
|
||
请根据图片内容判断其中的证件类型，并输出对应代号。
|
||
代号说明：
|
||
-1: 并非证件
|
||
0：身份证
|
||
1：银行卡
|
||
|
||
只输出代号数字，不要输出任何文字、标点或解释。
|
||
你的输出是：
|
||
""",
|
||
)
|
||
print("图片类型是", choice)
|
||
try:
|
||
choice = int(choice)
|
||
state["type"] = choice
|
||
if choice == -1:
|
||
state["content"] = '{"result": "暂不支持此证件"}'
|
||
except ValueError:
|
||
state["type"] = -1
|
||
state["content"] = '{"result": "暂不支持此证件"}'
|
||
return state
|
||
|
||
|
||
def idcard(state: State):
|
||
state["content"] = analyze_card(
|
||
state,
|
||
"""
|
||
你是一个 OCR 信息提取专家，请从图片中识别出身份证的全部可见信息，并输出严格的 JSON 格式。
|
||
要求：
|
||
1. 无论是正面还是反面，都要识别所有字段。
|
||
2. 如果某个字段无法辨认，请将值设为 null。
|
||
3. 只输出 JSON，不要输出任何解释或多余内容。
|
||
4. 严格保持 JSON 格式，键名统一为英文。
|
||
5. "side" 字段为身份证方向，填写：人像面 或 国徽面。
|
||
|
||
JSON 示例格式：
|
||
{
|
||
"side": "国徽面",
|
||
"name": "持证人姓名",
|
||
"gender": "性别",
|
||
"ethnicity": "民族",
|
||
"id_number": "身份证号",
|
||
"birth_date": "出生日期",
|
||
"address": "住址",
|
||
"issuing_authority": "签发机关",
|
||
"valid_period_start": "有效期开始日期",
|
||
"valid_period_end": "有效期结束日期",
|
||
"notes": "其他可见文字或备注"
|
||
}
|
||
|
||
请确保输出的 JSON 可以被严格解析。
|
||
""",
|
||
)
|
||
return state
|
||
|
||
|
||
def bankcard(state: State):
|
||
state["content"] = analyze_card(
|
||
state,
|
||
"""
|
||
你是一个 OCR 信息提取专家，请从图片中识别出银行卡上的全部可见信息，并输出严格的 JSON 格式。
|
||
要求：
|
||
1. 识别卡号、发卡行、持卡人、有效期等关键信息。
|
||
2. 如果某个字段无法辨认，请将值设为 null。
|
||
3. 只输出 JSON，不要输出任何解释或附加内容。
|
||
4. 严格保持 JSON 格式，键名统一为英文。
|
||
|
||
JSON 示例格式：
|
||
{
|
||
"bank_name": "发卡行名称",
|
||
"card_number": "卡号",
|
||
"card_holder": "持卡人姓名",
|
||
"expiry_date": "有效期",
|
||
"card_type": "卡种类型（如借记卡/信用卡/Visa/MasterCard）",
|
||
"issuer_country": "发卡国家",
|
||
"notes": "其他可见文字或符号"
|
||
}
|
||
|
||
请确保输出的 JSON 可以被严格解析。
|
||
""",
|
||
)
|
||
return state
|
||
|
||
|
||
# ------------------------------------------------------------------------ 构建有向图 --------
|
||
workflow = StateGraph(State)
|
||
workflow.add_node("decide", decide_source)
|
||
workflow.add_node("idcard", idcard)
|
||
workflow.add_node("bankcard", bankcard)
|
||
workflow.set_entry_point("decide")
|
||
# 条件边：根据 path 决定走向
|
||
workflow.add_conditional_edges(
|
||
"decide",
|
||
lambda state: state["type"],
|
||
{
|
||
-1: END,
|
||
0: "idcard",
|
||
1: "bankcard",
|
||
},
|
||
)
|
||
workflow.add_edge("idcard", END)
|
||
workflow.add_edge("bankcard", END)
|
||
graph = workflow.compile()
|
||
|
||
import re
|
||
|
||
|
||
# 执行函数
|
||
def get_license_response(image_url: str):
|
||
final_state = graph.invoke(
|
||
{
|
||
"image_url": image_url,
|
||
}
|
||
)
|
||
# 去掉 ```json 和 ``` 包裹
|
||
final_state["content"] = final_state["content"]
|
||
final_state["content"] = re.sub(
|
||
r"^```json\s*|\s*```$", "", final_state["content"].strip()
|
||
)
|
||
print("最终识别内容：", final_state["content"])
|
||
return final_state
|