o
    'jr                     @   s@   d Z ddlZddlZddlmZ ddlmZ G dd deZ	dS )z
This code is refer from:
https://github.com/open-mmlab/mmocr/blob/main/mmocr/datasets/pipelines/textdet_targets/drrg_targets.py
    N)check_install)normc                       s   e Zd Z														d3 fd
d	Zdd Zdd Zdd Zdd Zdd Zdd Z	dd Z
dd Zdd Zdd Zd d! Zd"d# Zd$d% Zd&d' Zd(d) Zd*d+ Zd,efd-d.Zd/d0 Zd1d2 Z  ZS )4DRRGTargets       @       @	   X        8@333333?      ?      ?皙?c                    sb   t    || _|| _|| _|| _|| _|| _|| _|| _	|	| _
|
| _|| _|| _|| _d| _d S N:0yE>)super__init__orientation_thrresample_stepnum_max_compsnum_min_comps	min_width	max_widthcenter_region_shrink_ratiocomp_shrink_ratiocomp_w_h_ratiotext_comp_nms_thrmin_rand_half_heightmax_rand_half_heightjitter_leveleps)selfr   r   r   r   r   r   r   r   r   r   r   r   r   kwargs	__class__ h/var/www/html/Deteccion_Ine/venv/lib/python3.10/site-packages/paddleocr/ppocr/data/imaug/drrg_targets.pyr      s   

zDRRGTargets.__init__c                 C   s   |j dkr|t|dd| j d }n|t|dd| j  }|j dkr3|t|dd| j d }n|t|dd| j  }tttj|| ddddS )N   Zaxisr'   r&   g      r   )ndimr   r   reshapenpZarccosclipsum)r    Zvec1Zvec2Z	unit_vec1Z	unit_vec2r$   r$   r%   vector_angle:   s   

zDRRGTargets.vector_anglec                 C   s*   t |dksJ t|d |d | j  S )N   r&   r   )lenabsr   r    Zvecr$   r$   r%   vector_slopeG   s   zDRRGTargets.vector_slopec                 C   &   t |dksJ |d t|| j  S )Nr0   r&   r1   r   r   r3   r$   r$   r%   
vector_sinK      zDRRGTargets.vector_sinc                 C   r5   )Nr0   r   r6   r3   r$   r$   r%   
vector_cosO   r8   zDRRGTargets.vector_cosc           "         s  |j dksJ |jd dksJ |jd dksJ t|ts J t|dkrt||d g}|dd  |d d   g }g }t D ]4\}} fdd|d |d fD } | }	t| 	||	}
| 	|	d |	d }|
|
 |
| qDt|tj }t|tj }tj|dd}tt|dd  | ddt|d d | dd}|t|| j  }tt }d	| d
|  }|d| 7 }t|d dkr|t|d d   d7  < |d  d7  < |d| 7 }t||g}tt|t|d f}tt|d tt|d  }dtdtj d	  tt|d	 d	 d d  }|t| }tt|D ]}|| ||d |t| d  | d  ||d d f< q;t| |j\}}|| d t| }|d t| }|d t| }||kr||}}||}}||g}||g}||fS | |d |d  | |d |d   | |d |d  | |d |d   k rddgddgg}ddgddgg}nddgddgg}ddgddgg}t||d d  ||d d   t||d d  ||d d    } t||d d  ||d d   t||d d  ||d d    }!| |!| krD|d }|d }||fS |d }|d }||fS )Nr0   r      r&   r'   c                    s   g | ]}|t   qS r$   )r1   ).0xZedge_vecr$   r%   
<listcomp>a   s    z.DRRGTargets.find_head_tail.<locals>.<listcomp>r(         ?g333333?gffffff?g?   r   r   r
   )r*   shape
isinstancefloatr1   r,   vstack	enumerater.   r/   appendarraypimeanmaximumr   maxr   zerosZconcatenatearangesqrtexppowerrangeZunravel_indexZargmaxr4   )"r    pointsr   
pad_pointsZ	theta_sumZadjacent_vec_thetaiZ	edge_vec1Zadjacent_indZadjacent_edge_vecZtemp_theta_sumZtemp_adjacent_thetaZtheta_sum_scoreZadjacent_theta_scoreZpoly_centerZ	edge_distZ
dist_scoreZposition_scorescoreZ	pad_scoreZscore_matrixr<   ZgaussianZ
head_startZtail_incrementZ
tail_startZhead_endZtail_end	head_inds	tail_indsZhorizontal_edge_indsZvertical_edge_indsZvertical_len_sumZhorizontal_len_sumr$   r=   r%   find_head_tailS   s   


