o
    {qi_                     @   s   d dl Z d dlZddlmZmZmZmZ ddlmZm	Z	m
Z
 ddlmZ ddlmZmZ ddlmZ d	d
ddddddddddZg dZG dd deZG dd deZdS )    N   )add_simple_inference_argsget_subcommand_argsperform_simple_inferencestr2bool)DeprecatedOptionAction
deprecatedwarn_deprecated_param)logger   )PaddleXPipelineWrapperPipelineCLISubcommandExecutor)create_config_from_structuretext_detection_model_dirtext_det_limit_side_lentext_det_limit_typetext_det_threshtext_det_box_threshtext_det_unclip_ratiotext_recognition_model_dirtext_recognition_batch_sizeuse_textline_orientationtextline_orientation_model_dirtextline_orientation_batch_sizeZdet_model_dirZdet_limit_side_lenZdet_limit_typeZdet_db_threshZdet_db_box_threshZdet_db_unclip_ratioZrec_model_dirZrec_batch_numZuse_angle_clsZcls_model_dirZcls_batch_num)PP-OCRv3PP-OCRv4PP-OCRv5c                       s   e Zd Z																										d fdd	Zedd Zddddddddddd
ddZddddddddddd
d	d
Zeddd Z	e
dd Zdd Zdd Z  ZS )	PaddleOCRNc           #   	      s  |d ur|t vrtd| dt  dttdd |||
|frL|d us(|d urK| ||\}}|d u s8|d u rGtdt| dt| d|}|}
n|d usT|d ur[tjdd	d
 i d|d|d|d|d|d|d|d|d|	d|
d|d|d|d|d|d|d||||||||d}i }| D ]8\} }!| t	v rt	|  }"t
| |" |"|v sJ t|" d||" d urtd|  d|" d |!||"< q|!|| < q|| _t jd!i | d S )"NzInvalid OCR version: z. Supported values are .c                 S   s   | d u S N )pr!   r!   _/home/app/PaddleOCR-VL/.venv_paddleocr/lib/python3.10/site-packages/paddleocr/_pipelines/ocr.py<lambda>]   s    z$PaddleOCR.__init__.<locals>.<lambda>z)No models are available for the language z and OCR version z^`lang` and `ocr_version` will be ignored when model names or model directories are not `None`.r   )
stacklevel#doc_orientation_classify_model_name"doc_orientation_classify_model_dirdoc_unwarping_model_namedoc_unwarping_model_dirtext_detection_model_namer   textline_orientation_model_namer   r   text_recognition_model_namer   r   use_doc_orientation_classifyuse_doc_unwarpingr   r   r   )r   r   r   text_det_input_shapetext_rec_score_threshreturn_word_boxtext_rec_input_shapez is not a valid parameter name.`z` and `z` are mutually exclusive.r!   )_SUPPORTED_OCR_VERSIONS
ValueErrorallmap_get_ocr_model_namesreprwarningswarnitems_DEPRECATED_PARAM_NAME_MAPPINGr	   _paramssuper__init__)#selfr&   r'   r(   r)   r*   r   r+   r   r   r,   r   r   r-   r.   r   r   r   r   r   r   r/   r0   r1   r2   langZocr_versionkwargsZdet_model_namerec_model_nameparamsZbase_paramsnamevalnew_name	__class__r!   r#   r@   8   s   	




zPaddleOCR.__init__c                 C      dS )NZOCRr!   rA   r!   r!   r#   _paddlex_pipeline_name      z PaddleOCR._paddlex_pipeline_name
r-   r.   r   r   r   r   r   r   r0   r1   c       
         C   s"   | j j|||||||||	|
|dS NrO   )Zpaddlex_pipelinepredictrA   inputr-   r.   r   r   r   r   r   r   r0   r1   r!   r!   r#   predict_iter   s   zPaddleOCR.predict_iterc       
         C   s$   t | j|||||||||	|
