From cf3cb5ef33fecc0b4528740af2c8797178373a1d Mon Sep 17 00:00:00 2001 From: Wenjie Zhang Date: Sun, 27 Jul 2025 04:12:49 +0800 Subject: [PATCH] =?UTF-8?q?feat(ocr):=20=E5=A2=9E=E5=BC=BAOCR=E5=A4=84?= =?UTF-8?q?=E7=90=86=E5=8A=9F=E8=83=BD=EF=BC=8C=E6=94=AF=E6=8C=81=E5=9B=BE?= =?UTF-8?q?=E5=83=8F=E8=A7=A3=E6=9E=90=E5=92=8C=E5=8F=82=E6=95=B0=E4=BC=A0?= =?UTF-8?q?=E9=80=92=EF=BC=8C=E4=BC=98=E5=8C=96Docker=E7=8E=AF=E5=A2=83?= =?UTF-8?q?=E5=8F=98=E9=87=8F=E9=85=8D=E7=BD=AE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- docker-compose.yml | 2 +- server/routers/system_router.py | 3 +- src/knowledge/indexing.py | 52 +++++++++++++++++++++++++++--- src/knowledge/knowledge_base.py | 6 ++-- src/knowledge/milvus_kb.py | 12 +++---- src/plugins/_ocr.py | 35 +++++++++++--------- web/src/views/DataBaseInfoView.vue | 18 +++++++---- 7 files changed, 90 insertions(+), 38 deletions(-) diff --git a/docker-compose.yml b/docker-compose.yml index b3b5dc10..c5cf26cf 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -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 diff --git a/server/routers/system_router.py b/server/routers/system_router.py index 588e2c02..73adbd1d 100644 --- a/server/routers/system_router.py +++ b/server/routers/system_router.py @@ -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") diff --git a/src/knowledge/indexing.py b/src/knowledge/indexing.py index 63c36fee..cc993a7c 100644 --- a/src/knowledge/indexing.py +++ b/src/knowledge/indexing.py @@ -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) diff --git a/src/knowledge/knowledge_base.py b/src/knowledge/knowledge_base.py index 7cce1ff1..bf176cb0 100644 --- a/src/knowledge/knowledge_base.py +++ b/src/knowledge/knowledge_base.py @@ -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']: diff --git a/src/knowledge/milvus_kb.py b/src/knowledge/milvus_kb.py index 1ae5bd3b..be4d6fae 100644 --- a/src/knowledge/milvus_kb.py +++ b/src/knowledge/milvus_kb.py @@ -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 ] diff --git a/src/plugins/_ocr.py b/src/plugins/_ocr.py index ae9d96af..ab53194c 100644 --- a/src/plugins/_ocr.py +++ b/src/plugins/_ocr.py @@ -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') diff --git a/web/src/views/DataBaseInfoView.vue b/web/src/views/DataBaseInfoView.vue index 10295639..f13fed45 100644 --- a/web/src/views/DataBaseInfoView.vue +++ b/web/src/views/DataBaseInfoView.vue @@ -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); + } } }