# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import queue import re import threading import time from itertools import chain from typing import Any, Dict, Optional, Tuple, Union import numpy as np from PIL import Image from ....utils import logging from ....utils.deps import pipeline_requires_extra from ...common.batch_sampler import ImageBatchSampler from ...common.reader import ReadImage from ...utils.benchmark import benchmark from ...utils.hpi import HPIConfig from ...utils.pp_option import PaddlePredictorOption from .._parallel import AutoParallelImageSimpleInferencePipeline from ..base import BasePipeline from ..components import CropByBoxes from ..layout_parsing.merge_table import merge_tables_across_pages from ..layout_parsing.title_level import assign_levels_to_parsing_res from ..layout_parsing.utils import construct_img_path, gather_imgs from .result import BaseResult, PaddleOCRVLBlock, PaddleOCRVLResult from .uilts import ( convert_otsl_to_html, crop_margin, filter_overlap_boxes, merge_blocks, post_process_for_spotting, pre_process_for_spotting, tokenize_figure_of_table, truncate_repetitive_content, untokenize_figure_of_table, ) IMAGE_LABELS = ["image", "header_image", "footer_image"] @benchmark.time_methods class _PaddleOCRVLPipeline(BasePipeline): """_PaddleOCRVLPipeline Pipeline""" def __init__( self, config: Dict, device: Optional[str] = None, pp_option: Optional[PaddlePredictorOption] = None, use_hpip: bool = False, hpi_config: Optional[Union[Dict[str, Any], HPIConfig]] = None, initial_predictor: bool = True, ) -> None: """ Initializes the class with given configurations and options. Args: config (Dict): Configuration dictionary containing various settings. device (str, optional): Device to run the predictions on. Defaults to None. pp_option (PaddlePredictorOption, optional): PaddlePredictor options. Defaults to None. use_hpip (bool, optional): Whether to use the high-performance inference plugin (HPIP) by default. Defaults to False. hpi_config (Optional[Union[Dict[str, Any], HPIConfig]], optional): The default high-performance inference configuration dictionary. Defaults to None. initial_predictor (bool, optional): Whether to initialize predictors. """ super().__init__( device=device, pp_option=pp_option, use_hpip=use_hpip, hpi_config=hpi_config ) if initial_predictor: self.use_doc_preprocessor = config.get("use_doc_preprocessor", True) if self.use_doc_preprocessor: doc_preprocessor_config = config.get("SubPipelines", {}).get( "DocPreprocessor", { "pipeline_config_error": "config error for doc_preprocessor_pipeline!" }, ) self.doc_preprocessor_pipeline = self.create_pipeline( doc_preprocessor_config ) self.use_layout_detection = config.get("use_layout_detection", True) if self.use_layout_detection: layout_det_config = config.get("SubModules", {}).get( "LayoutDetection", {"model_config_error": "config error for layout_det_model!"}, ) model_name = layout_det_config.get("model_name", None) assert model_name is not None and model_name in [ "PP-DocLayoutV2", "PP-DocLayoutV3", ], "model_name must be PP-DocLayoutV2 or PP-DocLayoutV3" layout_kwargs = {} if (threshold := layout_det_config.get("threshold", None)) is not None: layout_kwargs["threshold"] = threshold if ( layout_nms := layout_det_config.get("layout_nms", None) ) is not None: layout_kwargs["layout_nms"] = layout_nms if ( layout_unclip_ratio := layout_det_config.get( "layout_unclip_ratio", None ) ) is not None: layout_kwargs["layout_unclip_ratio"] = layout_unclip_ratio if ( layout_merge_bboxes_mode := layout_det_config.get( "layout_merge_bboxes_mode", None ) ) is not None: layout_kwargs["layout_merge_bboxes_mode"] = layout_merge_bboxes_mode self.layout_det_model = self.create_model( layout_det_config, **layout_kwargs ) self.use_chart_recognition = config.get("use_chart_recognition", False) self.use_seal_recognition = config.get("use_seal_recognition", False) vl_rec_config = config.get("SubModules", {}).get( "VLRecognition", {"model_config_error": "config error for vl_rec_model!"}, ) self.vl_rec_model = self.create_model(vl_rec_config) self.format_block_content = config.get("format_block_content", False) self.use_ocr_for_image_block = config.get("use_ocr_for_image_block", False) self.batch_sampler = ImageBatchSampler( batch_size=config.get("batch_size", 1) ) self.img_reader = ReadImage(format="BGR") self.crop_by_boxes = CropByBoxes() self.use_queues = config.get("use_queues", False) self.merge_layout_blocks = config.get("merge_layout_blocks", True) self.markdown_ignore_labels = config.get( "markdown_ignore_labels", [ "number", "footnote", "header", "header_image", "footer", "footer_image", "aside_text", ], ) def close(self): if hasattr(self, "vl_rec_model"): self.vl_rec_model.close() def get_model_settings( self, use_doc_orientation_classify: Union[bool, None], use_doc_unwarping: Union[bool, None], use_layout_detection: Union[bool, None], use_chart_recognition: Union[bool, None], use_seal_recognition: Union[bool, None], use_ocr_for_image_block: Union[bool, None], format_block_content: Union[bool, None], merge_layout_blocks: Union[bool, None], markdown_ignore_labels: Optional[list[str]] = None, ) -> dict: """ Get the model settings based on the provided parameters or default values. Args: use_doc_orientation_classify (Union[bool, None]): Enables document orientation classification if True. Defaults to system setting if None. use_doc_unwarping (Union[bool, None]): Enables document unwarping if True. Defaults to system setting if None. Returns: dict: A dictionary containing the model settings. """ if use_doc_orientation_classify is None and use_doc_unwarping is None: use_doc_preprocessor = self.use_doc_preprocessor else: if use_doc_orientation_classify is True or use_doc_unwarping is True: use_doc_preprocessor = True else: use_doc_preprocessor = False if use_layout_detection is None: use_layout_detection = self.use_layout_detection if use_chart_recognition is None: use_chart_recognition = self.use_chart_recognition if use_seal_recognition is None: use_seal_recognition = self.use_seal_recognition if use_ocr_for_image_block is None: use_ocr_for_image_block = self.use_ocr_for_image_block if format_block_content is None: format_block_content = self.format_block_content if merge_layout_blocks is None: merge_layout_blocks = self.merge_layout_blocks if markdown_ignore_labels is None: markdown_ignore_labels = self.markdown_ignore_labels return dict( use_doc_preprocessor=use_doc_preprocessor, use_layout_detection=use_layout_detection, use_chart_recognition=use_chart_recognition, use_seal_recognition=use_seal_recognition, use_ocr_for_image_block=use_ocr_for_image_block, format_block_content=format_block_content, merge_layout_blocks=merge_layout_blocks, markdown_ignore_labels=markdown_ignore_labels, ) def check_model_settings_valid(self, input_params: dict) -> bool: """ Check if the input parameters are valid based on the initialized models. Args: input_params (Dict): A dictionary containing input parameters. Returns: bool: True if all required models are initialized according to input parameters, False otherwise. """ if input_params["use_doc_preprocessor"] and not self.use_doc_preprocessor: logging.error( "Set use_doc_preprocessor, but the models for doc preprocessor are not initialized.", ) return False return True def get_layout_parsing_results( self, images, layout_det_results, imgs_in_doc, use_chart_recognition=False, use_seal_recognition=False, use_ocr_for_image_block=False, vlm_kwargs=None, merge_layout_blocks=True, layout_shape_mode="auto", ): blocks = [] has_spotting = False drop_figures_set = set() min_pixels = vlm_kwargs.pop("min_pixels", None) default_min_pixels = min_pixels if min_pixels is not None else 112896 max_pixels = vlm_kwargs.pop("max_pixels", None) default_max_pixels = max_pixels if max_pixels is not None else 1003520 batch_dict_by_pixel = {} id2pixel_key_map = {} image_path_to_obj_map = {} vis_image_labels = IMAGE_LABELS + ["seal"] image_labels = [] if use_ocr_for_image_block else IMAGE_LABELS.copy() if not use_chart_recognition: image_labels += ["chart"] vis_image_labels += ["chart"] if not use_seal_recognition: image_labels += ["seal"] for i, (image, layout_det_res, imgs_in_doc_for_img) in enumerate( zip(images, layout_det_results, imgs_in_doc) ): layout_det_res = filter_overlap_boxes(layout_det_res, layout_shape_mode) boxes = layout_det_res["boxes"] blocks_for_img = self.crop_by_boxes(image, boxes, layout_shape_mode) del layout_det_res, boxes if merge_layout_blocks: blocks_for_img = merge_blocks( blocks_for_img, non_merge_labels=image_labels + ["table"] ) blocks.append(blocks_for_img) for j, block in enumerate(blocks_for_img): block_img = block["img"] block_label = block["label"] if block_label not in image_labels and block_img is not None: figure_token_map = {} text_prompt = "OCR:" min_pixels = vlm_kwargs.pop("ocr_min_pixels", default_min_pixels) max_pixels = vlm_kwargs.pop("ocr_max_pixels", default_max_pixels) drop_figures = [] if block_label == "table": text_prompt = "Table Recognition:" block_img, figure_token_map, drop_figures = ( tokenize_figure_of_table( block_img, block["box"], imgs_in_doc_for_img ) ) min_pixels = vlm_kwargs.pop( "table_min_pixels", default_min_pixels ) max_pixels = vlm_kwargs.pop( "table_max_pixels", default_max_pixels ) elif block_label == "chart" and use_chart_recognition: text_prompt = "Chart Recognition:" min_pixels = vlm_kwargs.pop( "chart_min_pixels", default_min_pixels ) max_pixels = vlm_kwargs.pop( "chart_max_pixels", default_max_pixels ) elif "formula" in block_label and block_label != "formula_number": text_prompt = "Formula Recognition:" crop_img = crop_margin(block_img) w, h, _ = crop_img.shape if w > 2 and h > 2: block_img = crop_img min_pixels = vlm_kwargs.pop( "formula_min_pixels", default_min_pixels ) max_pixels = vlm_kwargs.pop( "formula_max_pixels", default_max_pixels ) elif block_label == "spotting": text_prompt = "Spotting:" has_spotting = True min_pixels = 112896 max_pixels = 1605632 block_img = pre_process_for_spotting(block_img) elif block_label == "seal" and use_seal_recognition: text_prompt = "Seal Recognition:" min_pixels = vlm_kwargs.pop( "seal_min_pixels", default_min_pixels ) max_pixels = vlm_kwargs.pop( "seal_max_pixels", default_max_pixels ) pixel_key = (min_pixels, max_pixels) if pixel_key not in batch_dict_by_pixel: batch_dict_by_pixel[pixel_key] = { "images": [], "queries": [], "figure_token_maps": [], "vlm_block_ids": [], "curr_vlm_block_idx": 0, } batch_dict_by_pixel[pixel_key]["images"].append(block_img) batch_dict_by_pixel[pixel_key]["queries"].append(text_prompt) batch_dict_by_pixel[pixel_key]["figure_token_maps"].append( figure_token_map ) batch_dict_by_pixel[pixel_key]["vlm_block_ids"].append((i, j)) id2pixel_key_map[(i, j)] = pixel_key drop_figures_set.update(drop_figures) del blocks_for_img del images, layout_det_results if vlm_kwargs is None: vlm_kwargs = {} elif vlm_kwargs.get("max_new_tokens", None) is None: vlm_kwargs["max_new_tokens"] = 4096 for pixel_key in batch_dict_by_pixel: min_pixels, max_pixels = pixel_key kwargs = { "use_cache": True, "min_pixels": min_pixels, "max_pixels": max_pixels, **vlm_kwargs, } images = batch_dict_by_pixel[pixel_key]["images"] queries = batch_dict_by_pixel[pixel_key]["queries"] batch_results = list( self.vl_rec_model.predict( [ { "image": image, "query": query, } for image, query in zip(images, queries) ], skip_special_tokens=False if has_spotting else True, **kwargs, ) ) del images, queries batch_dict_by_pixel[pixel_key]["vlm_results"] = batch_results parsing_res_lists = [] table_res_lists = [] spotting_res_list = [] table_blocks = [] for i, blocks_for_img in enumerate(blocks): parsing_res_list = [] table_res_list = [] spotting_res = {} for j, block in enumerate(blocks_for_img): block_img = block["img"] block_bbox = block["box"] block_label = block["label"] block_content = "" figure_token_map = {} if (i, j) in id2pixel_key_map: pixel_key = id2pixel_key_map[(i, j)] pixel_info = batch_dict_by_pixel[pixel_key] curr_vlm_block_idx = pixel_info["curr_vlm_block_idx"] assert curr_vlm_block_idx < len( pixel_info["vlm_block_ids"] ) and pixel_info["vlm_block_ids"][curr_vlm_block_idx] == (i, j) vl_rec_result = pixel_info["vlm_results"][curr_vlm_block_idx] block_img4vl = pixel_info["images"][curr_vlm_block_idx] figure_token_map = pixel_info["figure_token_maps"][ curr_vlm_block_idx ] curr_vlm_block_idx += 1 pixel_info["curr_vlm_block_idx"] = curr_vlm_block_idx vl_rec_result["image"] = block_img4vl result_str = vl_rec_result.get("result", "") if result_str is None: result_str = "" min_count = 5000 if block_label == "table" else 50 result_str = truncate_repetitive_content( result_str, min_count=min_count ) if ("\\(" in result_str and "\\)" in result_str) or ( "\\[" in result_str and "\\]" in result_str ): result_str = result_str.replace("$", "") result_str = ( result_str.replace("\\(", " $ ") .replace("\\)", " $") .replace("\\[\\[", "\\[") .replace("\\]\\]", "\\]") .replace("\\[", " $$ ") .replace("\\]", " $$ ") ) if block_label == "formula_number": result_str = result_str.replace("$", "") if block_label == "table": html_str = convert_otsl_to_html(result_str) if html_str != "": result_str = html_str if block_label == "spotting": h, w = block_img.shape[:2] result_str, spotting_res = post_process_for_spotting( result_str, w, h ) block_content = result_str block_info = PaddleOCRVLBlock( label=block_label, bbox=block_bbox, content=block_content, group_id=block.get("group_id", None), polygon_points=block.get("polygon_points", None), ) if block_label == "table": table_blocks.append( { "figure_token_map": figure_token_map, "block": block_info, } ) if block_label in vis_image_labels and block_img is not None: img_path = construct_img_path(block["label"], block["box"]) image_path_to_obj_map[img_path] = block_info if img_path not in drop_figures_set: import cv2 block_img = cv2.cvtColor(block_img, cv2.COLOR_BGR2RGB) block_info.image = { "path": img_path, "img": Image.fromarray(block_img), } else: continue parsing_res_list.append(block_info) del block_info, block_img # TODO(changdazhou): append table res to table_res_list for blk_info in table_blocks: block = blk_info["block"] figure_token_map = blk_info["figure_token_map"] block.content = untokenize_figure_of_table( block.content, figure_token_map, image_path_to_obj_map ) parsing_res_lists.append(parsing_res_list) table_res_lists.append(table_res_list) spotting_res_list.append(spotting_res) del parsing_res_list, table_res_list, spotting_res return ( parsing_res_lists, table_res_lists, spotting_res_list, imgs_in_doc, ) def predict( self, input: Union[str, list[str], np.ndarray, list[np.ndarray]], use_doc_orientation_classify: Union[bool, None] = False, use_doc_unwarping: Union[bool, None] = False, use_layout_detection: Union[bool, None] = None, use_chart_recognition: Union[bool, None] = None, use_seal_recognition: Union[bool, None] = None, use_ocr_for_image_block: Union[bool, None] = None, layout_threshold: Optional[Union[float, dict]] = None, layout_nms: Optional[bool] = None, layout_unclip_ratio: Optional[Union[float, Tuple[float, float], dict]] = None, layout_merge_bboxes_mode: Optional[str] = None, layout_shape_mode: Optional[str] = "auto", use_queues: Optional[bool] = None, prompt_label: Optional[Union[str, None]] = None, format_block_content: Union[bool, None] = None, repetition_penalty: Optional[float] = None, temperature: Optional[float] = None, top_p: Optional[float] = None, min_pixels: Optional[int] = None, max_pixels: Optional[int] = None, max_new_tokens: Optional[int] = None, merge_layout_blocks: Optional[bool] = None, markdown_ignore_labels: Optional[list[str]] = None, vlm_extra_args: Optional[dict] = None, **kwargs, ) -> PaddleOCRVLResult: """ Predicts the layout parsing result for the given input. Args: input (Union[str, list[str], np.ndarray, list[np.ndarray]]): Input image path, list of image paths, numpy array of an image, or list of numpy arrays. use_doc_orientation_classify (Optional[bool]): Whether to use document orientation classification. use_doc_unwarping (Optional[bool]): Whether to use document unwarping. use_layout_detection (Optional[bool]): Whether to use layout detection. Default is None. use_chart_recognition (Optional[bool]): Whether to use chart recognition. Default is None. use_seal_recognition (Optional[bool]): Whether to use seal recognition. Default is None. layout_threshold (Optional[float]): The threshold value to filter out low-confidence predictions. Default is None. layout_nms (bool, optional): Whether to use layout-aware NMS. Defaults to False. layout_unclip_ratio (Optional[Union[float, Tuple[float, float]]], optional): The ratio of unclipping the bounding box. Defaults to None. If it's a single number, then both width and height are used. If it's a tuple of two numbers, then they are used separately for width and height respectively. If it's None, then no unclipping will be performed. layout_merge_bboxes_mode (Optional[str], optional): The mode for merging bounding boxes. Defaults to None. layout_shape_mode (Optional[str], optional): The mode for layout shape. Defaults to "auto", [ "rect", "quad","poly", "auto"] are supported. use_queues (Optional[bool], optional): Whether to use queues. Defaults to None. prompt_label (Optional[Union[str, None]], optional): The label of the prompt in ['ocr', 'formula', 'table', 'chart']. Defaults to None. format_block_content (Optional[bool]): Whether to format the block content. Default is None. repetition_penalty (Optional[float]): The repetition penalty parameter used for VL model sampling. Default is None. temperature (Optional[float]): Temperature parameter used for VL model sampling. Default is None. top_p (Optional[float]): Top-p parameter used for VL model sampling. Default is None. min_pixels (Optional[int]): The minimum number of pixels allowed when the VL model preprocesses images. Default is None. max_pixels (Optional[int]): The maximum number of pixels allowed when the VL model preprocesses images. Default is None. max_new_tokens (Optional[int]): The maximum number of new tokens. Default is None. merge_layout_blocks (Optional[bool]): Whether to merge layout blocks. Default is None. markdown_ignore_labels (Optional[list[str]]): The list of ignored markdown labels. Default is None. **kwargs (Any): Additional settings to extend functionality. Returns: PaddleOCRVLResult: The predicted layout parsing result. """ model_settings = self.get_model_settings( use_doc_orientation_classify, use_doc_unwarping, use_layout_detection, use_chart_recognition, use_seal_recognition, use_ocr_for_image_block, format_block_content, merge_layout_blocks, markdown_ignore_labels, ) model_settings["return_layout_polygon_points"] = ( False if layout_shape_mode == "rect" else True ) if not self.check_model_settings_valid(model_settings): yield {"error": "the input params for model settings are invalid!"} if use_queues is None: use_queues = self.use_queues if vlm_extra_args is None: vlm_extra_args = {} if not model_settings["use_layout_detection"]: prompt_label = prompt_label if prompt_label else "ocr" if prompt_label.lower() == "chart": model_settings["use_chart_recognition"] = True elif prompt_label.lower() == "seal": model_settings["use_seal_recognition"] = True assert prompt_label.lower() in [ "ocr", "formula", "table", "chart", "spotting", "seal", ], f"Layout detection is disabled (use_layout_detection=False). 'prompt_label' must be one of ['ocr', 'formula', 'table', 'chart'], but got '{prompt_label}'." def _process_cv(batch_data, new_batch_size=None): if not new_batch_size: new_batch_size = len(batch_data) for idx in range(0, len(batch_data), new_batch_size): instances = batch_data.instances[idx : idx + new_batch_size] input_paths = batch_data.input_paths[idx : idx + new_batch_size] page_indexes = batch_data.page_indexes[idx : idx + new_batch_size] page_counts = batch_data.page_counts[idx : idx + new_batch_size] image_arrays = self.img_reader(instances) if model_settings["use_doc_preprocessor"]: doc_preprocessor_results = list( self.doc_preprocessor_pipeline( image_arrays, use_doc_orientation_classify=use_doc_orientation_classify, use_doc_unwarping=use_doc_unwarping, ) ) else: doc_preprocessor_results = [ {"output_img": arr} for arr in image_arrays ] doc_preprocessor_images = [ item["output_img"] for item in doc_preprocessor_results ] if model_settings["use_layout_detection"]: layout_det_results = list( self.layout_det_model( doc_preprocessor_images, threshold=layout_threshold, layout_nms=layout_nms, layout_unclip_ratio=layout_unclip_ratio, layout_merge_bboxes_mode=layout_merge_bboxes_mode, layout_shape_mode=layout_shape_mode, filter_overlap_boxes=False, ) ) imgs_in_doc = [ gather_imgs(doc_pp_img, layout_det_res["boxes"]) for doc_pp_img, layout_det_res in zip( doc_preprocessor_images, layout_det_results ) ] else: layout_det_results = [] for doc_preprocessor_image in doc_preprocessor_images: layout_det_results.append( { "input_path": None, "page_index": None, "boxes": [ { "cls_id": 0, "label": prompt_label.lower(), "score": 1, "coordinate": [ 0, 0, doc_preprocessor_image.shape[1], doc_preprocessor_image.shape[0], ], } ], } ) imgs_in_doc = [[] for _ in layout_det_results] yield input_paths, page_indexes, page_counts, doc_preprocessor_images, doc_preprocessor_results, layout_det_results, imgs_in_doc def _process_vlm(results_cv): ( input_paths, page_indexes, page_counts, doc_preprocessor_images, doc_preprocessor_results, layout_det_results, imgs_in_doc, ) = results_cv ( parsing_res_lists, table_res_lists, spotting_res_list, imgs_in_doc, ) = self.get_layout_parsing_results( images=doc_preprocessor_images, layout_det_results=layout_det_results, imgs_in_doc=imgs_in_doc, use_chart_recognition=model_settings["use_chart_recognition"], use_seal_recognition=model_settings["use_seal_recognition"], use_ocr_for_image_block=model_settings["use_ocr_for_image_block"], vlm_kwargs={ "repetition_penalty": repetition_penalty, "temperature": temperature, "top_p": top_p, "min_pixels": min_pixels, "max_pixels": max_pixels, "max_new_tokens": max_new_tokens, **vlm_extra_args, }, merge_layout_blocks=model_settings["merge_layout_blocks"], layout_shape_mode=layout_shape_mode, ) for ( input_path, page_index, page_count, doc_preprocessor_image, doc_preprocessor_res, layout_det_res, table_res_list, parsing_res_list, spotting_res, imgs_in_doc_for_img, ) in zip( input_paths, page_indexes, page_counts, doc_preprocessor_images, doc_preprocessor_results, layout_det_results, table_res_lists, parsing_res_lists, spotting_res_list, imgs_in_doc, ): single_img_res = { "input_path": input_path, "page_index": page_index, "page_count": page_count, "width": doc_preprocessor_image.shape[1], "height": doc_preprocessor_image.shape[0], "doc_preprocessor_res": doc_preprocessor_res, "layout_det_res": layout_det_res, "table_res_list": table_res_list, "parsing_res_list": parsing_res_list, "spotting_res": spotting_res, "imgs_in_doc": imgs_in_doc_for_img, "model_settings": model_settings, } yield PaddleOCRVLResult(single_img_res) if use_queues: max_num_batches_in_process = 64 queue_input = queue.Queue(maxsize=max_num_batches_in_process) queue_cv = queue.Queue(maxsize=max_num_batches_in_process) queue_vlm = queue.Queue( maxsize=self.batch_sampler.batch_size * max_num_batches_in_process ) event_shutdown = threading.Event() event_data_loading_done = threading.Event() event_cv_processing_done = threading.Event() event_vlm_processing_done = threading.Event() def _worker_input(input_): all_batch_data = self.batch_sampler(input_) while not event_shutdown.is_set(): try: batch_data = next(all_batch_data) except StopIteration: break except Exception as e: queue_input.put((False, "input", e)) break else: queue_input.put((True, batch_data)) del batch_data event_data_loading_done.set() def _worker_cv(): while not event_shutdown.is_set(): try: item = queue_input.get(timeout=0.5) except queue.Empty: if event_data_loading_done.is_set(): event_cv_processing_done.set() break continue if not item[0]: queue_cv.put(item) break try: for results_cv in _process_cv( item[1], ( self.layout_det_model.batch_sampler.batch_size if model_settings["use_layout_detection"] else None ), ): queue_cv.put((True, results_cv)) del results_cv del item except Exception as e: queue_cv.put((False, "cv", e)) break def _worker_vlm(): MAX_QUEUE_DELAY_SECS = 0.5 MAX_NUM_BOXES = self.vl_rec_model.batch_sampler.batch_size while not event_shutdown.is_set(): results_cv_list = [] start_time = time.time() should_break = False num_boxes = 0 while True: remaining_time = MAX_QUEUE_DELAY_SECS - ( time.time() - start_time ) if remaining_time <= 0: break try: item = queue_cv.get(timeout=remaining_time) except queue.Empty: break if not item[0]: queue_vlm.put(item) should_break = True break results_cv_list.append(item[1]) del item for res in results_cv_list[-1][5]: num_boxes += len(res["boxes"]) if num_boxes >= MAX_NUM_BOXES: break if should_break: break if not results_cv_list: if event_cv_processing_done.is_set(): event_vlm_processing_done.set() break continue merged_results_cv = [ list(chain.from_iterable(lists)) for lists in zip(*results_cv_list) ] del results_cv_list try: for result_vlm in _process_vlm(merged_results_cv): queue_vlm.put((True, result_vlm)) del result_vlm del merged_results_cv except Exception as e: queue_vlm.put((False, "vlm", e)) break thread_input = threading.Thread( target=_worker_input, args=(input,), daemon=False ) thread_input.start() thread_cv = threading.Thread(target=_worker_cv, daemon=False) thread_cv.start() thread_vlm = threading.Thread(target=_worker_vlm, daemon=False) thread_vlm.start() try: if use_queues: while not (event_vlm_processing_done.is_set() and queue_vlm.empty()): try: item = queue_vlm.get(timeout=0.5) except queue.Empty: if event_vlm_processing_done.is_set(): break continue if not item[0]: raise RuntimeError( f"Exception from the '{item[1]}' worker: {item[2]}" ) else: yield item[1] else: for batch_data in self.batch_sampler(input): results_cv_list = list(_process_cv(batch_data)) assert len(results_cv_list) == 1, len(results_cv_list) results_cv = results_cv_list[0] for res in _process_vlm(results_cv): yield res del res, results_cv, results_cv_list, batch_data finally: if use_queues: event_shutdown.set() thread_input.join(timeout=5) if thread_input.is_alive(): logging.warning("Input worker did not terminate in time") thread_cv.join(timeout=5) if thread_cv.is_alive(): logging.warning("CV worker did not terminate in time") thread_vlm.join(timeout=5) if thread_vlm.is_alive(): logging.warning("VLM worker did not terminate in time") def concatenate_markdown_pages(self, markdown_list: list) -> tuple: """ Concatenate Markdown content from multiple pages into a single document. Args: markdown_list (list): A list containing Markdown data for each page. Returns: tuple: A tuple containing the processed Markdown text. """ markdown_texts = "" for res in markdown_list: markdown_texts += "\n\n" + res["markdown_texts"] return markdown_texts def concatenate_pages( self, res_list: list, merge_table: bool = True, title_level: bool = True, merge_pages: bool = False, ): """Concatenate layout parsing results from multiple pages. Args: res_list: List of page parsing results merge_talble: Whether to merge tables across pages title_level: Whether to assign title levels merge_pages: Whether to concatenate pages using the new consolidate_pages() logic Returns: PaddleOCRVLResult: Combined OCR-VL result after merge_table or title_level policy """ logging.warning( f"DeprecationWarning: `concatenate_pages()` is deprecated as of v3.3.14 and will be removed in v3.4.0. Please use `restructure_pages()` instead. It provides better support for table merging and title restructuring." ) return self.restructure_pages(res_list, merge_table, title_level, merge_pages) def restructure_pages( self, res_list: list, merge_tables: bool = True, relevel_titles: bool = True, concatenate_pages: bool = False, ): """Restructure layout parsing results from multiple pages. Args: res_list: List of page parsing results merge_tables: Whether to merge tables across pages relevel_titles: Whether to relevel titles concatenate_pages: Whether to concatenate pages to a single document Returns: PaddleOCRVLResult: Combined OCR-VL result after merge_tables or relevel_titles policy """ if len(res_list) == 0: return [] def _get_img_obj(block, model_settings): if block.get("image", None): return block["image"] if block["block_label"] in ("image", "seal") or ( block["block_label"] == "chart" and not model_settings.get("use_chart_recognition", False) ): path = construct_img_path(block["block_label"], block["block_bbox"]) return {"path": path, "img": None} return None def _conver_blocks_to_obj(blocks, model_settings): res = [] for block in blocks: obj = PaddleOCRVLBlock( label=block["block_label"], bbox=block["block_bbox"], polygon_points=block.get("block_polygon_points", None), content=re.sub(r"^#+\s", "", block["block_content"]), group_id=block.get("group_id", None), ) if img := _get_img_obj(block, model_settings): obj.image = img res.append(obj) return res global_block_id = 0 obj_res_list = [] for one_page_res in res_list: if not isinstance(one_page_res, BaseResult): one_page_res = one_page_res["res"] one_page_res["imgs_in_doc"] = [] blocks = one_page_res.get("parsing_res_list", []) model_settings = one_page_res.get("model_settings", {}) blocks = _conver_blocks_to_obj(blocks, model_settings) else: blocks = one_page_res["parsing_res_list"] model_settings = one_page_res.get("model_settings", {}) parsing_res_list = [] for block in blocks: block.global_block_id = global_block_id block.global_group_id = global_block_id global_block_id += 1 parsing_res_list.append(block) one_page_res["parsing_res_list"] = parsing_res_list obj_res_list.append(one_page_res) res_list = obj_res_list blocks_by_page = [res["parsing_res_list"] for res in res_list] if merge_tables: blocks_by_page = merge_tables_across_pages(blocks_by_page) if relevel_titles: blocks_by_page = assign_levels_to_parsing_res(blocks_by_page) concatenate_res = [] if concatenate_pages: all_page_res = res_list[0] all_page_res["parsing_res_list"] = [ blk for blks in blocks_by_page for blk in blks ] all_page_res["page_index"] = None all_page_res["page_count"] = len(res_list) if model_settings["use_layout_detection"]: all_page_res["layout_det_res"] = [ res["layout_det_res"] for res in res_list ] if model_settings["use_doc_preprocessor"]: all_page_res["doc_preprocessor_res"] = [ res["doc_preprocessor_res"] for res in res_list ] concatenate_res.append(PaddleOCRVLResult(all_page_res)) else: for page_idx, one_page_res in enumerate(res_list): one_page_res["parsing_res_list"] = blocks_by_page[page_idx] concatenate_res.append(PaddleOCRVLResult(one_page_res)) yield from concatenate_res class _BasePaddleOCRVLPipeline(AutoParallelImageSimpleInferencePipeline): @property def _pipeline_cls(self): return _PaddleOCRVLPipeline def _get_batch_size(self, config): return config.get("batch_size", 1) @pipeline_requires_extra("ocr") class PaddleOCRVLPipeline(_BasePaddleOCRVLPipeline): entities = "PaddleOCR-VL" @pipeline_requires_extra("ocr") class PaddleOCRVL15Pipeline(_BasePaddleOCRVLPipeline): entities = "PaddleOCR-VL-1.5"