|dS rP   )listrT   rR   r!   r!   r#   rQ      s   zPaddleOCR.predictzPlease use `predict` instead.c                 K   s   | j |fi |S r    )rQ   )rA   imgrC   r!   r!   r#   ocr   s   zPaddleOCR.ocrc                 C   s   t  S r    )PaddleOCRCLISubcommandExecutor)clsr!   r!   r#   get_cli_subcommand_executor   s   z%PaddleOCR.get_cli_subcommand_executorc              
   C   s(  i d| j d d| j d d| j d d| j d d	| j d
 d| j d d| j d d| j d d| j d d| j d d| j d d| j d d| j d d| j d d| j d p_| j d d| j d d| j d  | j d! | j d" | j d# | j d$ | j d% | j d& | j d' | j d( d)}t|S )*NzISubPipelines.DocPreprocessor.SubModules.DocOrientationClassify.model_namer&   zHSubPipelines.DocPreprocessor.SubModules.DocOrientationClassify.model_dirr'   z?SubPipelines.DocPreprocessor.SubModules.DocUnwarping.model_namer(   z>SubPipelines.DocPreprocessor.SubModules.DocUnwarping.model_dirr)   z#SubModules.TextDetection.model_namer*   z"SubModules.TextDetection.model_dirr   z)SubModules.TextLineOrientation.model_namer+   z(SubModules.TextLineOrientation.model_dirr   z)SubModules.TextLineOrientation.batch_sizer   z%SubModules.TextRecognition.model_namer,   z$SubModules.TextRecognition.model_dirr   z%SubModules.TextRecognition.batch_sizer   z9SubPipelines.DocPreprocessor.use_doc_orientation_classifyr-   z.SubPipelines.DocPreprocessor.use_doc_unwarpingr.   Zuse_doc_preprocessorr   z'SubModules.TextDetection.limit_side_lenr   r   r   r   r   r/   r0   r1   r2   )z#SubModules.TextDetection.limit_typezSubModules.TextDetection.threshz#SubModules.TextDetection.box_threshz%SubModules.TextDetection.unclip_ratioz$SubModules.TextDetection.input_shapez'SubModules.TextRecognition.score_threshz*SubModules.TextRecognition.return_word_boxz&SubModules.TextRecognition.input_shape)r>   r   )rA   Z	STRUCTUREr!   r!   r#   _get_paddlex_config_overrides   s   
"%(+
-.1Dz'PaddleOCR._get_paddlex_config_overridesc                 C   s  g d}g d}g d}g d}g d}g d}|d u rd}|d u r>|g d| | | | | v r5d	}n	||v r<d
}ndS |d	krd\}	}
|dv rMd}
nL|dkrTd}
nE||v r[d}	n>||v rbd}	n7||v rid}	n0||v rpd}	n)||v rwd}	n"|dkr~d}	n|dkrd}	n|dkrd}	n|dkrd}	n|dkrd}	|	d ur|	 d}
d|
fS |dkr|dkrdS |dkrdS dS d }	||v rd}	n||v rd}	n||v rd}	n||v rd}	n||v r|}	d }
|	dkrd}
d!|
fS |	d ur|	 d }
d!|
fS )"N)1afazbscscydadeesetfrgahrhuidisitkuZlaltlvmimsmtnlnoocpiplptroZrs_latinskslsqsvswtltruzvifrenchgermanfieuglZlbZrmcaZqu)arfaZugurZpsrl   sdZbal)rubeuk)!r   Zrs_cyrillicr   bgr   ZmnZabqZadyZkbdZavaZdarZinhZcheZlbeZleztabkkkytgmkttZcvbaZmhrmoZudmkvosZbuaZxalZtyvZsahZkaa)himrneZbhmaiangZbhoZmahZscknewZgomsaZbgc)chenkoreanjapanchinese_chttekatar   )	r   r   r   r   r   thelr   r   r   r   )NN)r   r   r   zPP-OCRv5_server_recr   zen_PP-OCRv5_mobile_reclatinZeslavarabiccyrillicZ
devanagarir   r   r   r   r   z_PP-OCRv5_mobile_reczPP-OCRv5_server_detr   )PP-OCRv4_mobile_detzPP-OCRv4_mobile_rec)r   zen_PP-OCRv4_mobile_reczPP-OCRv3_mobile_recz_PP-OCRv3_mobile_reczPP-OCRv3_mobile_detr!   )rA   rB   Zppocr_versionZLATIN_LANGSZARABIC_LANGSZESLAV_LANGSZCYRILLIC_LANGSZDEVANAGARI_LANGSZSPECIFIC_LANGSZrec_langrD   r!   r!   r#   r8   4  s   3#

