init

app.py

@@ -0,0 +1,231 @@
+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)