o 'jO' @sddlZddlZddlZejejeZejeej dejej eddej d<ddl Z ddl Z ddlZddlZddlZddlZddlmZddlmmZddlmmZddlmmZddlmmZddlmZm Z ddl!m"Z"ddlm#Z#m$Z$m%Z%e"Z&Gd d d e'Z(d d Z)d dZ*e+dkre,Z-e-j.rgZ/e-j0Z0e1e0D]#Z2ej3dgej4d5e2d5dgZ6ej7e6ej8ej8dZ9e/e9qe/D]Z9e9:qdSe*e-dSdS)Nz../..Z auto_growthZFLAGS_allocator_strategy)Image)get_image_file_listcheck_and_read) get_logger)draw_ocr_box_txtget_rotate_crop_imageget_minarea_rect_cropc@s&eZdZddZddZd ddZdS) TextSystemcCs\|js ttjt||_t ||_ |j |_ |j |_ |j r&t ||_||_d|_dS)Nr)Zshow_logloggersetLevelloggingINFO predict_detZ TextDetector text_detector predict_recZTextRecognizertext_recognizer use_angle_cls drop_score predict_clsZTextClassifiertext_classifierargscrop_image_res_index)selfrre/var/www/html/Deteccion_Ine/venv/lib/python3.10/site-packages/paddleocr/tools/infer/predict_system.py__init__*s     zTextSystem.__init__c Csvtj|ddt|}t|D]"}ttj|d||jd||t |d||q|j|7_dS)NTexist_okZmg_crop_z.jpgz, ) osmakedirslenrangecv2imwritepathjoinrr debug)r output_dir img_crop_listrec_resZbbox_numbnorrrdraw_crop_rec_res8s zTextSystem.draw_crop_rec_resTcCsddddd}|durtddd|fSt}|}||\}}||d<|durCtd|t}|||d<dd|fStdt||g} t|}tt|D]} t || } |j j dkrot || } nt || } | | qZ|jr|r|| \} } }||d <td t| ||| \}}||d <td t|||j jr||j j| |gg}}t||D]\}}|\}}||jkr||||qt}|||d<|||fS) Nr)detrecclsallzno valid image providedr,zno dt_boxes found, elapsed : {}r/zdt_boxes num : {}, elapsed : {}Zquadr.zcls num : {}, elapsed : {}r-zrec_res num : {}, elapsed : {})r r&timecopyrformatr sorted_boxesr!deepcopyrZ det_box_typerrappendrrrZ save_crop_resr+Zcrop_res_save_dirzipr)rimgr. time_dictstartZori_imdt_boxeselapseendr(r*Ztmp_boxZimg_cropZ angle_listr)Z filter_boxesZfilter_rec_resboxZ rec_resulttextscorerrr__call__Csf                 zTextSystem.__call__N)T)__name__ __module__ __qualname__rr+r@rrrrr )s r cCs|jd}t|ddd}t|}t|dD]F}t|ddD]=}t||ddd||dddkr\||ddd||ddkr\||}||d||<|||d<qq|S)z Sort text boxes in order from top to bottom, left to right args: dt_boxes(array):detected text boxes with shape [4, 2] return: sorted boxes(array) with shape [4, 2] rcSs|dd|ddfS)Nrr)xrrrzsorted_boxes..)keyrD )shapesortedlistr!abs)r:Z num_boxesr3Z_boxesijtmprrrr3{s ,$r3c% st|j}||jd|j}t|}d}|j}|j}|j}tj |ddg}t d|j rGt jddgdt j}tdD]} ||} q@d} d\} } }t}d}t|D]T\}}t|\}}}|sm|smt|}|s|dur|t d |qX|g}n|j}|t|ks|dkrt|}|d|}t|D]\}}t}||\}t|}| |7} t|d krt t|d t|d ||fn t t|d ||fD] \}}t d ||qއfddttD} t|d krtj|d t|dtj | ddd}ntj|dtj | ddd}|!||rt"#t$|tj%}}fddttD} fddttD}!t&||| |!||d}"|rn|ddd}#n|r~|'dd t|d}#n|}#t(tj)|tj|#|"dddddddft dtj)|tj|#qqXt dt||j*r|j+j,-|j.j,-t/tj)|ddd d!}$|$0|WddS1swYdS)"NTrzIn PP-OCRv3, rec_image_shape parameter defaults to '3, 48, 320', if you are using recognition model with PP-OCRv2 or an older version, please set --rec_image_shape='3,32,320r)rSrJ)rrrzerror in loading image:{}rD_z Predict time of %s: %.3fsz {}, {:.3f}cs4g|]}|dt|tjdqS)r)Z transcriptionZpoints)nparrayastypeZint32tolist.0rOr:r)rr s   zmain.. F) ensure_ascii cg|]}|dqS)rrrZr)rrr]rGcra)rDrrZrbrrr]rG)r font_pathZpngz.pdfz.pngrIz The visualized image saved in {}zThe predict total time is {}zsystem_results.txtwzutf-8)encoding)1rZ image_dir process_idtotal_process_numr Z vis_font_pathrdraw_img_save_dirrrr infoZwarmuprVrandomuniformrXZuint8r!r0 enumeraterr"Zimreadr&r2page_numr strr$basenamejsondumpsr5rZ fromarrayZcvtColorZ COLOR_BGR2RGBrreplacer#r%Z benchmarkrZautologreportropen writelines)%rZimage_file_listZtext_sysZ is_visualizercrriZ save_resultsr7rOres total_timeZcpu_memZgpu_memZgpu_utilZ_stcountidxZ image_fileZflag_gifZflag_pdfZimgsrnindexZ starttimer8r;r>r?Z save_predimageZboxesZtxtsZscoresZdraw_imgZ save_filefrr\rmains              9    $r~__main__z-uz--process_id={}z --use_mp={}F)stdoutstderr);rsys subprocessr$dirnameabspath__file____dir__r5insertr%environr"r1numpyrVrqr0r ZPILrZtools.infer.utilityZinferZutilityZtools.infer.predict_recrZtools.infer.predict_detrZtools.infer.predict_clsrZppocr.utils.utilityrrZppocr.utils.loggingrrrrr objectr r3r~rA parse_argsrZuse_mpZp_listrhr!rg executableargvr2cmdPopenrpwaitrrrrsT      Rn