zPaddleOCR._get_ocr_model_names)NNNNNNNNNNNNNNNNNNNNNNNNNN)__name__
__module____qualname__r@   propertyrM   rT   rQ   r   rW   classmethodrZ   r[   r8   __classcell__r!   r!   rI   r#   r   7   st    m
!

Gr   c                   @   s(   e Zd Zedd Zdd Zdd ZdS )rX   c                 C   rK   )NrW   r!   rL   r!   r!   r#   subparser_name  rN   z-PaddleOCRCLISubcommandExecutor.subparser_namec                 C   s  t | |jdtdd |jdtdd |jdtdd |jdtd	d |jd
tdd |jdtdd |jdtdd |jdtdd |jdtdd |jdtdd |jdtdd |jdtdd |jdtdd |jdtdd |jdtdd |jd td!d |jd"td#d |jd$td%d |jd&td'd |jd(td)d |jd*d+td,d-d. |jd/td0d |jd1td2d |jd3d+td,d4d. |jd5td6d |jd7ttd8d9 tttttttttttd:}t D ]\}}||v sJ ||jd;| t	td<| d=d> qd S )?Nz%--doc_orientation_classify_model_namez<Name of the document image orientation classification model.)typehelpz$--doc_orientation_classify_model_dirzFPath to the document image orientation classification model directory.z--doc_unwarping_model_namez'Name of the text image unwarping model.z--doc_unwarping_model_dirz,Path to the image unwarping model directory.z--text_detection_model_namez!Name of the text detection model.z--text_detection_model_dirz+Path to the text detection model directory.z!--textline_orientation_model_namez7Name of the text line orientation classification model.z --textline_orientation_model_dirzAPath to the text line orientation classification model directory.z!--textline_orientation_batch_sizez>Batch size for the text line orientation classification model.z--text_recognition_model_namez#Name of the text recognition model.z--text_recognition_model_dirz-Path to the text recognition model directory.z--text_recognition_batch_sizez*Batch size for the text recognition model.z--use_doc_orientation_classifyz9Whether to use document image orientation classification.z--use_doc_unwarpingz$Whether to use text image unwarping.z--use_textline_orientationz4Whether to use text line orientation classification.z--text_det_limit_side_lenzUThis sets a limit on the side length of the input image for the text detection model.z--text_det_limit_typezxThis determines how the side length limit is applied to the input image before feeding it into the text deteciton model.z--text_det_threshzDetection pixel threshold for the text detection model. Pixels with scores greater than this threshold in the output probability map are considered text pixels.z--text_det_box_threshzDetection box threshold for the text detection model. A detection result is considered a text region if the average score of all pixels within the border of the result is greater than this threshold.z--text_det_unclip_ratiozText detection expansion coefficient, which expands the text region using this method. The larger the value, the larger the expansion area.z--text_det_input_shape   )CHWz(Input shape of the text detection model.)nargsr   metavarr   z--text_rec_score_threshz^Text recognition threshold. Text results with scores greater than this threshold are retained.z--return_word_boxz<Whether to return the coordinates of the recognition result.z--text_rec_input_shapez*Input shape of the text recognition model.z--langz/Language in the input image for OCR processing.z--ocr_versionzPP-OCR version to use.)r   choicesr   r   z--z[Deprecated] Please use `--z
` instead.)actionr   r   )
r   add_argumentstrintr   floatr4   r=   r<   r   )rA   Z	subparserZdeprecated_arg_typesrF   rH   r!   r!   r#   _update_subparser  s6  
z0PaddleOCRCLISubcommandExecutor._update_subparserc                 C   s~   t |}t D ]/\}}||v sJ || }|| }|d ur.|d ur.td|| td |d u r7|| qtt	| d S )Nz)`--%s` and `--%s` are mutually exclusive.r   )
r   r=   r<   r
   errorsysexitpopr   r   )rA   argsrE   rF   rH   rG   new_valr!   r!   r#   execute_with_args  s   

z0PaddleOCRCLISubcommandExecutor.execute_with_argsN)r   r   r   r   r   r   r   r!   r!   r!   r#   rX     s    
 !rX   )r   r:   Z
_utils.clir   r   r   r   Z_utils.deprecationr   r   r	   Z_utils.loggingr
   baser   r   utilsr   r=   r4   r   rX   r!   r!   r!   r#   <module>   s2      N