432 lines
16 KiB
Python
432 lines
16 KiB
Python
"""身份证识别 API 服务（极限低配版：1 核 CPU / 2G 内存）。
|
||
|
||
POST /api/v1/idcard/recognize 上传图片（multipart）
|
||
POST /api/v1/idcard/recognize/base64 base64 JSON
|
||
POST /api/v1/auth/token 用 client_secret 换短期 token（可选）
|
||
GET /api/v1/idcard/health 健康检查（含内存读数）
|
||
GET / 服务信息
|
||
|
||
鉴权三层（全部标准库实现，无额外依赖）：
|
||
1. 静态 token：API_TOKENS=名称:token,名称2:token2，每个调用方一把，可单独吊销
|
||
2. 请求头：Authorization: Bearer <token>（标准），兼容 X-API-Key
|
||
3. 短期 token：设 CLIENT_SECRET 后由 /api/v1/auth/token 签发，HMAC-SHA256 签名带过期
|
||
|
||
低配适配策略见 README，核心五点：
|
||
1. 单 worker + 引擎全局单例，绝不加载多份模型
|
||
2. 线程三重钉死：OMP/OPENBLAS/MKL=1 + cv2.setNumThreads(0)
|
||
—— Docker 的 cpus 限额不改变容器内 os.cpu_count()，onnxruntime 会按
|
||
宿主机核数建线程池，1 核机器上直接线程风暴，必须显式限死
|
||
3. 关闭方向分类器（身份证基本正向）
|
||
4. 输入图最长边 1200px + det_limit_side_len=960
|
||
5. content-length 预检 + 惰性尺寸检查，防大图/解压炸弹打爆内存
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import base64
|
||
import binascii
|
||
import hashlib
|
||
import hmac
|
||
import io
|
||
import json
|
||
import logging
|
||
import os
|
||
import time
|
||
from typing import Optional
|
||
|
||
# —— 线程限制必须早于 numpy / onnxruntime import 才会生效 ——
|
||
# setdefault：不覆盖用户显式设置，未设置时兜底为 1
|
||
for _v in ("OMP_NUM_THREADS", "OPENBLAS_NUM_THREADS", "MKL_NUM_THREADS"):
|
||
os.environ.setdefault(_v, "1")
|
||
|
||
import numpy as np
|
||
from fastapi import FastAPI, File, Header, HTTPException, Request, UploadFile
|
||
from fastapi.middleware.cors import CORSMiddleware
|
||
from fastapi.responses import JSONResponse
|
||
from PIL import Image
|
||
from pydantic import BaseModel, Field
|
||
|
||
try: # OpenCV 自建线程池，1 核机器上是纯负担
|
||
import cv2
|
||
|
||
cv2.setNumThreads(0)
|
||
except Exception: # pragma: no cover - headless 缺失时不影响主流程
|
||
cv2 = None
|
||
|
||
from . import parser as id_parser
|
||
|
||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
||
log = logging.getLogger("idcard-api")
|
||
|
||
MAX_IMAGE_BYTES = int(os.environ.get("MAX_IMAGE_BYTES", 5 * 1024 * 1024))
|
||
MAX_SIDE = int(os.environ.get("MAX_SIDE", 1200)) # 输入图最长边，超过则等比缩小
|
||
MAX_DIMENSION = int(os.environ.get("MAX_DIMENSION", 6000)) # header 尺寸硬上限，超过直接拒
|
||
DET_LIMIT_SIDE_LEN = int(os.environ.get("DET_LIMIT_SIDE_LEN", 960))
|
||
REC_BATCH_NUM = int(os.environ.get("REC_BATCH_NUM", 1)) # 低内存必须为 1
|
||
MODEL_DIR = os.environ.get("RAPIDOCR_MODEL_DIR") or None
|
||
# 允许跨域的前端来源，逗号分隔；默认 * 方便本地前端开发。
|
||
# 生产环境建议收紧为实际地址，例如 http://10.1.1.163:5173,http://nas.local
|
||
CORS_ORIGINS = os.environ.get("CORS_ORIGINS", "*").strip()
|
||
|
||
# —— 鉴权配置 ——
|
||
# 1) 静态 token（主力用法）：每个调用方发一把，可单独吊销。
|
||
# 格式：名称:token,名称2:token2 例如 frontend:abc123,win-service:def456
|
||
API_TOKENS_RAW = os.environ.get("API_TOKENS", "").strip()
|
||
# 2) 兼容旧配置：单个 X-API-Key
|
||
API_KEY = os.environ.get("API_KEY", "").strip()
|
||
# 3) 可选短期 token：设了 CLIENT_SECRET 后启用 /api/v1/auth/token 换 token
|
||
CLIENT_ID = os.environ.get("CLIENT_ID", "default").strip()
|
||
CLIENT_SECRET = os.environ.get("CLIENT_SECRET", "").strip()
|
||
TOKEN_TTL = int(os.environ.get("TOKEN_TTL", 3600)) # 短期 token 有效期（秒）
|
||
|
||
|
||
def _parse_tokens() -> dict:
|
||
"""解析 API_TOKENS 为 {token: 调用方名称}。token 里不要含逗号/冒号。"""
|
||
tokens = {}
|
||
for item in API_TOKENS_RAW.split(","):
|
||
item = item.strip()
|
||
if not item:
|
||
continue
|
||
if ":" in item:
|
||
name, tok = item.split(":", 1)
|
||
if tok.strip():
|
||
tokens[tok.strip()] = name.strip() or "client"
|
||
else:
|
||
tokens[item] = "client"
|
||
if API_KEY:
|
||
tokens[API_KEY] = "legacy-key"
|
||
return tokens
|
||
|
||
|
||
_TOKEN_MAP = _parse_tokens()
|
||
|
||
# 弱口令提醒：示例占位符（REPLACE_ME_*）或过短的 token 不该出现在正式环境。
|
||
# 只告警不阻止启动，避免影响联调；换成真随机串后这条日志自动消失。
|
||
_weak_clients = sorted(
|
||
{name for tok, name in _TOKEN_MAP.items() if "REPLACE_ME" in tok.upper() or len(tok) < 12}
|
||
)
|
||
if _weak_clients:
|
||
log.warning(
|
||
"token 过弱或仍是示例占位符，请尽快更换：%s（生成：openssl rand -hex 24）",
|
||
", ".join(_weak_clients),
|
||
)
|
||
|
||
_engine = None
|
||
|
||
|
||
def get_engine():
|
||
"""懒加载 OCR 引擎（单例）。
|
||
|
||
rapidocr-onnxruntime 不同版本的构造参数略有差异，这里逐级降级尝试，
|
||
避免因某个参数不存在直接启动失败。生产环境请锁定 requirements.txt 版本。
|
||
"""
|
||
global _engine
|
||
if _engine is None:
|
||
from rapidocr_onnxruntime import RapidOCR
|
||
|
||
log.info("loading RapidOCR engine (model_dir=%s)", MODEL_DIR)
|
||
base = {"model_dir": MODEL_DIR} if MODEL_DIR else {}
|
||
for kwargs in (
|
||
{
|
||
**base,
|
||
"rec_batch_num": REC_BATCH_NUM,
|
||
"det_limit_side_len": DET_LIMIT_SIDE_LEN,
|
||
"use_angle_cls": False,
|
||
},
|
||
{**base, "rec_batch_num": REC_BATCH_NUM, "use_angle_cls": False},
|
||
{**base, "rec_batch_num": REC_BATCH_NUM},
|
||
base,
|
||
):
|
||
try:
|
||
_engine = RapidOCR(**kwargs)
|
||
log.info("engine ready with kwargs=%s", sorted(kwargs))
|
||
break
|
||
except TypeError as ex:
|
||
log.warning("RapidOCR rejected %s: %s", sorted(kwargs), ex)
|
||
if _engine is None:
|
||
raise RuntimeError("cannot initialize RapidOCR engine")
|
||
return _engine
|
||
|
||
|
||
app = FastAPI(
|
||
title="IDCard Recognition API",
|
||
version="2.0.0",
|
||
description="离线身份证信息识别服务，返回结构化 JSON。极限低配硬件友好（1C/2G）。",
|
||
)
|
||
|
||
# CORS：浏览器里的前端页面跨源调用时必须开启，否则预检（OPTIONS）直接失败。
|
||
# 后端到后端调用（curl / requests / 其他容器）不受 CORS 影响，开不开都能用。
|
||
# 鉴权走 X-API-Key 请求头，不用 cookie，所以 allow_credentials 保持 False
|
||
# （allow_origins=["*"] 与 credentials=True 是非法组合，浏览器会拒绝）。
|
||
_cors_origins = [o.strip() for o in CORS_ORIGINS.split(",") if o.strip()] or ["*"]
|
||
app.add_middleware(
|
||
CORSMiddleware,
|
||
allow_origins=_cors_origins,
|
||
allow_credentials=False,
|
||
allow_methods=["GET", "POST", "OPTIONS"],
|
||
allow_headers=["*"],
|
||
expose_headers=["*"],
|
||
max_age=600,
|
||
)
|
||
if "*" in _cors_origins and not _TOKEN_MAP and not CLIENT_SECRET:
|
||
log.warning(
|
||
"CORS 允许全部来源且未配置任何 token：局域网内任意网页都能读取识别结果。"
|
||
" 生产环境请设置 API_TOKENS 并收紧 CORS_ORIGINS。"
|
||
)
|
||
|
||
|
||
@app.middleware("http")
|
||
async def _limit_payload(request: Request, call_next):
|
||
"""第一道防线：按 content-length 预检，别让超大请求先读进内存。"""
|
||
cl = request.headers.get("content-length")
|
||
if cl:
|
||
try:
|
||
if int(cl) > MAX_IMAGE_BYTES:
|
||
return JSONResponse(
|
||
status_code=413,
|
||
content={"detail": "payload too large, max %d bytes" % MAX_IMAGE_BYTES},
|
||
)
|
||
except ValueError:
|
||
pass
|
||
return await call_next(request)
|
||
|
||
|
||
class Base64Body(BaseModel):
|
||
image_base64: str = Field(..., description="图片 base64，可带 data:image/... 前缀")
|
||
side: str = Field("auto", description="auto / front / back")
|
||
mask: bool = Field(False, description="是否脱敏返回")
|
||
|
||
|
||
class TokenRequest(BaseModel):
|
||
client_id: str = Field("default", description="调用方标识")
|
||
client_secret: str = Field(..., description="客户端密钥（与 CLIENT_SECRET 一致）")
|
||
|
||
|
||
def _b64u(raw: bytes) -> str:
|
||
return base64.urlsafe_b64encode(raw).rstrip(b"=").decode()
|
||
|
||
|
||
def _b64u_decode(text: str) -> bytes:
|
||
return base64.urlsafe_b64decode(text + "=" * (-len(text) % 4))
|
||
|
||
|
||
def _sign(payload: str) -> str:
|
||
"""HMAC-SHA256 签名（标准库，无额外依赖）。"""
|
||
digest = hmac.new(CLIENT_SECRET.encode(), payload.encode(), hashlib.sha256).digest()
|
||
return _b64u(digest)
|
||
|
||
|
||
def issue_token(subject: str) -> dict:
|
||
"""签发短期 token：base64(payload).hmac —— 结构类似 JWT，但不依赖 PyJWT。"""
|
||
payload = _b64u(
|
||
json.dumps(
|
||
{"sub": subject, "iat": int(time.time()), "exp": int(time.time()) + TOKEN_TTL},
|
||
separators=(",", ":"),
|
||
).encode()
|
||
)
|
||
return {"token": payload + "." + _sign(payload), "expires_in": TOKEN_TTL, "token_type": "Bearer"}
|
||
|
||
|
||
def _verify_short_token(token: str) -> Optional[str]:
|
||
"""校验短期 token 的签名和有效期，成功返回 subject。"""
|
||
if not CLIENT_SECRET or "." not in token:
|
||
return None
|
||
payload, sig = token.rsplit(".", 1)
|
||
if not hmac.compare_digest(sig, _sign(payload)): # 常数时间比较，防时序侧信道
|
||
return None
|
||
try:
|
||
data = json.loads(_b64u_decode(payload))
|
||
except Exception:
|
||
return None
|
||
# 到达 exp 时刻即失效（用 >=，避免整数截断让 token 多活最多 1 秒）
|
||
if int(time.time()) >= int(data.get("exp", 0)):
|
||
return None
|
||
return str(data.get("sub") or "token")
|
||
|
||
|
||
def _authenticate(authorization: Optional[str], x_api_key: Optional[str]) -> str:
|
||
"""统一鉴权，返回调用方名称用于日志。
|
||
|
||
支持两种请求头（同时兼容）：
|
||
Authorization: Bearer <token> ← 标准写法，推荐
|
||
X-API-Key: <token> ← 旧写法，保留兼容
|
||
token 可以是 API_TOKENS 里的静态 token，也可以是 /auth/token 签发的短期 token。
|
||
完全没配置鉴权时放行（返回 anonymous），方便本地调试。
|
||
"""
|
||
if not _TOKEN_MAP and not CLIENT_SECRET:
|
||
return "anonymous"
|
||
|
||
token = ""
|
||
if authorization and authorization.lower().startswith("bearer "):
|
||
token = authorization[7:].strip()
|
||
elif x_api_key:
|
||
token = x_api_key.strip()
|
||
|
||
if not token:
|
||
raise HTTPException(
|
||
status_code=401,
|
||
detail="missing token",
|
||
headers={"WWW-Authenticate": "Bearer"},
|
||
)
|
||
if token in _TOKEN_MAP:
|
||
return _TOKEN_MAP[token]
|
||
subject = _verify_short_token(token)
|
||
if subject:
|
||
return subject
|
||
raise HTTPException(
|
||
status_code=401,
|
||
detail="invalid or expired token",
|
||
headers={"WWW-Authenticate": "Bearer"},
|
||
)
|
||
|
||
|
||
def _load_image(data: bytes) -> Image.Image:
|
||
"""读图 + 尺寸收敛。低配关键：图越大，解码缓冲和 det 张量越吃内存。"""
|
||
if len(data) > MAX_IMAGE_BYTES:
|
||
raise HTTPException(status_code=413, detail="image too large, max %d bytes" % MAX_IMAGE_BYTES)
|
||
try:
|
||
img = Image.open(io.BytesIO(data)) # 惰性：只读 header，不解码
|
||
except Exception:
|
||
raise HTTPException(status_code=400, detail="invalid image file")
|
||
# 解压炸弹防护：解码前按 header 尺寸拦截。
|
||
# 6000px 是硬上限（解码后约 108MB 缓冲，可接受）；手机常见 3000-4000px 拍照
|
||
# 必须放行、交给下面的 thumbnail 缩小，不能用 MAX_SIDE 当阈值。
|
||
if max(img.size) > MAX_DIMENSION:
|
||
raise HTTPException(
|
||
status_code=413, detail="image dimension too large, max %d px" % MAX_DIMENSION
|
||
)
|
||
try:
|
||
img.load()
|
||
except Exception:
|
||
raise HTTPException(status_code=400, detail="invalid image file")
|
||
if img.mode != "RGB":
|
||
img = img.convert("RGB")
|
||
if max(img.size) > MAX_SIDE:
|
||
img.thumbnail((MAX_SIDE, MAX_SIDE), Image.LANCZOS)
|
||
return img
|
||
|
||
|
||
def _run_ocr(img: Image.Image) -> dict:
|
||
engine = get_engine()
|
||
started = time.perf_counter()
|
||
result, _elapse = engine(np.array(img))
|
||
lines = id_parser._lines_from_result(result)
|
||
data = id_parser.parse(lines)
|
||
data["elapsed_ms"] = round((time.perf_counter() - started) * 1000, 1)
|
||
if not lines:
|
||
raise HTTPException(status_code=422, detail="no text detected in image")
|
||
if data["field_count"] == 0:
|
||
raise HTTPException(status_code=422, detail="no id card fields detected")
|
||
return data
|
||
|
||
|
||
def _mem_mb() -> Optional[int]:
|
||
"""容器/进程内存读数，方便 docker stats 之外直接看接口。"""
|
||
for path in (
|
||
"/sys/fs/cgroup/memory.current", # cgroup v2
|
||
"/sys/fs/cgroup/memory/memory.usage_in_bytes", # cgroup v1
|
||
):
|
||
try:
|
||
with open(path) as f:
|
||
return int(f.read().strip()) // (1024 * 1024)
|
||
except (OSError, ValueError):
|
||
continue
|
||
try:
|
||
with open("/proc/self/status") as f:
|
||
for line in f:
|
||
if line.startswith("VmRSS:"):
|
||
return int(line.split()[1]) // 1024
|
||
except OSError:
|
||
pass
|
||
return None
|
||
|
||
|
||
@app.get("/")
|
||
def root():
|
||
"""根路径给个可读的入口说明，避免直接访问域名时看到 404。"""
|
||
return {
|
||
"service": "idcard-api",
|
||
"version": app.version,
|
||
"status": "running",
|
||
"docs": "/docs",
|
||
"endpoints": {
|
||
"health": "GET /api/v1/idcard/health",
|
||
"recognize": "POST /api/v1/idcard/recognize (multipart, 字段名 file)",
|
||
"recognize_base64": "POST /api/v1/idcard/recognize/base64 (JSON)",
|
||
"token": "POST /api/v1/auth/token (需配置 CLIENT_SECRET)",
|
||
},
|
||
"auth": {
|
||
"enabled": bool(_TOKEN_MAP) or bool(CLIENT_SECRET),
|
||
"scheme": "Authorization: Bearer <token>（兼容 X-API-Key 头）",
|
||
"short_lived_token": bool(CLIENT_SECRET),
|
||
},
|
||
}
|
||
|
||
|
||
@app.get("/api/v1/idcard/health")
|
||
def health():
|
||
return {
|
||
"status": "ok",
|
||
"engine_loaded": _engine is not None,
|
||
"mem_mb": _mem_mb(),
|
||
"auth_enabled": bool(_TOKEN_MAP) or bool(CLIENT_SECRET),
|
||
}
|
||
|
||
|
||
@app.post("/api/v1/auth/token")
|
||
def create_token(body: TokenRequest):
|
||
"""用 client_id + client_secret 换一个短期 token（默认 1 小时）。
|
||
|
||
适合前端场景：不在页面里硬编码长期密钥，只持有短时 token。
|
||
"""
|
||
if not CLIENT_SECRET:
|
||
raise HTTPException(status_code=400, detail="token endpoint disabled, CLIENT_SECRET not set")
|
||
if not hmac.compare_digest(body.client_secret, CLIENT_SECRET):
|
||
raise HTTPException(status_code=401, detail="invalid client credentials")
|
||
payload = issue_token(body.client_id or CLIENT_ID)
|
||
log.info("token issued for client=%s", body.client_id or CLIENT_ID)
|
||
return payload
|
||
|
||
|
||
@app.post("/api/v1/idcard/recognize")
|
||
async def recognize(
|
||
file: UploadFile = File(..., description="身份证图片 jpg/png"),
|
||
side: str = "auto",
|
||
mask: bool = False,
|
||
authorization: Optional[str] = Header(None),
|
||
x_api_key: Optional[str] = Header(None),
|
||
):
|
||
caller = _authenticate(authorization, x_api_key)
|
||
if file.content_type and not file.content_type.startswith("image/"):
|
||
raise HTTPException(status_code=400, detail="content type must be image/*")
|
||
data = _run_ocr(_load_image(await file.read()))
|
||
if side in ("front", "back"):
|
||
data["side"] = side
|
||
# 只记调用方和耗时，不记姓名/身份证号（合规要求）
|
||
log.info("recognize by=%s side=%s %sms", caller, data["side"], data["elapsed_ms"])
|
||
return _respond(data, mask)
|
||
|
||
|
||
@app.post("/api/v1/idcard/recognize/base64")
|
||
def recognize_base64(
|
||
body: Base64Body,
|
||
authorization: Optional[str] = Header(None),
|
||
x_api_key: Optional[str] = Header(None),
|
||
):
|
||
caller = _authenticate(authorization, x_api_key)
|
||
raw = body.image_base64
|
||
if "," in raw and raw.strip().startswith("data:"):
|
||
raw = raw.split(",", 1)[1]
|
||
try:
|
||
payload = base64.b64decode(raw, validate=True)
|
||
except (binascii.Error, ValueError):
|
||
raise HTTPException(status_code=400, detail="invalid base64 payload")
|
||
out = _run_ocr(_load_image(payload))
|
||
if body.side in ("front", "back"):
|
||
out["side"] = body.side
|
||
log.info("recognize_b64 by=%s side=%s %sms", caller, out["side"], out["elapsed_ms"])
|
||
return _respond(out, body.mask)
|
||
|
||
|
||
def _respond(data: dict, mask: bool) -> JSONResponse:
|
||
return JSONResponse({"code": 0, "data": id_parser.mask(data) if mask else data})
|