From f6e0f6d89ebbae500f9223215311fc27d2f8c296 Mon Sep 17 00:00:00 2001 From: cin-niko Date: Sat, 7 Mar 2026 16:38:23 +0000 Subject: [PATCH] fix: paddleocr output adapter --- .../kotaemon/indices/ingests/files.py | 6 +- .../loaders/paddleocr_loader/adapter.py | 164 +++++++++++++++- .../paddleocr_loader/paddleocr_vl_loader.py | 157 +-------------- .../paddleocr_loader/ppstructure_v3_loader.py | 179 +++--------------- libs/ktem/ktem/index/file/pipelines.py | 24 ++- libs/ktem/ktem/pages/chat/__init__.py | 6 +- 6 files changed, 216 insertions(+), 320 deletions(-) diff --git a/libs/kotaemon/kotaemon/indices/ingests/files.py b/libs/kotaemon/kotaemon/indices/ingests/files.py index 3b88e1d5..4ab5b588 100644 --- a/libs/kotaemon/kotaemon/indices/ingests/files.py +++ b/libs/kotaemon/kotaemon/indices/ingests/files.py @@ -40,9 +40,9 @@ adobe_reader.vlm_endpoint = ( azure_reader.vlm_endpoint ) = docling_reader.vlm_endpoint = getattr(flowsettings, "KH_VLM_ENDPOINT", "") -paddle_device = str(config("PADDLE_DEVICE", default="gpu:0")) -paddle_struct_reader = PPStructureV3Reader(device=paddle_device) -paddle_vl_reader = PaddleOCRVLReader(device=paddle_device) +# paddle_device = str(config("PADDLE_DEVICE", default="gpu:0")) +paddle_struct_reader = PPStructureV3Reader() +paddle_vl_reader = PaddleOCRVLReader() KH_DEFAULT_FILE_EXTRACTORS: dict[str, BaseReader] = { diff --git a/libs/kotaemon/kotaemon/loaders/paddleocr_loader/adapter.py b/libs/kotaemon/kotaemon/loaders/paddleocr_loader/adapter.py index 3b18d630..2389ea57 100644 --- a/libs/kotaemon/kotaemon/loaders/paddleocr_loader/adapter.py +++ b/libs/kotaemon/kotaemon/loaders/paddleocr_loader/adapter.py @@ -1,31 +1,183 @@ +"""PaddleOCR result adapter for converting raw output to Documents.""" + import re -from abc import abstractmethod -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path from typing import Any from kotaemon.base import Document +# Label categories for PaddleOCR models (both PPStructureV3 and PaddleOCRVL) +# These are the block_label values from PaddleOCR output + +TEXT_LABELS: set[str] = { + "text", + "paragraph_title", + "doc_title", + "abstract", + "content", + "footnote", + "reference", + "reference_content", + "aside_text", + "algorithm", +} + +TABLE_LABELS: set[str] = {"table"} + +IMAGE_LABELS: set[str] = { + "image", + "chart", +} + +FORMULA_LABELS: set[str] = { + "formula", + "display_formula", + "inline_formula", +} + +# Labels to ignore (not useful for RAG) +IGNORE_LABELS: set[str] = { + "footer", + "footer_image", + "formula_number", + "figure_title", + "figure_table_chart_title", + "header", + "header_image", + "number", + "seal", + "vision_footnote", +} + @dataclass class PaddleOCRResult: - """Base adapter for PaddleOCR results. + """Unified adapter for PaddleOCR results (PPStructureV3 and PaddleOCRVL). Converts raw PaddleOCR output to kotaemon Documents. + + Both PPStructureV3 and PaddleOCRVL have similar output structure: + - List of page results + - Each page has parsing_res_list with blocks + - Each block has block_label and block_content """ raw_result: Any file_path: Path extra_info: dict + text_labels: set[str] = field(default_factory=lambda: TEXT_LABELS.copy()) + table_labels: set[str] = field(default_factory=lambda: TABLE_LABELS.copy()) + image_labels: set[str] = field(default_factory=lambda: IMAGE_LABELS.copy()) + formula_labels: set[str] = field(default_factory=lambda: FORMULA_LABELS.copy()) + ignore_labels: set[str] = field(default_factory=lambda: IGNORE_LABELS.copy()) + @property def file_name(self) -> str: return self.file_path.name - @abstractmethod def to_documents(self) -> list[Document]: - """Convert the result to a list of Documents.""" - ... + """Convert PaddleOCR results to Documents.""" + texts: list[Document] = [] + tables: list[Document] = [] + figures: list[Document] = [] + print("WTFFFF", self.raw_result) + + for page_result in self.raw_result: + result_dict = page_result.json["res"] + page_index = result_dict.get("page_index") + print("WTFFFF", page_index) + page_label = (page_index + 1) if page_index is not None else 1 + + page_texts, page_tables, page_figures = self._parse_page( + result_dict, page_label + ) + texts.extend(page_texts) + tables.extend(page_tables) + figures.extend(page_figures) + + return texts + tables + figures + + def _parse_page( + self, + result_dict: dict, + page_label: int, + ) -> tuple[list[Document], list[Document], list[Document]]: + """Parse a single page result.""" + parsing_list = result_dict.get("parsing_res_list", []) + + text_blocks: list[str] = [] + tables: list[Document] = [] + figures: list[Document] = [] + + for block in parsing_list: + label = block.get("block_label", "") + content = block.get("block_content", "") + + if not content or label in self.ignore_labels: + continue + + bbox = block.get("block_bbox") + polygon = block.get("block_polygon_points") + + base_metadata = { + "page_label": page_label, + "file_name": self.file_name, + "file_path": str(self.file_path), + **self.extra_info, + } + if bbox is not None: + base_metadata["bbox"] = bbox + if polygon is not None: + base_metadata["polygon"] = polygon + + if label in self.text_labels: + text_blocks.append(content) + elif label in self.table_labels: + table_content = self._clean_table_html(content) + tables.append( + Document( + text=table_content, + metadata={ + "type": "table", + "table_origin": table_content, + **base_metadata, + }, + ) + ) + + elif label in self.image_labels: + figures.append( + Document( + text=content, + metadata={ + "type": "image", + **base_metadata, + }, + ) + ) + + elif label in self.formula_labels: + text_blocks.append(f"$${content}$$") + else: + text_blocks.append(content) + + text_docs: list[Document] = [] + if text_blocks: + text_docs.append( + Document( + text="\n\n".join(text_blocks), + metadata={ + "page_label": page_label, + "file_name": self.file_name, + "file_path": str(self.file_path), + **self.extra_info, + }, + ) + ) + + return text_docs, tables, figures def _clean_table_html(self, html_content: str) -> str: """Clean HTML table content for better readability.""" diff --git a/libs/kotaemon/kotaemon/loaders/paddleocr_loader/paddleocr_vl_loader.py b/libs/kotaemon/kotaemon/loaders/paddleocr_loader/paddleocr_vl_loader.py index 72049cc9..c44a9243 100644 --- a/libs/kotaemon/kotaemon/loaders/paddleocr_loader/paddleocr_vl_loader.py +++ b/libs/kotaemon/kotaemon/loaders/paddleocr_loader/paddleocr_vl_loader.py @@ -1,8 +1,3 @@ -"""PaddleOCRVL document loader and result adapter.""" - -from __future__ import annotations - -from dataclasses import dataclass, field from pathlib import Path from kotaemon.base import Document, Param @@ -11,149 +6,16 @@ from kotaemon.loaders.base import BaseReader from .adapter import PaddleOCRResult -@dataclass -class PaddleOCRVLResult(PaddleOCRResult): - """Adapter for PaddleOCRVL results. - - PaddleOCRVL uses vision-language models for OCR with: - - Layout detection with polygon points - - Merged layout blocks - - Better handling of complex layouts - """ - - text_labels: set[str] = field( - default_factory=lambda: { - "text", - "paragraph_title", - "doc_title", - "abstract", - "content", - } - ) - table_labels: set[str] = field(default_factory=lambda: {"table"}) - figure_labels: set[str] = field( - default_factory=lambda: {"chart", "figure", "image"} - ) - - def to_documents(self) -> list[Document]: - """Convert PaddleOCRVL results to Documents.""" - texts: list[Document] = [] - tables: list[Document] = [] - figures: list[Document] = [] - - for page_result in self.raw_result: - result_dict = page_result.json - page_index = result_dict.get("page_index") - page_label = (page_index + 1) if page_index is not None else 1 - - page_texts, page_tables, page_figures = self._parse_page( - result_dict, page_label - ) - texts.extend(page_texts) - tables.extend(page_tables) - figures.extend(page_figures) - - return texts + tables + figures - - def _parse_page( - self, - result_dict: dict, - page_label: int, - ) -> tuple[list[Document], list[Document], list[Document]]: - """Parse a single page result from PaddleOCRVL.""" - parsing_list = result_dict.get("parsing_res_list", []) - - text_blocks: list[str] = [] - tables: list[Document] = [] - figures: list[Document] = [] - - for block in parsing_list: - label = block.get("block_label", "") - content = block.get("block_content", "") - - if not content: - continue - - bbox = block.get("block_bbox") - polygon = block.get("block_polygon_points") - - if label in self.text_labels: - text_blocks.append(content) - - elif label in self.table_labels: - table_content = self._clean_table_html(content) - tables.append( - Document( - text=table_content, - metadata={ - "type": "table", - "page_label": page_label, - "table_origin": table_content, - "bbox": bbox, - "polygon": polygon, - "file_name": self.file_name, - "file_path": str(self.file_path), - **self.extra_info, - }, - ) - ) - - elif label in self.figure_labels: - figures.append( - Document( - text=content, - metadata={ - "type": "image", - "page_label": page_label, - "bbox": bbox, - "polygon": polygon, - "file_name": self.file_name, - "file_path": str(self.file_path), - **self.extra_info, - }, - ) - ) - - text_docs: list[Document] = [] - if text_blocks: - text_docs.append( - Document( - text="\n\n".join(text_blocks), - metadata={ - "page_label": page_label, - "file_name": self.file_name, - "file_path": str(self.file_path), - **self.extra_info, - }, - ) - ) - - return text_docs, tables, figures - - class PaddleOCRVLReader(BaseReader): - """Document reader using PaddleOCR Vision-Language model. + """Multilingual document parsing via PaddleOCR-VL-1.5 (0.9B VLM). - PaddleOCRVL uses vision-language models for enhanced OCR with better - understanding of complex document layouts. - - Example: - ```python - from kotaemon.loaders import PaddleOCRVLReader - - # GPU mode (default) - reader = PaddleOCRVLReader() - documents = reader.load_data("path/to/image.png") - - # CPU mode - reader = PaddleOCRVLReader(device="cpu") - ``` - - Args: - device: Device for inference - "gpu:0", "cpu", "npu:0", "xpu:0" + Handles text, tables, formulas, charts, seal recognition, and text spotting. + Robust to skew, warping, scanning, lighting, and screen photography. + Supports cross-page table merging and paragraph heading recognition. + Model: https://huggingface.co/PaddlePaddle/PaddleOCR-VL-1.5 """ - _dependencies = ["paddleocr"] + _dependencies = ["paddleocr[all]"] device: str = Param( "gpu:0", @@ -239,6 +101,7 @@ class PaddleOCRVLReader(BaseReader): "markdown_ignore_labels": self.markdown_ignore_labels, "use_queues": self.use_queues, } + kwargs = {k: v for k, v in kwargs.items() if v is not None} return PaddleOCRVL(**kwargs) def run( @@ -275,10 +138,8 @@ class PaddleOCRVLReader(BaseReader): raw_result = self.pipeline_.predict(str(file_path)) - result = PaddleOCRVLResult( + return PaddleOCRResult( raw_result=raw_result, file_path=file_path, extra_info=extra_info or {}, - ) - - return result.to_documents() + ).to_documents() diff --git a/libs/kotaemon/kotaemon/loaders/paddleocr_loader/ppstructure_v3_loader.py b/libs/kotaemon/kotaemon/loaders/paddleocr_loader/ppstructure_v3_loader.py index f255fffa..71800e54 100644 --- a/libs/kotaemon/kotaemon/loaders/paddleocr_loader/ppstructure_v3_loader.py +++ b/libs/kotaemon/kotaemon/loaders/paddleocr_loader/ppstructure_v3_loader.py @@ -1,8 +1,3 @@ -"""PPStructureV3 document loader and result adapter.""" - -from __future__ import annotations - -from dataclasses import dataclass, field from pathlib import Path from kotaemon.base import Document, Param @@ -11,145 +6,14 @@ from kotaemon.loaders.base import BaseReader from .adapter import PaddleOCRResult -@dataclass -class PPStructureV3Result(PaddleOCRResult): - """Adapter for PPStructureV3 results. - - PPStructureV3 provides structured document parsing with: - - Layout detection - - Table recognition (HTML format) - - Formula recognition - - Chart/figure detection - """ - - text_labels: set[str] = field( - default_factory=lambda: { - "text", - "paragraph_title", - "doc_title", - "abstract", - "content", - } - ) - table_labels: set[str] = field(default_factory=lambda: {"table"}) - figure_labels: set[str] = field( - default_factory=lambda: {"chart", "figure", "image"} - ) - - def to_documents(self) -> list[Document]: - """Convert PPStructureV3 results to Documents.""" - texts: list[Document] = [] - tables: list[Document] = [] - figures: list[Document] = [] - - for page_result in self.raw_result: - result_dict = page_result.json - page_index = result_dict.get("page_index") - page_label = (page_index + 1) if page_index is not None else 1 - - page_texts, page_tables, page_figures = self._parse_page( - result_dict, page_label - ) - texts.extend(page_texts) - tables.extend(page_tables) - figures.extend(page_figures) - - return texts + tables + figures - - def _parse_page( - self, - result_dict: dict, - page_label: int, - ) -> tuple[list[Document], list[Document], list[Document]]: - """Parse a single page result.""" - parsing_list = result_dict.get("parsing_res_list", []) - - text_blocks: list[str] = [] - tables: list[Document] = [] - figures: list[Document] = [] - - for block in parsing_list: - label = block.get("block_label", "") - content = block.get("block_content", "") - - if not content: - continue - - if label in self.text_labels: - text_blocks.append(content) - - elif label in self.table_labels: - table_content = self._clean_table_html(content) - tables.append( - Document( - text=table_content, - metadata={ - "type": "table", - "page_label": page_label, - "table_origin": table_content, - "file_name": self.file_name, - "file_path": str(self.file_path), - **self.extra_info, - }, - ) - ) - - elif label in self.figure_labels: - figures.append( - Document( - text=content, - metadata={ - "type": "image", - "page_label": page_label, - "file_name": self.file_name, - "file_path": str(self.file_path), - **self.extra_info, - }, - ) - ) - - text_docs: list[Document] = [] - if text_blocks: - text_docs.append( - Document( - text="\n\n".join(text_blocks), - metadata={ - "page_label": page_label, - "file_name": self.file_name, - "file_path": str(self.file_path), - **self.extra_info, - }, - ) - ) - - return text_docs, tables, figures - - class PPStructureV3Reader(BaseReader): - """Document reader using PaddleOCR PPStructureV3. + """Document structure extraction via PaddleOCR PPStructureV3. - PPStructureV3 provides comprehensive document structure extraction with - layout detection, table recognition, and formula recognition. - - Example: - ```python - from kotaemon.loaders import PPStructureV3Reader - - # GPU mode (default) - reader = PPStructureV3Reader() - documents = reader.load_data("path/to/document.pdf") - - # CPU mode - reader = PPStructureV3Reader(device="cpu") - ``` - - Args: - device: Device for inference - "gpu:0", "cpu", "npu:0", "xpu:0" - use_doc_orientation_classify: Enable document orientation classification - use_doc_unwarping: Enable document unwarping preprocessing + Layout detection, OCR pipeline, table/chart/formula/seal recognition. + Model: https://huggingface.co/PaddlePaddle/PP-DocLayout-L """ - _dependencies = ["paddleocr"] + _dependencies = ["paddleocr[all]"] device: str = Param( "gpu:0", @@ -157,7 +21,7 @@ class PPStructureV3Reader(BaseReader): ) supported_file_types: list[str] = Param( - [".pdf", ".png", ".jpg", ".jpeg", ".tiff", ".tif", ".bmp"], + [".pdf", ".jpg", ".jpeg", ".png", ".tiff", ".tif"], help="Supported file extensions", ) @@ -267,15 +131,15 @@ class PPStructureV3Reader(BaseReader): "text_det_thresh": self.text_det_thresh, "text_det_box_thresh": self.text_det_box_thresh, "text_det_unclip_ratio": self.text_det_unclip_ratio, - "textline_orientation_model_name": (self.textline_orientation_model_name), - "textline_orientation_model_dir": (self.textline_orientation_model_dir), - "textline_orientation_batch_size": (self.textline_orientation_batch_size), + "textline_orientation_model_name": self.textline_orientation_model_name, + "textline_orientation_model_dir": self.textline_orientation_model_dir, + "textline_orientation_batch_size": self.textline_orientation_batch_size, "text_recognition_model_name": self.text_recognition_model_name, "text_recognition_model_dir": self.text_recognition_model_dir, "text_recognition_batch_size": self.text_recognition_batch_size, "text_rec_score_thresh": self.text_rec_score_thresh, - "table_classification_model_name": (self.table_classification_model_name), - "table_classification_model_dir": (self.table_classification_model_dir), + "table_classification_model_name": self.table_classification_model_name, + "table_classification_model_dir": self.table_classification_model_dir, "wired_table_structure_recognition_model_name": ( self.wired_table_structure_recognition_model_name ), @@ -306,20 +170,20 @@ class PPStructureV3Reader(BaseReader): "table_orientation_classify_model_dir": ( self.table_orientation_classify_model_dir ), - "seal_text_detection_model_name": (self.seal_text_detection_model_name), + "seal_text_detection_model_name": self.seal_text_detection_model_name, "seal_text_detection_model_dir": self.seal_text_detection_model_dir, "seal_det_limit_side_len": self.seal_det_limit_side_len, "seal_det_limit_type": self.seal_det_limit_type, "seal_det_thresh": self.seal_det_thresh, "seal_det_box_thresh": self.seal_det_box_thresh, "seal_det_unclip_ratio": self.seal_det_unclip_ratio, - "seal_text_recognition_model_name": (self.seal_text_recognition_model_name), - "seal_text_recognition_model_dir": (self.seal_text_recognition_model_dir), - "seal_text_recognition_batch_size": (self.seal_text_recognition_batch_size), + "seal_text_recognition_model_name": self.seal_text_recognition_model_name, + "seal_text_recognition_model_dir": self.seal_text_recognition_model_dir, + "seal_text_recognition_batch_size": self.seal_text_recognition_batch_size, "seal_rec_score_thresh": self.seal_rec_score_thresh, - "formula_recognition_model_name": (self.formula_recognition_model_name), - "formula_recognition_model_dir": (self.formula_recognition_model_dir), - "formula_recognition_batch_size": (self.formula_recognition_batch_size), + "formula_recognition_model_name": self.formula_recognition_model_name, + "formula_recognition_model_dir": self.formula_recognition_model_dir, + "formula_recognition_batch_size": self.formula_recognition_batch_size, "use_doc_orientation_classify": self.use_doc_orientation_classify, "use_doc_unwarping": self.use_doc_unwarping, "use_textline_orientation": self.use_textline_orientation, @@ -333,6 +197,7 @@ class PPStructureV3Reader(BaseReader): "lang": self.lang, "ocr_version": self.ocr_version, } + kwargs = {k: v for k, v in kwargs.items() if v is not None} return PPStructureV3(**kwargs) def run( @@ -368,11 +233,11 @@ class PPStructureV3Reader(BaseReader): ) raw_result = self.pipeline_.predict(str(file_path)) + print("WTFFFF", file_path) + print("WTFFFF", raw_result) - result = PPStructureV3Result( + return PaddleOCRResult( raw_result=raw_result, file_path=file_path, extra_info=extra_info or {}, - ) - - return result.to_documents() + ).to_documents() diff --git a/libs/ktem/ktem/index/file/pipelines.py b/libs/ktem/ktem/index/file/pipelines.py index 70212d0d..7d05e4d1 100644 --- a/libs/ktem/ktem/index/file/pipelines.py +++ b/libs/ktem/ktem/index/file/pipelines.py @@ -683,9 +683,27 @@ class IndexDocumentPipeline(BaseFileIndexIndexing): elif self.reader_mode == "docling": readers[".pdf"] = docling_reader elif self.reader_mode == "paddle-struct": - readers[".pdf"] = paddle_struct_reader + readers.update( + { + ".pdf": paddle_struct_reader, + ".png": paddle_struct_reader, + ".jpeg": paddle_struct_reader, + ".jpg": paddle_struct_reader, + ".tiff": paddle_struct_reader, + ".tif": paddle_struct_reader, + } + ) elif self.reader_mode == "paddle-vl": - readers[".pdf"] = paddle_vl_reader + readers.update( + { + ".pdf": paddle_vl_reader, + ".png": paddle_vl_reader, + ".jpeg": paddle_vl_reader, + ".jpg": paddle_vl_reader, + ".tiff": paddle_vl_reader, + ".tif": paddle_vl_reader, + } + ) dev_readers, _, _ = dev_settings() readers.update(dev_readers) @@ -710,7 +728,7 @@ class IndexDocumentPipeline(BaseFileIndexIndexing): "PaddleOCR PPStructureV3 (table+figure extraction)", "paddle-struct", ), - ("PaddleOCR VL (vision-language OCR)", "paddle-vl"), + ("PaddleOCR-VL (VLM document parsing)", "paddle-vl"), ], "component": "dropdown", }, diff --git a/libs/ktem/ktem/pages/chat/__init__.py b/libs/ktem/ktem/pages/chat/__init__.py index 2acb636a..e050e423 100644 --- a/libs/ktem/ktem/pages/chat/__init__.py +++ b/libs/ktem/ktem/pages/chat/__init__.py @@ -924,17 +924,17 @@ class ChatPage(BasePage): urls, chat_input_text = get_urls(chat_input_text) if urls and self.first_indexing_url_fn: print("Detected URLs", urls) - url_file_ids = self.first_indexing_url_fn( + indexed_url_ids = self.first_indexing_url_fn( "\n".join(urls), True, settings, user_id, request=None, ) - file_ids.extend(url_file_ids) + file_ids.extend(indexed_url_ids) # Add new file ids to the first selector choices for display - first_selector_choices.extend(zip(urls, url_file_ids)) + first_selector_choices.extend(zip(urls, indexed_url_ids)) # if file_ids is not empty and chat_input_text is empty # set the input to summary