_UniRec Pipeline
| 60 | |
| 61 | |
| 62 | class OpenDoc(BasePipeline): |
| 63 | """_UniRec Pipeline""" |
| 64 | |
| 65 | def __init__( |
| 66 | self, |
| 67 | gpuId: Optional[int] = 0, |
| 68 | pp_option: Optional[PaddlePredictorOption] = None, |
| 69 | use_hpip: bool = False, |
| 70 | hpi_config: Optional[Union[Dict[str, Any], HPIConfig]] = None, |
| 71 | ) -> None: |
| 72 | """ |
| 73 | Initializes the class with given configurations and options. |
| 74 | |
| 75 | Args: |
| 76 | config (Dict): Configuration dictionary containing various settings. |
| 77 | gpuId (int, optional): GPU ID to run the predictions on. Defaults to 0. |
| 78 | pp_option (PaddlePredictorOption, optional): PaddlePredictor options. Defaults to None. |
| 79 | use_hpip (bool, optional): Whether to use the high-performance |
| 80 | inference plugin (HPIP) by default. Defaults to False. |
| 81 | hpi_config (Optional[Union[Dict[str, Any], HPIConfig]], optional): |
| 82 | The default high-performance inference configuration dictionary. |
| 83 | Defaults to None. |
| 84 | """ |
| 85 | if gpuId < 0: |
| 86 | device = None |
| 87 | else: |
| 88 | device = f'gpu:{gpuId}' |
| 89 | super().__init__(device=device, |
| 90 | pp_option=pp_option, |
| 91 | use_hpip=use_hpip, |
| 92 | hpi_config=hpi_config) |
| 93 | config = load_pipeline_config( |
| 94 | str(root_dir / '../configs/rec/unirec/opendoc_pipeline.yml')) |
| 95 | |
| 96 | self.use_doc_preprocessor = config.get('use_doc_preprocessor', True) |
| 97 | if self.use_doc_preprocessor: |
| 98 | doc_preprocessor_config = config.get('SubPipelines', {}).get( |
| 99 | 'DocPreprocessor', |
| 100 | { |
| 101 | 'pipeline_config_error': |
| 102 | 'config error for doc_preprocessor_pipeline!' |
| 103 | }, |
| 104 | ) |
| 105 | self.doc_preprocessor_pipeline = self.create_pipeline( |
| 106 | doc_preprocessor_config) |
| 107 | |
| 108 | self.use_layout_detection = config.get('use_layout_detection', True) |
| 109 | if self.use_layout_detection: |
| 110 | layout_det_config = config.get('SubModules', {}).get( |
| 111 | 'LayoutDetection', |
| 112 | {'model_config_error': 'config error for layout_det_model!'}, |
| 113 | ) |
| 114 | model_name = layout_det_config.get('model_name', None) |
| 115 | assert (model_name is not None and model_name |
| 116 | == 'PP-DocLayoutV2'), 'model_name must be PP-DocLayoutV2' |
| 117 | layout_kwargs = {} |
| 118 | if (threshold := layout_det_config.get('threshold', |
| 119 | None)) is not None: |