o 0jF<@sddlmZmZmZmZmZmZddlZddl m Z ddl m Z ddl mZddlmZdd lmZmZdd lmZdd lmZd d lmZd dlmZd dlmZddlmZej GdddeZ!e dGdddeZ"dS))AnyDictListOptionalTupleUnionN)logging)pipeline_requires_extra)ImageBatchSampler) ReadImage) HPIConfigPaddlePredictorOption) DetResult) benchmark)(AutoParallelImageSimpleInferencePipeline) BasePipeline) CropByBoxes)FormulaRecognitionResultcsfeZdZdZddddddddedeedeedeeeefd eed e d ee eeefe fd dffd dZ dee dee dee d e fddZdede eeefd e fddZ        dde eeeejeejfdee dee dee dee eeefdee ee fdee dee eeeeffdeed efddZZS) _FormulaRecognitionPipelinezFormula Recognition PipelineNFdeviceengine engine_config pp_optionuse_hpip hpi_configconfigrrrrrrreturnc s\tjd||||||d||dd|_|jr,|didddi} || |_|dd|_|jr|d id d d i} i} | d d} durR| | d <| dd} dur`| | d<| dd}durn|| d<| dd}dur||| d<|j| fi| |_|d idd di}|||_ t |_ t |ddd|_ tdd|_dS)aInitializes the formula recognition pipeline. Args: config (Dict): Configuration dictionary containing various settings. device (Optional[str], optional): The device to use for prediction. Defaults to `None`. engine (Optional[str], optional): Inference engine. Defaults to `None`. engine_config (Optional[Dict[str, Any]], optional): Engine-specific config. Defaults to `None`. pp_option (Optional[PaddlePredictorOption], optional): Paddle predictor options. Defaults to `None`. use_hpip (bool, optional): Whether to use HPIP. Defaults to `False`. hpi_config (Optional[Union[Dict[str, Any], HPIConfig]], optional): HPIP configuration. Defaults to `None`. ruse_doc_preprocessorTZ SubPipelinesZDocPreprocessorZpipeline_config_errorz+config error for doc_preprocessor_pipeline!use_layout_detectionZ SubModulesZLayoutDetectionZmodel_config_errorz"config error for layout_det_model! thresholdN layout_nmslayout_unclip_ratiolayout_merge_bboxes_modeZFormulaRecognitionz#config error for formula_rec_model! batch_sizer)r(ZBGR)format)super__init__getr"Zcreate_pipelinedoc_preprocessor_pipeliner#Z create_modellayout_det_modelformula_recognition_modelr_crop_by_boxesr batch_samplerr img_reader)selfr rrrrrrkwargsZdoc_preprocessor_configZlayout_det_configZ layout_kwargsr$r%r&r'Zformula_recognition_config __class__r*y/var/www/html/Deteccion_Ine/venv/lib/python3.10/site-packages/paddlex/inference/pipelines/formula_recognition/pipeline.pyr,$sp     z$_FormulaRecognitionPipeline.__init__use_doc_orientation_classifyuse_doc_unwarpingr#cCsL|dur |dur |j}n |dus|durd}nd}|dur |j}t||dS)a Get the model settings based on the provided parameters or default values. Args: 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. Returns: dict: A dictionary containing the model settings. NTF)r"r#)r"r#dict)r4r9r:r#r"r*r*r8get_model_settingszsz._FormulaRecognitionPipeline.get_model_settingsmodel_settingslayout_det_rescCsR|dr|jstddS|dr'|durtddS|js'tddSdS) a Check if the input parameters are valid based on the initialized models. Args: model_settings (Dict): A dictionary containing input parameters. layout_det_res (Union[DetResult, List[DetResult]]): The layout detection result(s). Returns: bool: True if all required models are initialized according to input parameters, False otherwise. r"zRSet use_doc_preprocessor, but the models for doc preprocessor are not initialized.Fr#Nz^The layout detection model has already been initialized, please set use_layout_detection=FalsezRSet use_layout_detection, but the models for layout detection are not initialized.T)r"r errorr#)r4r=r>r*r*r8check_model_settings_valids" z6_FormulaRecognitionPipeline.check_model_settings_validinputlayout_thresholdr%r&r'c ( ks||||} || |sddiV|} | dur%t| ts!| g} t| } t||D]@\} }||j}| drFt|j |||d}ndd|D}dd|D}g}| d s| durd d|D}t| |}|D]}g}d |d <| || |qnn| d rt|j ||||| d }ng}|D]} zt | }Wn tytdw| |qg}g}dg}t||D]C\}}d }|dD]1}|ddvr|||g}|d}| |di}||d <|d|d<| ||d 7}q| t|qt| |}tt|d D];}|||||d } |||||d }!t| |!D]\}"}|"d }|"d}#||d <|#|d<q/| |!qt|j|j|||D]\}$}%}}&}|$|%||&|| d}'t|'VqVq,dS)aH This function predicts the layout parsing result for the given input. Args: input (Union[str, list[str], np.ndarray, list[np.ndarray]]): The input image(s) of pdf(s) to be processed. use_layout_detection (Optional[bool]): Whether to use layout detection. use_doc_orientation_classify (Optional[bool]): Whether to use document orientation classification. use_doc_unwarping (Optional[bool]): Whether to use document unwarping. layout_det_res (Optional[Union[DetResult, List[DetResult]]]): The layout detection result(s). It will be used if it is not None and use_layout_detection is False. **kwargs: Additional keyword arguments. Returns: formulaRecognitionResult: The predicted formula recognition result. r?z0the input params for model settings are invalid!Nr")r9r:cSsg|]}d|iqSZ output_imgr*).0Zarrr*r*r8 sz7_FormulaRecognitionPipeline.predict..cSsg|]}|dqSrCr*)rDitemr*r*r8rEsr#cSsg|]}iqSr*r*)rD_r*r*r8rEsrformula_region_id)r$r%r&r'zNo more layout det resultsrZboxeslabel)ZformulaZimgZ coordinatedt_polys) input_path page_indexr>doc_preprocessor_resZformula_res_listr=)r<r@ isinstancelistiter enumerater2r3Z instancesr.r0appendr/next StopIteration ValueErrorziplowerr1lenrangeZ input_pathsZ page_indexesr)(r4rAr#r9r:r>rBr%r&r'r5r=Zexternal_layout_det_resultsrGZ batch_dataZ image_arraysZdoc_preprocessor_resultsZdoc_preprocessor_imagesZformula_resultsZlayout_det_resultsZformula_rec_resultsZformula_rec_resZformula_results_for_imgZformula_crop_imgsZformula_det_resultsZ chunk_indicesZdoc_preprocessor_imagerHZbox_infoZ crop_img_inforesidxZformula_det_results_for_idxZformula_rec_results_for_idxZformula_det_resrJrKrLrMZsingle_img_resr*r*r8predicts                 z#_FormulaRecognitionPipeline.predict)NNNNNNNN)__name__ __module__ __qualname____doc__rrstrrrboolrrr,r;r<rrr@npZndarrayfloatrrr\ __classcell__r*r*r6r8r s  V ! %   rZocrc@s&eZdZdgZeddZddZdS)FormulaRecognitionPipelineZformula_recognitioncCstS)N)r)r4r*r*r8 _pipeline_clsasz(FormulaRecognitionPipeline._pipeline_clscCs |ddS)Nr(r)r-)r4r r*r*r8_get_batch_sizees z*FormulaRecognitionPipeline._get_batch_sizeN)r]r^r_entitiespropertyrgrhr*r*r*r8rf]s   rf)#typingrrrrrrnumpyrcutilsr Z utils.depsr Zcommon.batch_samplerr Z common.readerr modelsrrZmodels.object_detection.resultrZutils.benchmarkrZ _parallelrbaser componentsrresultrZ time_methodsrrfr*r*r*r8s&           >