from fastapi import FastAPI, UploadFile, File, Form, HTTPException, Body from rapidocr_onnxruntime import RapidOCR from pdf2image import convert_from_bytes from PIL import Image import numpy as np import io import time import logging import requests import os import tempfile import subprocess # 新增库引入 from openpyxl import load_workbook from pptx import Presentation import xlrd # 用于处理 .xls # 配置日志 logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) app = FastAPI(title="OCR Service") # 1. 全局初始化 OCR 模型 try: logger.info("正在加载OCR模型...") ocr_engine = RapidOCR() logger.info("OCR模型加载完成！") except Exception as e: logger.error(f"模型加载失败: {e}") raise e # --- 核心处理逻辑函数 --- def ocr_single_image(img_np): """辅助函数：处理单张图片的 Numpy 数组并返回文本""" try: result, _ = ocr_engine(img_np) txt_res = [] if result: for line in result: txt_res.append(line[1]) return "\n".join(txt_res) except Exception as e: logger.error(f"单图识别出错: {e}") return "" def extract_from_excel(file_bytes): """解析 Excel (.xlsx) 文本""" try: wb = load_workbook(io.BytesIO(file_bytes), data_only=True) texts = [] for sheet in wb.worksheets: texts.append(f"--- Sheet: {sheet.title} ---") for row in sheet.iter_rows(values_only=True): row_text = " ".join([str(cell) for cell in row if cell is not None]) if row_text.strip(): texts.append(row_text) return "\n".join(texts) except Exception as e: logger.error(f"Excel(.xlsx)解析失败: {e}") raise HTTPException(status_code=400, detail=f"xlsx解析失败: {str(e)}") def extract_from_xls(file_bytes): """解析旧版 Excel (.xls) 文本""" try: # xlrd 支持直接从内存读取 wb = xlrd.open_workbook(file_contents=file_bytes) texts = [] for sheet in wb.sheets(): texts.append(f"--- Sheet: {sheet.name} ---") for row_idx in range(sheet.nrows): row = sheet.row(row_idx) # xlrd 的 cell 对象需要取 .value row_text = " ".join([str(c.value) for c in row if c.value not in ('', None)]) if row_text.strip(): texts.append(row_text) return "\n".join(texts) except Exception as e: logger.error(f"Excel(.xls)解析失败: {e}") raise HTTPException(status_code=400, detail=f"xls解析失败: {str(e)}") def extract_from_ppt(file_bytes): """解析 PPT (.pptx) 文本""" try: prs = Presentation(io.BytesIO(file_bytes)) texts = [] for i, slide in enumerate(prs.slides): page_content = [] for shape in slide.shapes: if hasattr(shape, "text") and shape.text: page_content.append(shape.text) if page_content: texts.append(f"--- Slide {i+1} ---") texts.append("\n".join(page_content)) return "\n\n".join(texts) except Exception as e: logger.error(f"PPT(.pptx)解析失败: {e}") raise HTTPException(status_code=400, detail="pptx解析失败") def extract_from_ppt_legacy(file_bytes): """解析旧版 PPT (.ppt) 文本""" # .ppt 是二进制 OLE 格式，纯 Python 库支持极差 # 这里使用系统级工具 catppt (来自 catdoc 包) 进行提取 temp_file = None try: # 创建临时文件，因为 catppt 需要文件路径 with tempfile.NamedTemporaryFile(suffix=".ppt", delete=False) as tmp: tmp.write(file_bytes) temp_file = tmp.name # 调用系统命令 catppt # -d utf-8 尝试强制输出 utf-8 (视 catdoc 版本而定，通常默认即可) process = subprocess.Popen(['catppt', temp_file], stdout=subprocess.PIPE, stderr=subprocess.PIPE) stdout, stderr = process.communicate() if process.returncode != 0: logger.error(f"catppt error: {stderr.decode()}") raise Exception("解析进程异常退出") # 尝试解码 return stdout.decode('utf-8', errors='ignore') except Exception as e: logger.error(f"PPT(.ppt)解析失败: {e}") raise HTTPException(status_code=400, detail="ppt解析失败 (请检查文件是否损坏)") finally: # 清理临时文件 if temp_file and os.path.exists(temp_file): os.remove(temp_file) def process_file_bytes(file_bytes, filename): """通用文件流处理入口""" filename = filename.lower() # 1. PDF 处理 if filename.endswith('.pdf'): try: images = convert_from_bytes(file_bytes, dpi=200, fmt='jpeg') texts = [] for i, img in enumerate(images): img_np = np.array(img) texts.append(f"--- Page {i+1} ---\n" + ocr_single_image(img_np)) return "\n\n".join(texts) except Exception as e: raise HTTPException(status_code=500, detail="PDF解析失败") # 2. Excel 处理 elif filename.endswith('.xlsx'): return extract_from_excel(file_bytes) elif filename.endswith('.xls'): return extract_from_xls(file_bytes) # 3. PPT 处理 elif filename.endswith('.pptx'): return extract_from_ppt(file_bytes) elif filename.endswith('.ppt'): return extract_from_ppt_legacy(file_bytes) # 4. 图片处理 else: try: img = Image.open(io.BytesIO(file_bytes)) img_np = np.array(img) return ocr_single_image(img_np) except Exception: raise HTTPException(status_code=400, detail=f"不支持的文件格式: {filename}") # --- 接口定义 --- @app.get("/health") def health_check(): return {"status": "ok"} @app.post("/ocr") async def recognize_file(file: UploadFile = File(...)): """ 文件上传接口 支持: 图片, PDF, Excel(.xlsx/.xls), PPT(.pptx/.ppt) """ start_time = time.time() try: content = await file.read() result_text = process_file_bytes(content, file.filename) cost = (time.time() - start_time) * 1000 return {"code": 200, "data": result_text, "cost_time_ms": cost} except HTTPException as he: raise he except Exception as e: logger.error(f"处理异常: {e}") return {"code": 500, "msg": str(e)} @app.post("/ocr/url") async def recognize_url(url: str = Body(..., embed=True)): """ URL 识别接口 Body参数: {"url": "http://example.com/file.png"} """ start_time = time.time() try: logger.info(f"正在下载文件: {url}") resp = requests.get(url, timeout=15) if resp.status_code != 200: raise HTTPException(status_code=400, detail=f"下载失败，状态码: {resp.status_code}") filename = os.path.basename(url.split("?")[0]) if not filename or "." not in filename: content_type = resp.headers.get('Content-Type', '') if 'pdf' in content_type: filename = 'temp.pdf' elif 'sheet' in content_type or 'excel' in content_type: # 简单判断，这里默认给 xlsx，如果是 xls 可能需要更复杂的 magic number 判断 filename = 'temp.xlsx' elif 'presentation' in content_type: filename = 'temp.pptx' else: filename = 'temp.jpg' result_text = process_file_bytes(resp.content, filename) cost = (time.time() - start_time) * 1000 return {"code": 200, "data": result_text, "cost_time_ms": cost} except HTTPException as he: raise he except Exception as e: logger.error(f"URL处理异常: {e}") return {"code": 500, "msg": str(e)} if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=9000)