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

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

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")
 

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)

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
+            from src.knowledge.indexing import parse_image_async
-            text = ocr.process_image(str(file_path_obj))
+            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']:

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["id"] for chunk in chunks],
-                        [chunk["content"] for chunk in chunks], 
+                        [chunk["content"] for chunk in chunks],
-                        [chunk["source"] for chunk in chunks], 
+                        [chunk["source"] for chunk in chunks],
-                        [chunk["chunk_id"] for chunk in chunks], 
+                        [chunk["chunk_id"] for chunk in chunks],
-                        [chunk["file_id"] for chunk in chunks], 
+                        [chunk["file_id"] for chunk in chunks],
-                        [chunk["chunk_index"] for chunk in chunks], 
+                        [chunk["chunk_index"] for chunk in chunks],
                         embeddings
                     ]
 

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文件
+        使用Mineru OCR处理文件
-        :param pdf_path: PDF文件路径
+        :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.")
+            logger.debug(f"Mineru OCR result: {text[:50]}(...) total {len(text)} characters.")
-            return pdf_text
+            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')

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 {
+      &:hover {
-    background-color: var(--main-light-5);
+        background-color: var(--main-color) !important;
-    border-color: var(--main-color);
+        color: #fff !important;
-    color: var(--main-color);
+      }
+    }
   }
 
   &.active.expanded {
     transform: scaleX(1);
+    &:hover {
+      background-color: var(--main-light-5);
+      border-color: var(--main-color);
+      color: var(--main-color);
+    }
   }
 }
 </style>