"







zDRRGTargets.find_head_tailc                 C   s   |j dksJ |jd dksJ |jd dksJ | || j\}}|| || }}t||g}|d dk r>t||d< ||d |d  }||d |d t|  }tj|ddtj|dd }	|	d dkrp||}
}n||}
}|||
|fS )Nr0   r   r:   r&   r(   )r*   rA   rX   r   r,   rD   r1   rI   )r    rR   rV   rW   Z	head_edgeZ	tail_edgerS   	sideline1	sideline2Zsideline_mean_shiftZtop_sidelineZbot_sideliner$   r$   r%   reorder_poly_edge   s*   
zDRRGTargets.reorder_poly_edgec                 C   s~   |j dksJ t|dksJ t|dd df |d ddf  d |dd df |d ddf  d  }t|}||fS )Nr0   r&   r   r'   )r*   r1   r,   rN   r.   )r    lineedges_lengthtotal_lengthr$   r$   r%   cal_curve_length   s   (

zDRRGTargets.cal_curve_lengthc                 C   sr  |j dksJ |jd dksJ |jd dksJ t|ts J |dks&J | |\}}tt|dd}||d  }tjd|d tj	d| }d}|d g}	|D ]W}
|t
|d k ry|
||d  kry|d7 }|t
|d k ry|
||d  kse|| ||d  }}tj||
 |
| gtj	d|| | j  }t||||d g }|	| qS|	|d  t|	}|S )Nr0   r   r&   Zdtyper'   )r*   rA   rB   intr_   r,   insertZcumsumrM   float32r1   rG   r   dotrF   rD   )r    r\   nr]   r^   Zt_orgZunit_tZt_equidistantZedge_indrR   tZt_lZt_rweightZp_coordsZresampled_liner$   r$   r%   resample_line   s4   
  
zDRRGTargets.resample_linec                 C   s   |j |j   krdksJ  J |jd |jd   kr!dks$J  J |jd dks-J |jd dks6J t|ts=J | |\}}| |\}}|| d }ttt|| d d}| ||}	| ||}
|	|
fS )Nr0   r&   r   r@   )r*   rA   rB   rC   r_   rK   ra   rh   )r    rY   rZ   r   _Zlength1Zlength2Z
avg_lengthZresample_point_numZresampled_line1Zresampled_line2r$   r$   r%   resample_sidelines   s    (zDRRGTargets.resample_sidelinesc                 C   sB   t |tsJ |\}}tt|| || t|| d  }|S r   )rB   tupler2   r,   crossr   )r    pointr\   Zpoint1Zpoint2dr$   r$   r%   dist_point2line   s   zDRRGTargets.dist_point2linec
                 C   s  |j |j   kr|j ksJ  J |j |j   kr+|j   kr+|j   kr+|j ks.J  J t|	ts5J |j \}
}tdt|d D ]G}|| ||d   d }|| ||d   d }| || }| || }|| || ||  |	  }||d  ||d  ||d   |	  }||d  ||d  ||d   |	  }|| || ||  |	  }t||||g	tj
}tj||gdd tj||g|d tj||g|d t|d d df d|d |d d df< t|d d df d|
d |d d df< tj|dd	tj
}tj|dd	tj
}|| }|| d }tj|d |d ftjd}tj||gdd t|dk}||d |d f }t|}| ||| ||d  f||d d df |d d df f< | ||| ||d  f||d d df |d d df f< qCd S )Nr   r&   r0   )colorr(   r`   )rA   rB   rC   rQ   r1   r7   r9   r,   rD   astypeint32cv2fillPolyr-   minrK   rL   uint8argwhereZfliplrro   )r    top_linebot_linecenter_linecenter_region_masktop_height_mapbot_height_mapsin_mapcos_mapZregion_shrink_ratiohwrT   Ztop_mid_pointZbot_mid_pointZ	sin_thetaZ	cos_thetatltrbrblZcurrent_center_boxZ	min_coordZ	max_coordZbox_szZcenter_box_maskZindsZinds_xyr$   r$   r%   draw_center_region_maps  sj   "
  
$&z#DRRGTargets.draw_center_region_mapsc                 C   sn  t |tsJ |\}}g }t||ftj}tj||ftjd}tj||ftjd}tj||ftjd}	tj||ftjd}
|D ]}|}| |\}}}}| ||| j\}}|d d d }|| d }| 	|d |d  dkr|d |d  d dk r|d d d }|d d d }|d d d }n!|d |d  d dk r|d d d }|d d d }|d d d }t
t|d |d  | j | j| jd }t
t|d |d  | j | j| jd }t|| j }t|| j }t||| d kr||t||  }||t||  }||t||  }||tj | |||||||	|
| j	 q@|||||	|
fS )Nr`   r'   r0   r   r&   )rB   rk   r,   rL   rv   rc   r[   rj   r   r4   r-   r   r   r   r   ra   r1   rF   rq   rr   r   r   )r    img_size
text_polysr   r   center_linesr{   r|   r}   r~   r   polyZpolygon_pointsri   rx   ry   Zresampled_top_lineZresampled_bot_linerz   Zline_head_shrink_lenZline_tail_shrink_lenZnum_head_shrinkZnum_tail_shrinkr$   r$   r%    generate_center_mask_attrib_maps8  s~   
z,DRRGTargets.generate_center_mask_attrib_mapsc              
   C   s  t |tsJ |dksJ |jdksJ |j\}}| j}| j}|d }t|| j | j	| j
}tt|d d |d d  d }	d|	 d t||krt||td| j	d  ks^J tt||d | j	d d }t|d | j	d }|d }t|| j | j	| j
}tt|d d |d d  d }	t|}
||	||	 |	||	 f |
|	||	 |	||	 f< tt|dd}t|
t||ftj}
t|
dk}t|}tj||}|| }tjj||t|dfd}tjj||t|dfd}dtjjt|dfd d }dtjjt|dfd d }td|d |d  d	  }|| }|| }|| }t|| j | j	| j
}t|d d d d d
f ||||t|gtj}|S )Nr   r0   r&   r:         sizer   r   r'   )rB   ra   r*   rA   r   r   r,   r-   r   r   r   rN   ru   rK   
zeros_likers   Zerodeonesrv   rw   r1   randomchoicerandinthstackrq   rc   )r    num_rand_compsZcenter_sample_maskr   r   r   r   Zmax_rand_heightZmax_rand_widthmarginZinner_center_sample_maskZkernel_sizeZcenter_candidatesZnum_center_candidatesZsample_indsZrand_centersZrand_top_heightZrand_bot_heightZrand_cosZrand_sinscaleheightwidthrand_comp_attribsr$   r$   r%   generate_rand_comp_attribsq  s   
"


