kris/ForcePilot
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

feat(ocr): 增强OCR处理功能，支持图像解析和参数传递，优化Docker环境变量配置

Browse Source
This commit is contained in:
Wenjie Zhang 2025-07-27 04:12:49 +08:00
parent c14bd06db6
commit cf3cb5ef33
7 changed files with 90 additions and 38 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
docker-compose.yml
View File

@ -29,7 +29,7 @@ services:
- MINERU_OCR_URI=${MINERU_OCR_URI:-http://mineru:30000}
- PADDLEX_URI=${PADDLEX_URI:-http://paddlex:8080}
- MINIO_URI=${MINIO_URI:-http://milvus-minio:9000}
- MODEL_DIR=/models
- MODEL_DIR_IN_DOCKER=/models
- RUNNING_IN_DOCKER=true
- NO_PROXY=localhost,127.0.0.1,milvus,graph,milvus-minio,milvus-etcd-dev,etcd,minio,mineru,paddlex
- no_proxy=localhost,127.0.0.1,milvus,graph,milvus-minio,milvus-etcd-dev,etcd,minio,mineru,paddlex

3
server/routers/system_router.py
View File

@ -192,7 +192,8 @@ async def check_ocr_services_health(current_user: User = Depends(get_admin_user)
# 检查 RapidOCR (ONNX) 模型
try:
model_dir = os.path.join(os.getenv("MODEL_DIR", ""), "SWHL/RapidOCR")
model_dir_root = os.getenv("MODEL_DIR") if not os.getenv("RUNNING_IN_DOCKER") else os.getenv("MODEL_DIR_IN_DOCKER")
model_dir = os.path.join(model_dir_root, "SWHL/RapidOCR")
det_model_path = os.path.join(model_dir, "PP-OCRv4/ch_PP-OCRv4_det_infer.onnx")
rec_model_path = os.path.join(model_dir, "PP-OCRv4/ch_PP-OCRv4_rec_infer.onnx")

52
src/knowledge/indexing.py
View File

@ -146,18 +146,18 @@ def parse_pdf(file, params=None):
try:
if opt_ocr == "onnx_rapid_ocr":
from src.plugins import ocr
return ocr.process_pdf(file)
return ocr.process_pdf(file, params=params)
elif opt_ocr == "mineru_ocr":
from src.plugins import ocr
return ocr.process_pdf_mineru(file)
return ocr.process_file_mineru(file, params=params)
elif opt_ocr == "paddlex_ocr":
from src.plugins import ocr
return ocr.process_pdf_paddlex(file)
return ocr.process_file_paddlex(file, params=params)
else:
return pdfreader(file, params=params)
raise ValueError(f"不支持的OCR方式: {opt_ocr}")
except OCRServiceException as e:
logger.error(f"OCR service failed: {e.service_name} - {str(e)}")
@ -170,5 +170,49 @@ def parse_pdf(file, params=None):
"parsing_failed"
)
def parse_image(file, params=None):
"""
解析图像文件，支持多种OCR方式
"""
from src.plugins._ocr import OCRServiceException
params = params or {}
opt_ocr = params.get("enable_ocr", "disable")
if opt_ocr == "disable":
logger.warning(f"OCR is disabled for image file: {file}, Using `onnx_rapid_ocr` instead")
opt_ocr = "onnx_rapid_ocr"
try:
if opt_ocr == "onnx_rapid_ocr":
from src.plugins import ocr
return ocr.process_image(file, params=params)
elif opt_ocr == "mineru_ocr":
from src.plugins import ocr
return ocr.process_file_mineru(file, params=params)
elif opt_ocr == "paddlex_ocr":
from src.plugins import ocr
return ocr.process_file_paddlex(file, params=params)
else:
raise ValueError(f"不支持的OCR方式: {opt_ocr}")
except OCRServiceException as e:
logger.error(f"OCR service failed: {e.service_name} - {str(e)}")
raise
except Exception as e:
logger.error(f"Image parsing failed: {str(e)}")
raise OCRServiceException(
f"Image解析失败: {str(e)}",
opt_ocr,
"parsing_failed"
)
async def parse_pdf_async(file, params=None):
return await asyncio.to_thread(parse_pdf, file, params=params)
async def parse_image_async(file, params=None):
return await asyncio.to_thread(parse_image, file, params=params)

6
src/knowledge/knowledge_base.py
View File

@ -411,10 +411,10 @@ class KnowledgeBase(ABC):
text = '\n'.join([para.text for para in doc.paragraphs])
return f"# {file_path_obj.name}\n\n{text}"
elif file_ext in ['.jpg', '.jpeg', '.png', '.bmp']:
elif file_ext in ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif']:
# 使用 OCR 处理图片
from src.plugins import ocr
text = ocr.process_image(str(file_path_obj))
from src.knowledge.indexing import parse_image_async
text = await parse_image_async(str(file_path_obj), params=params)
return f"# {file_path_obj.name}\n\n{text}"
elif file_ext in ['.html', '.htm']:

12
src/knowledge/milvus_kb.py
View File

@ -270,12 +270,12 @@ class MilvusKB(KnowledgeBase):
embeddings = await embedding_function(texts)
entities = [
[chunk["id"] for chunk in chunks],
[chunk["content"] for chunk in chunks],
[chunk["source"] for chunk in chunks],
[chunk["chunk_id"] for chunk in chunks],
[chunk["file_id"] for chunk in chunks],
[chunk["chunk_index"] for chunk in chunks],
[chunk["id"] for chunk in chunks],
[chunk["content"] for chunk in chunks],
[chunk["source"] for chunk in chunks],
[chunk["chunk_id"] for chunk in chunks],
[chunk["file_id"] for chunk in chunks],
[chunk["chunk_index"] for chunk in chunks],
embeddings
]

35
src/plugins/_ocr.py
View File

@ -70,11 +70,12 @@ class OCRPlugin:
def __init__(self, **kwargs):
self.ocr = None
self.det_box_thresh = kwargs.get('det_box_thresh', 0.3)
self.model_dir_root = os.getenv("MODEL_DIR") if not os.getenv("RUNNING_IN_DOCKER") else os.getenv("MODEL_DIR_IN_DOCKER")
def _check_rapid_ocr_availability(self):
"""检查RapidOCR模型是否可用"""
try:
model_dir = os.path.join(os.getenv("MODEL_DIR", ""), "SWHL/RapidOCR")
model_dir = os.path.join(self.model_dir_root, "SWHL/RapidOCR")
det_model_dir = os.path.join(model_dir, "PP-OCRv4/ch_PP-OCRv4_det_infer.onnx")
rec_model_dir = os.path.join(model_dir, "PP-OCRv4/ch_PP-OCRv4_rec_infer.onnx")
@ -111,7 +112,7 @@ class OCRPlugin:
# 先检查模型可用性
self._check_rapid_ocr_availability()
model_dir = os.path.join(os.getenv("MODEL_DIR", ""), "SWHL/RapidOCR")
model_dir = os.path.join(self.model_dir_root, "SWHL/RapidOCR")
det_model_dir = os.path.join(model_dir, "PP-OCRv4/ch_PP-OCRv4_det_infer.onnx")
rec_model_dir = os.path.join(model_dir, "PP-OCRv4/ch_PP-OCRv4_rec_infer.onnx")
@ -125,7 +126,7 @@ class OCRPlugin:
"load_failed"
)
def process_image(self, image):
def process_image(self, image, params=None):
"""
对单张图像执行OCR并提取文本
@ -134,7 +135,7 @@ class OCRPlugin:
- str: 图像文件路径
- PIL.Image: PIL图像对象
- numpy.ndarray: numpy图像数组
params: 参数
Returns:
str: 提取的文本内容
"""
@ -207,10 +208,11 @@ class OCRPlugin:
return image_path
def process_pdf(self, pdf_path):
def process_pdf(self, pdf_path, params=None):
"""
处理PDF文件并提取文本
:param pdf_path: PDF文件路径
:param params: 参数
:return: 提取的文本
"""
@ -242,10 +244,11 @@ class OCRPlugin:
logger.error(f"PDF processing error: {str(e)}")
return ""
def process_pdf_mineru(self, pdf_path):
def process_file_mineru(self, file_path, params=None):
"""
使用Mineru OCR处理PDF文件
:param pdf_path: PDF文件路径
使用Mineru OCR处理文件
:param file_path: 文件路径
:param params: 参数
:return: 提取的文本
"""
import requests
@ -281,23 +284,23 @@ class OCRPlugin:
try:
start_time = time.time()
pdf_path_list = [pdf_path]
file_path_list = [file_path]
output_dir = os.path.join(os.getcwd(), "tmp", "mineru_ocr")
pdf_text = parse_doc(pdf_path_list, output_dir,
text = parse_doc(file_path_list, output_dir,
backend="vlm-sglang-client",
server_url=mineru_ocr_uri)[0]
processing_time = time.time() - start_time
log_ocr_request("mineru_ocr", pdf_path, True, processing_time)
log_ocr_request("mineru_ocr", file_path, True, processing_time)
logger.debug(f"Mineru OCR result: {pdf_text[:50]}(...) total {len(pdf_text)} characters.")
return pdf_text
logger.debug(f"Mineru OCR result: {text[:50]}(...) total {len(text)} characters.")
return text
except Exception as e:
processing_time = time.time() - start_time
error_msg = f"MinerU OCR处理失败: {str(e)}"
log_ocr_request("mineru_ocr", pdf_path, False, processing_time, error_msg)
log_ocr_request("mineru_ocr", file_path, False, processing_time, error_msg)
raise OCRServiceException(
error_msg,
@ -305,10 +308,11 @@ class OCRPlugin:
"processing_failed"
)
def process_pdf_paddlex(self, pdf_path):
def process_file_paddlex(self, pdf_path, params=None):
"""
使用Paddlex OCR处理PDF文件
:param pdf_path: PDF文件路径
:param params: 参数
:return: 提取的文本
"""
from .paddlex import analyze_document, check_paddlex_health
@ -382,7 +386,6 @@ def plainreader(file_path):
return text
if __name__ == "__main__":
parser = ArgumentParser()
parser.add_argument('--pdf-path', type=str, required=True, help='Path to the PDF file')

18
web/src/views/DataBaseInfoView.vue
View File

@ -2933,20 +2933,24 @@ const getFileIconColor = (filename) => {
&.auto-refresh-btn {
&.ant-btn-primary {
background-color: var(--main-color); /* Ant Design primary color */
background-color: var(--main-color);
border-color: var(--main-color);
color: #fff;
}
}
&:hover {
background-color: var(--main-light-5);
border-color: var(--main-color);
color: var(--main-color);
&:hover {
background-color: var(--main-color) !important;
color: #fff !important;
}
}
}
&.active.expanded {
transform: scaleX(1);
&:hover {
background-color: var(--main-light-5);
border-color: var(--main-color);
color: var(--main-color);
}
}
}
</style>