o
    'j                      @   sX  d dl mZ d dlmZmZ d dlmZ d dlmZm	Z	 d dl
mZ ddlmZ d dlmZ G d	d
 d
eZG dd deZG dd deZG dd deZG dd deZedkrd dlZd dlZedZeeZW d   n1 svw   Y  edZeeZW d   n1 sw   Y  eddZeeeZe  Z!e!e dS dS )    )Levenshtein)APTEDConfig)Tree)etreehtml)deque   )parallel_process)tqdmc                   @   s   e Zd ZdddZdd ZdS )	TableTreeNc                 G   s&   || _ || _|| _|| _t|| _d S N)tagcolspanrowspancontentlistchildren)selfr   r   r   r   r    r   v/var/www/html/Deteccion_Ine/venv/lib/python3.10/site-packages/paddleocr/ppstructure/table/table_metric/table_metric.py__init__   s
   zTableTree.__init__c                 C   sP   | j dkrd| j | j| j| jf }nd| j  }| jD ]}|| 7 }qd|S )z!Show tree using brackets notationtdz3"tag": %s, "colspan": %d, "rowspan": %d, "text": %sz	"tag": %sz{{{}}})r   r   r   r   r   bracketformat)r   resultchildr   r   r   r      s   



zTableTree.bracket)NNN)__name__
__module____qualname__r   r   r   r   r   r   r      s    
r   c                   @      e Zd Zdd ZdS )CustomConfigc                 C   sR   |j |j ks|j|jks|j|jkrdS |j dkr'|js|jr't|j|jS dS )Compares attributes of trees      ?r           )r   r   r   r   r   normalized_distance)r   node1node2r   r   r   rename*   s   $
zCustomConfig.renameNr   r   r   r(   r   r   r   r   r!   )       r!   c                   @   r    )CustomConfig_del_shortc                 C   s~   |j |j ks|j|jks|j|jkrdS |j dkr=|js|jr=|j}|j}t|dk r.dg}t|dk r7dg}t||S dS )r"   r#   r      z####r$   )r   r   r   r   lenr   r%   r   r&   r'   Znode1_contentZnode2_contentr   r   r   r(   8   s   $
zCustomConfig_del_short.renameNr)   r   r   r   r   r+   7   r*   r+   c                   @   r    )CustomConfig_del_blockc                 C   s   |j |j ks|j|jks|j|jkrdS |j dkrY|js|jrY|j}|j}d|v r<t|d ||d d|v s)d|v rSt|d ||d d|v s@t||S dS )r"   r#   r    r$   )	r   r   r   r   printindexpopr   r%   r.   r   r   r   r(   K   s    $
zCustomConfig_del_block.renameNr)   r   r   r   r   r/   J   r*   r/   c                   @   sD   e Zd ZdZdddZdd Zdd	d
Zdd Zdd Zdd Z	dS )TEDSz* Tree Edit Distance basead Similarity
    Fr	   Nc                 C   s6   t |tr	|dksJ d|| _|| _|| _g | _d S )Nr	   z)n_jobs must be an integer greather than 1)
isinstanceintstructure_onlyn_jobsignore_nodes
__tokens__)r   r7   r8   r9   r   r   r   r   a   s   

zTEDS.__init__c                 C   s   | j d|j  |jdur|  j t|j7  _ | D ]}| | q|jdkr2| j d|j  |jdkrH|jdurJ|  j t|j7  _ dS dS dS )z Tokenizes table cells
        z<%s>NZunkz</%s>r   )r:   appendr   textr   getchildrentokenizetail)r   nodenr   r   r   r>   i   s   

zTEDS.tokenizec                 C   s   |j dkr8| jrg }ng | _| | | jdd  }t|j t|jddt|jdd|gt	 R  }nt|j dddgt	 R  }|durO|j
| |j dkra| D ]}| || qX|du rg|S dS )z< Converts HTML tree to the format required by apted
        r   r	   r   1r   N)r   r7   r:   r>   copyr   r6   Zattribgetr   r   r;   r=   load_html_tree)r   r@   parentcellnew_noderA   r   r   r   rF   v   s,   



zTEDS.load_html_treec           
      C   s   |r|sdS t jddd}t j||d}t j||d}|drx|drx|dd }|dd }| jrJtj|g| jR   tj|g| jR   t|d}t|d}t||}| 	|}| 	|}t
||t  }	d	t|	|  S dS )
zg Computes TEDS score between the prediction and the ground truth of a
            given sample
        r$   Tzutf-8)Zremove_commentsencoding)parserz
body/tabler   z.//*r#   )r   
HTMLParserZ
fromstringZxpathr9   r   Z
strip_tagsr-   maxrF   r   r!   Zcompute_edit_distancefloat)
r   predtruerK   Zn_nodes_predZn_nodes_trueZn_nodesZ	tree_predZ	tree_trueZdistancer   r   r   evaluate   s.   


zTEDS.evaluatec                    sj     }jdkr fddt|D }n fdd|D }t|jdjdd}tt||}|S )a&   Computes TEDS score between the prediction and the ground truth of
            a batch of samples
            @params pred_json: {'FILENAME': 'HTML CODE', ...}
            @params true_json: {'FILENAME': {'html': 'HTML CODE'}, ...}
            @output: {'FILENAME': 'TEDS SCORE', ...}
        r	   c                    s(   g | ]}  |d | d qS ) r   )rQ   rE   .0filename	pred_jsonr   	true_jsonr   r   
<listcomp>   s    
z'TEDS.batch_evaluate.<locals>.<listcomp>c                    s&   g | ]}  |d | d dqS )rR   r   rO   rP   )rE   rS   )rW   rX   r   r   rY      s    

TZ
use_kwargsr8   Z	front_num)keysr8   r   r
   rQ   dictzip)r   rW   rX   Zsamplesscoresinputsr   rV   r   batch_evaluate   s   
zTEDS.batch_evaluatec                    sT    j dkr fddt||D }|S dd t||D }t| jd j dd}|S )zk Computes TEDS score between the prediction and the ground truth of
            a batch of samples
        r	   c                    s   g | ]
\}}  ||qS r   )rQ   rT   Z	pred_htmlZ	true_htmlr   r   r   rY      s    z,TEDS.batch_evaluate_html.<locals>.<listcomp>c                 S   s   g | ]	\}}||d qS )rZ   r   rb   r   r   r   rY      s    Tr[   )r8   r^   r
   rQ   )r   Z
pred_htmlsZ
true_htmlsr_   r`   r   rc   r   batch_evaluate_html   s   

zTEDS.batch_evaluate_html)Fr	   Nr   )
r   r   r   __doc__r   r>   rF   rQ   ra   rd   r   r   r   r   r4   ]   s    

r4   __main__Nzsample_pred.jsonzsample_gt.json   )r8   )"Zrapidfuzz.distancer   Zaptedr   r   Zapted.helpersr   Zlxmlr   r   collectionsr   Zparallelr
   r   r   r!   r+   r/   objectr4   r   jsonpprintopenfploadrW   rX   Ztedsra   r_   ZPrettyPrinterppr   r   r   r   <module>   s4   o