z&DRRGTargets.generate_rand_comp_attribsc              	   C   s*  |j d dks	J |j d dksJ t|tsJ |dddf d}|dddf d}|dddf d}|dddf d}|dddf d}|ddd	f d}|ddd
f d}	|tjjt|dfdd |t| |t|   | 7 }|tjjt|dfdd |t| |t|   | 7 }|tjjt|dfdd | | 7 }|tjjt|dfdd | | 7 }|tjjt|dfdd d | 7 }|tjjt|dfdd d | 7 }td|d |d  d  }
||
 }||
 }t	|||||||	g}|S )ay  Jitter text components attributes.

        Args:
            comp_attribs (ndarray): The text component attributes.
            jitter_level (float): The jitter level of text components
                attributes.

        Returns:
            jittered_comp_attribs (ndarray): The jittered text component
                attributes (x, y, h, w, cos, sin, comp_label).
        r&   r   r   Nr)   r0   r@   r:         r   r?   r   r   )
rA   rB   rC   r+   r,   r   r1   r2   rN   r   )r    comp_attribsr   r<   yr   r   cossincomp_labelsr   Zjittered_comp_attribsr$   r$   r%   jitter_comp_attribs  sZ   zDRRGTargets.jitter_comp_attribsc           &   	   C   s  t |tsJ |j|j  kr'|j  kr'|j  kr'|j  kr'|jks*J  J t|}t||ddd || }t|dk}	|	dddf }
|	dddf }||
|f d| j	 }||
|f d| j	 }||
|f d}||
|f d}|	t
|| || g }|	t
|| || g }|| | j }t|| j| j}|d }|dddddf t
| | || g }|dddddf t
| | || g }|dddddf t
| | || g }|dddddf t
| | || g }t
||||gtj}tj|jd dftjd}t
||g}tdd	 dd
lm} ||| j}|jd dkr+|j\}}t|dddddf d|d |dddddf< t|dddddf d|d |dddddf< tj|ddddf dddtj}	|	dddf }|	dddf }
||
|f ||
|f  d}t|| j | j| j}||
|f d}||
|f d}tj|dd\}}||
|f dtj} |dtj}|
dtj}
t
||
||||| g}!| |!| j}!|!jd | jk r*| j|!jd  }"| |"d| }#t|!|#g}!n	| | jd| }!tj|!jd dftjd|!jd  }$t
|$|!g}!|!jd | jkrj|!d| jddf }!| j|!dddf< tj| j|!jd ftjd}%|!|%d|!jd ddf< |%S )a$  Generate text component attributes.

        Args:
            center_lines (list[ndarray]): The list of text center lines .
            text_mask (ndarray): The text region mask.
            center_region_mask (ndarray): The text center region mask.
            top_height_map (ndarray): The map on which the distance from points
                to top side lines will be drawn for each pixel in text center
                regions.
            bot_height_map (ndarray): The map on which the distance from points
                to bottom side lines will be drawn for each pixel in text
                center regions.
            sin_map (ndarray): The sin(theta) map where theta is the angle
                between vector (top point - bottom point) and vector (1, 0).
            cos_map (ndarray): The cos(theta) map where theta is the angle
                between vector (top point - bottom point) and vector (1, 0).

        Returns:
            pad_comp_attribs (ndarray): The padded text component attributes
                of a fixed size.
        r   r&   Nr)   r0   r'   r`   lanmsz	lanms-neo)merge_quadrangle_n9   )r'   r:   r0   r(   )Zconnectivity) rB   listrA   r,   r   rs   Z	polylinesrw   r+   r   r   r   r-   r   r   rq   rc   r   r   r   r   r   rI   rr   ZconnectedComponentsr   r   r   r   rD   r   rL   )&r    r   Z	text_maskr{   r|   r}   r~   r   Zcenter_lines_maskZcomp_centersr   r<   Z
