import os import hashlib import tempfile import threading import traceback from collections import OrderedDict from fastapi import FastAPI, UploadFile, File, HTTPException, Request from fastapi.responses import JSONResponse from starlette.middleware.base import BaseHTTPMiddleware import wcocr app = FastAPI(title="WeChat OCR API", version="1.0.0") # 全局变量 ocr_cache = None ocr_lock = None API_TOKEN = os.getenv("TOKEN", "my-secret-token") # 不需要验证的路径 EXCLUDED_PATHS = ["/health"] class TokenAuthMiddleware(BaseHTTPMiddleware): """简单的Token验证中间件""" async def dispatch(self, request: Request, call_next): # 跳过不需要验证的路径 if request.url.path in EXCLUDED_PATHS: return await call_next(request) # 从请求头获取 Authorization auth_header = request.headers.get("Authorization") # 兼容 Bearer 前缀 if auth_header and auth_header.startswith("Bearer "): auth_header = auth_header[7:] # 移除 "Bearer " 前缀 # 验证token是否匹配 if auth_header != API_TOKEN: return JSONResponse( status_code=401, content={"detail": "未授权：无效的认证令牌"} ) # token验证通过，继续处理请求 response = await call_next(request) return response # 添加中间件 app.add_middleware(TokenAuthMiddleware) class LRUCache: """简单的 LRU 缓存，防止内存溢出""" def __init__(self, capacity: int = 100): self.cache = OrderedDict() self.capacity = capacity def get(self, key): if key not in self.cache: return None self.cache.move_to_end(key) return self.cache[key] def put(self, key, value): if key in self.cache: self.cache.move_to_end(key) self.cache[key] = value if len(self.cache) > self.capacity: self.cache.popitem(last=False) @app.on_event("startup") async def startup_event(): global ocr_cache, ocr_lock print("====== 开始初始化 WeChat OCR ======") try: # 初始化 WeChat OCR WECHAT_PATH = os.environ.get("WECHAT_PATH", "/code/wechat") WXOCR_PATH = os.environ.get("WXOCR_PATH", "/code/wechat/wxocr") wcocr.init(WXOCR_PATH, WECHAT_PATH) # 全局缓存对象 (缓存 100 张最近的图片结果) ocr_cache = LRUCache(capacity=100) # OCR 引擎并发锁 ocr_lock = threading.Lock() print("====== 🎉 WeChat OCR 初始化成功，API 已就绪！ ======") except Exception as e: print("❌ 初始化失败, 错误如下:\n", str(e)) @app.get("/health") def health_check(): return {"status": "healthy"} def process_ocr(file_bytes: bytes) -> dict: """处理图片并执行 OCR，包含缓存机制""" # 计算图片 Hash 用于缓存 Key img_hash = hashlib.md5(file_bytes).hexdigest() cached_result = ocr_cache.get(img_hash) if ocr_cache else None if cached_result is not None: return cached_result if not file_bytes or len(file_bytes) < 8: return {"errcode": 400, "errmsg": "Empty or invalid file content"} # 使用轻量级的 Magic Bytes（文件头魔数）检查代替 cv2，判断是否是常见图片格式 header = file_bytes[:12] is_image = ( header.startswith(b'\x89PNG\r\n\x1a\n') or # PNG header.startswith(b'\xff\xd8\xff') or # JPEG header.startswith(b'BM') or # BMP header.startswith(b'GIF8') or # GIF (header.startswith(b'RIFF') and header[8:12] == b'WEBP') # WEBP ) if not is_image: return {"errcode": 400, "errmsg": "Unsupported or invalid image format"} # 写入临时文件供 wcocr 调用 with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp_file: tmp_file.write(file_bytes) tmp_file.flush() tmp_path = tmp_file.name try: # 加锁执行 OCR，防止底层组件崩溃 if ocr_lock: with ocr_lock: result = wcocr.ocr(tmp_path) else: result = wcocr.ocr(tmp_path) finally: # 清理临时文件 if os.path.exists(tmp_path): os.remove(tmp_path) # 存入缓存 if ocr_cache: ocr_cache.put(img_hash, result) return result def format_structured_text(ocr_response: list) -> str: """基于坐标的文本行合并算法，还原原始文本的换行结构""" if not ocr_response: return "" # 按 Y 轴 (top) 排序 boxes = sorted(ocr_response, key=lambda x: x['top']) lines = [] current_line = [] for box in boxes: if not current_line: current_line.append(box) continue # 计算当前行的中心基线 line_top_avg = sum(b['top'] for b in current_line) / len(current_line) line_bottom_avg = sum(b['bottom'] for b in current_line) / len(current_line) line_center = (line_top_avg + line_bottom_avg) / 2 # 如果当前 box 的 top 小于上一行的中心点，认为在同一行 if box['top'] < line_center: current_line.append(box) else: lines.append(current_line) current_line = [box] if current_line: lines.append(current_line) # 将同一行的文本按 X 轴 (left) 排序并拼接 structured_text = [] for line in lines: line.sort(key=lambda x: x['left']) line_text = " ".join([b['text'] for b in line]) structured_text.append(line_text) return "\n".join(structured_text) # ========================================== # 接口 1：OCR 识别后的原始结构 # ========================================== @app.post("/api/ocr/raw") async def api_ocr_raw(file: UploadFile = File(...)): try: contents = await file.read() raw_result = process_ocr(contents) if raw_result.get('errcode', 0) != 0: return JSONResponse(status_code=400, content=raw_result) return raw_result except Exception as e: print(f"❌ [api_ocr_raw] 发生异常: {str(e)}") traceback.print_exc() raise HTTPException( status_code=500, detail=f"OCR raw extraction error: {str(e)}") # ========================================== # 接口 2：识别后的纯文本 (无分隔符直接拼接) # ========================================== @app.post("/api/ocr/text") async def api_ocr_text(file: UploadFile = File(...)): try: contents = await file.read() raw_result = process_ocr(contents) if raw_result.get('errcode', 0) != 0: return JSONResponse(status_code=400, content=raw_result) ocr_response = raw_result.get('ocr_response', []) pure_text = "".join([item['text'] for item in ocr_response]) return { "errcode": 0, "text": pure_text } except Exception as e: print(f"❌ [api_ocr_text] 发生异常: {str(e)}") traceback.print_exc() raise HTTPException( status_code=500, detail=f"OCR text extraction error: {str(e)}") # ========================================== # 接口 3：识别后的纯文本，保留原始结构换行 # ========================================== @app.post("/api/ocr/structured") async def api_ocr_structured(file: UploadFile = File(...)): try: contents = await file.read() raw_result = process_ocr(contents) if raw_result.get('errcode', 0) != 0: return JSONResponse(status_code=400, content=raw_result) ocr_response = raw_result.get('ocr_response', []) structured_text = format_structured_text(ocr_response) return { "errcode": 0, "text": structured_text } except Exception as e: print(f"❌ [api_ocr_structured] 发生异常: {str(e)}") traceback.print_exc() raise HTTPException( status_code=500, detail=f"OCR structured extraction error: {str(e)}")