mirror of
https://github.com/Tencent/WeKnora.git
synced 2026-09-01 14:53:07 +08:00
begoniezhao
279 lines
10 KiB
Python
279 lines
10 KiB
Python
import os
|
|
import logging
|
|
import base64
|
|
from typing import Optional, Union, Dict, Any
|
|
from abc import ABC, abstractmethod
|
|
from PIL import Image
|
|
import io
|
|
import numpy as np
|
|
from .image_utils import image_to_base64
|
|
|
|
logger = logging.getLogger(__name__)
|
|
logger.setLevel(logging.INFO)
|
|
|
|
class OCRBackend(ABC):
|
|
"""Base class for OCR backends"""
|
|
|
|
@abstractmethod
|
|
def predict(self, image: Union[str, bytes, Image.Image]) -> str:
|
|
"""Extract text from an image
|
|
|
|
Args:
|
|
image: Image file path, bytes, or PIL Image object
|
|
|
|
Returns:
|
|
Extracted text
|
|
"""
|
|
pass
|
|
|
|
class PaddleOCRBackend(OCRBackend):
|
|
"""PaddleOCR backend implementation"""
|
|
|
|
def __init__(self, **kwargs):
|
|
"""Initialize PaddleOCR backend"""
|
|
self.ocr = None
|
|
try:
|
|
import os
|
|
import paddle
|
|
|
|
# Set PaddlePaddle to use CPU and disable GPU
|
|
os.environ['CUDA_VISIBLE_DEVICES'] = ''
|
|
paddle.set_device('cpu')
|
|
|
|
# 尝试检测CPU是否支持AVX指令集
|
|
try:
|
|
import subprocess
|
|
import platform
|
|
|
|
# 检测CPU是否支持AVX
|
|
if platform.system() == "Linux":
|
|
try:
|
|
result = subprocess.run(['grep', '-o', 'avx', '/proc/cpuinfo'],
|
|
capture_output=True, text=True, timeout=5)
|
|
has_avx = 'avx' in result.stdout.lower()
|
|
if not has_avx:
|
|
logger.warning("CPU does not support AVX instructions, using compatibility mode")
|
|
# 进一步限制指令集使用
|
|
os.environ['FLAGS_use_avx2'] = '0'
|
|
os.environ['FLAGS_use_avx'] = '1'
|
|
except (subprocess.TimeoutExpired, FileNotFoundError, subprocess.SubprocessError):
|
|
logger.warning("Could not detect AVX support, using compatibility mode")
|
|
os.environ['FLAGS_use_avx2'] = '0'
|
|
os.environ['FLAGS_use_avx'] = '1'
|
|
except Exception as e:
|
|
logger.warning(f"Error detecting CPU capabilities: {e}, using compatibility mode")
|
|
os.environ['FLAGS_use_avx2'] = '0'
|
|
os.environ['FLAGS_use_avx'] = '1'
|
|
|
|
from paddleocr import PaddleOCR
|
|
# OCR configuration with text orientation classification enabled
|
|
ocr_config = {
|
|
"use_gpu": False,
|
|
"text_det_limit_type": "max",
|
|
"text_det_limit_side_len": 960,
|
|
"use_doc_orientation_classify": True, # 启用文档方向分类
|
|
"use_doc_unwarping": False,
|
|
"use_textline_orientation": True, # 启用文本行方向检测
|
|
"text_recognition_model_name": "PP-OCRv4_server_rec",
|
|
"text_detection_model_name": "PP-OCRv4_server_det",
|
|
"text_det_thresh": 0.3,
|
|
"text_det_box_thresh": 0.6,
|
|
"text_det_unclip_ratio": 1.5,
|
|
"text_rec_score_thresh": 0.0,
|
|
"ocr_version": "PP-OCRv4",
|
|
"lang": "ch",
|
|
"show_log": False,
|
|
"use_dilation": True, # improves accuracy
|
|
"det_db_score_mode": "slow", # improves accuracy
|
|
}
|
|
|
|
self.ocr = PaddleOCR(**ocr_config)
|
|
logger.info("PaddleOCR engine initialized successfully")
|
|
|
|
except ImportError as e:
|
|
logger.error(f"Failed to import paddleocr: {str(e)}. Please install it with 'pip install paddleocr'")
|
|
except OSError as e:
|
|
if "Illegal instruction" in str(e) or "core dumped" in str(e):
|
|
logger.error(f"PaddlePaddle crashed due to CPU instruction set incompatibility: {str(e)}")
|
|
logger.error("This usually happens when the CPU doesn't support AVX instructions.")
|
|
logger.error("Please try installing a CPU-only version of PaddlePaddle or use a different OCR backend.")
|
|
else:
|
|
logger.error(f"Failed to initialize PaddleOCR due to OS error: {str(e)}")
|
|
except Exception as e:
|
|
logger.error(f"Failed to initialize PaddleOCR: {str(e)}")
|
|
|
|
def predict(self, image):
|
|
"""Perform OCR recognition on the image
|
|
|
|
Args:
|
|
image: Image object (PIL.Image or numpy array)
|
|
|
|
Returns:
|
|
Extracted text string
|
|
"""
|
|
try:
|
|
# Ensure image is in RGB format
|
|
if hasattr(image, "convert") and image.mode != "RGB":
|
|
image = image.convert("RGB")
|
|
|
|
# Convert to numpy array if needed
|
|
if hasattr(image, "convert"):
|
|
image_array = np.array(image)
|
|
else:
|
|
image_array = image
|
|
|
|
# Perform OCR
|
|
ocr_result = self.ocr.ocr(image_array, cls=False)
|
|
|
|
# Extract text
|
|
ocr_text = ""
|
|
if ocr_result and ocr_result[0]:
|
|
for line in ocr_result[0]:
|
|
if line and len(line) >= 2:
|
|
text = line[1][0] if line[1] else ""
|
|
if text:
|
|
ocr_text += text + " "
|
|
|
|
text_length = len(ocr_text.strip())
|
|
if text_length > 0:
|
|
logger.info(f"OCR extracted {text_length} characters")
|
|
return ocr_text.strip()
|
|
else:
|
|
logger.warning("OCR returned empty result")
|
|
return ""
|
|
|
|
except Exception as e:
|
|
logger.error(f"OCR recognition error: {str(e)}")
|
|
return ""
|
|
|
|
class NanonetsOCRBackend(OCRBackend):
|
|
"""Nanonets OCR backend implementation using OpenAI API format"""
|
|
|
|
def __init__(self, **kwargs):
|
|
"""Initialize Nanonets OCR backend
|
|
|
|
Args:
|
|
api_key: API key for OpenAI API
|
|
base_url: Base URL for OpenAI API
|
|
model: Model name
|
|
"""
|
|
try:
|
|
from openai import OpenAI
|
|
self.api_key = kwargs.get("api_key", "123")
|
|
self.base_url = kwargs.get("base_url", "http://localhost:8000/v1")
|
|
self.model = kwargs.get("model", "nanonets/Nanonets-OCR-s")
|
|
self.temperature = kwargs.get("temperature", 0.0)
|
|
self.max_tokens = kwargs.get("max_tokens", 15000)
|
|
|
|
self.client = OpenAI(api_key=self.api_key, base_url=self.base_url)
|
|
self.prompt = """
|
|
## 任务说明
|
|
|
|
请从上传的文档中提取文字内容，严格按自然阅读顺序（从上到下，从左到右）输出，并遵循以下格式规范。
|
|
|
|
### 1. **文本处理**
|
|
|
|
* 按正常阅读顺序提取文字，语句流畅自然。
|
|
|
|
### 2. **表格**
|
|
|
|
* 所有表格统一转换为 **Markdown 表格格式**。
|
|
* 内容保持清晰、对齐整齐，便于阅读。
|
|
|
|
### 3. **公式**
|
|
|
|
* 所有公式转换为 **LaTeX 格式**，使用 `$$公式$$` 包裹。
|
|
|
|
### 4. **图片**
|
|
|
|
* 忽略图片信息
|
|
|
|
### 5. **链接**
|
|
|
|
* 不要猜测或补全不确定的链接地址。
|
|
"""
|
|
logger.info(f"Nanonets OCR engine initialized with model: {self.model}")
|
|
except ImportError:
|
|
logger.error("Failed to import openai. Please install it with 'pip install openai'")
|
|
self.client = None
|
|
except Exception as e:
|
|
logger.error(f"Failed to initialize Nanonets OCR: {str(e)}")
|
|
self.client = None
|
|
|
|
def predict(self, image: Union[str, bytes, Image.Image]) -> str:
|
|
"""Extract text from an image using Nanonets OCR
|
|
|
|
Args:
|
|
image: Image file path, bytes, or PIL Image object
|
|
|
|
Returns:
|
|
Extracted text
|
|
"""
|
|
if self.client is None:
|
|
logger.error("Nanonets OCR client not initialized")
|
|
return ""
|
|
|
|
try:
|
|
# Encode image to base64
|
|
img_base64 = image_to_base64(image)
|
|
if not img_base64:
|
|
return ""
|
|
|
|
# Call Nanonets OCR API
|
|
logger.info(f"Calling Nanonets OCR API with model: {self.model}")
|
|
response = self.client.chat.completions.create(
|
|
model=self.model,
|
|
messages=[
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"type": "image_url",
|
|
"image_url": {"url": f"data:image/png;base64,{img_base64}"},
|
|
},
|
|
{
|
|
"type": "text",
|
|
"text": self.prompt,
|
|
},
|
|
],
|
|
}
|
|
],
|
|
temperature=self.temperature,
|
|
max_tokens=self.max_tokens
|
|
)
|
|
|
|
return response.choices[0].message.content
|
|
except Exception as e:
|
|
logger.error(f"Nanonets OCR prediction error: {str(e)}")
|
|
return ""
|
|
|
|
class OCREngine:
|
|
"""OCR Engine factory class"""
|
|
|
|
_instance = None
|
|
|
|
@classmethod
|
|
def get_instance(cls, backend_type="paddle", **kwargs) -> Optional[OCRBackend]:
|
|
"""Get OCR engine instance
|
|
|
|
Args:
|
|
backend_type: OCR backend type, one of: "paddle", "nanonets"
|
|
**kwargs: Additional arguments for the backend
|
|
|
|
Returns:
|
|
OCR engine instance or None if initialization fails
|
|
"""
|
|
if cls._instance is None:
|
|
logger.info(f"Initializing OCR engine with backend: {backend_type}")
|
|
|
|
if backend_type.lower() == "paddle":
|
|
cls._instance = PaddleOCRBackend(**kwargs)
|
|
elif backend_type.lower() == "nanonets":
|
|
cls._instance = NanonetsOCRBackend(**kwargs)
|
|
else:
|
|
logger.error(f"Unknown OCR backend type: {backend_type}")
|
|
return None
|
|
|
|
return cls._instance
|
|
|