top_heightZ
bot_heightr   r   Ztop_mid_pointsZbot_mid_pointsr   rr   r   r   r   Z
text_compsrU   Zla_nmsZimg_hZimg_wr   ri   Zcomp_label_maskr   r   r   r   Z	num_compsZpad_comp_attribsr$   r$   r%   generate_comp_attribs  s   
....

66
z!DRRGTargets.generate_comp_attribsc                 C   s\   t |tsJ |\}}tj||ftjd}|D ]}tj|tjdd}t	||d q|S )a,  Generate text center region mask and geometry attribute maps.

        Args:
            img_size (tuple): The image size (height, width).
            text_polys (list[list[ndarray]]): The list of text polygons.

        Returns:
            text_region_mask (ndarray): The text region mask.
        r`   )r&   r'   r0   r&   )
rB   rk   r,   rL   rv   rG   rr   r+   rs   rt   )r    r   r   r   r   Ztext_region_maskr   polygonr$   r$   r%   generate_text_region_mask\  s   z%DRRGTargets.generate_text_region_mask	mask_sizec                 C   sB   t j|t jd}|D ]}|t jddd}t||d q
|S )a_  Generate effective mask by setting the ineffective regions to 0 and
        effective regions to 1.

        Args:
            mask_size (tuple): The mask size.
            polygons_ignore (list[[ndarray]]: The list of ignored text
                polygons.

        Returns:
            mask (ndarray): The effective mask of (height, width).
        r`   r&   r'   r0   r   )r,   r   rv   rq   rr   r+   rs   rt   )r    r   Zpolygons_ignoremaskr   instancer$   r$   r%   generate_effective_maskr  s
   z#DRRGTargets.generate_effective_maskc              	   C   s   t |tsJ |d }|d }|d }|j\}}}g }g }	t||D ]\}
}|
du r0|	| q"|| q"| ||f|}| ||f|	}| ||f|\}}}}}}| |||||||}|||||||d}|	| ||d< |S )zGenerate the gt targets for DRRG.

        Args:
            data (dict): The input result dictionary.

        Returns:
            data (dict): The output result dictionary.
        imageZpolysignore_tagsT)gt_text_maskgt_center_region_maskgt_maskgt_top_height_mapgt_bot_height_map
gt_sin_map
gt_cos_mapgt_comp_attribs)
rB   dictrA   ziprF   r   r   r   r   update)r    datar   Zpolygonsr   r   r   ri   Zpolygon_masksZpolygon_masks_ignoretagr   r   r   r   r   r   r   r   r   r   mappingr$   r$   r%   generate_targets  sD   



zDRRGTargets.generate_targetsc                 C   s   |  |}|S )N)r   )r    r   r$   r$   r%   __call__  s   
zDRRGTargets.__call__)r   r   r   r   r   r	   r
   r   r
   r   r   r	   r   )__name__
__module____qualname__r   r/   r4   r7   r9   rX   r[   r_   rh   rj   ro   r   r   r   r   r   r   rk   r   r   r   __classcell__r$   r$   r"   r%   r      sD     W
59D0w2r   )
__doc__rs   numpyr,   Zppocr.utils.utilityr   Znumpy.linalgr   objectr   r$   r$   r$   r%   <module>   s   