
    ` j9%                        S SK r S SKrS SKJr   " S S5      r " S S\R
                  5      r " S S\R
                  5      r " S S	\R
                  5      r " S
 S\R
                  5      r	 " S S\R
                  5      r
 " S S\
5      r " S S\R
                  5      r " S S\R
                  5      r " S S\R
                  5      rg)    Nc                   $    \ rS rSrSrS rS rSrg)SegDetectorLossBuilder   a  
Build loss functions for SegDetector.
Details about the built functions:
    Input:
        pred: A dict which contains predictions.
            thresh: The threshold prediction
            binary: The text segmentation prediction.
            thresh_binary: Value produced by `step_function(binary - thresh)`.
        batch:
            gt: Text regions bitmap gt.
            mask: Ignore mask,
                pexels where value is 1 indicates no contribution to loss.
            thresh_mask: Mask indicates regions cared by thresh supervision.
            thresh_map: Threshold gt.
    Return:
        (loss, metrics).
        loss: A scalar loss value.
        metrics: A dict contraining partial loss values.
c                 (    Xl         X l        X0l        g N)
loss_class	loss_argsloss_kwargs)selfr   argskwargss       y/root/GenerationalWealth/GenerationalWealth/venv/lib/python3.13/site-packages/easyocr/DBNet/decoders/seg_detector_loss.py__init__SegDetectorLossBuilder.__init__   s    $!    c                     [        [        R                  [           U R                  5      " U R
                  0 U R                  D6$ r   )getattrsysmodules__name__r   r	   r
   )r   s    r   buildSegDetectorLossBuilder.build!   s1    s{{8,doo>cRVRbRbccr   )r	   r   r
   N)r   
__module____qualname____firstlineno____doc__r   r   __static_attributes__ r   r   r   r      s    ("
dr   r   c                   6   ^  \ rS rSrSrSU 4S jjrS rSrU =r$ )DiceLoss%   >
DiceLoss on binary.
For SegDetector without adaptive module.
c                 P   > [         [        U ]  5         SSKJn  U" U5      U l        g N   r    )superr    r   	dice_lossloss)r   epsLoss	__class__s      r   r   DiceLoss.__init__+   s    h&(/I	r   c                 P    U R                  US   US   US   5      nU[        US94$ Nbinarygtmask)r(   r)   dictr   predbatchr)   s       r   forwardDiceLoss.forward0   0    yyhteFmDTD)))r   r)   ư>	r   r   r   r   r   r   r8   r   __classcell__r,   s   @r   r    r    %   s    

* *r   r    c                   6   ^  \ rS rSrSrSU 4S jjrS rSrU =r$ )BalanceBCELoss5   r"   c                 N   > [         [        U ]  5         SSKJn  U" 5       U l        g )Nr%   BalanceCrossEntropyLoss)r'   rB   r   balance_cross_entropy_lossrF   r)   )r   r*   rF   r,   s      r   r   BalanceBCELoss.__init__;   s    nd,.G+-	r   c                 P    U R                  US   US   US   5      nU[        US94$ r/   r3   r5   s       r   r8   BalanceBCELoss.forward@   r:   r   r;   r<   r>   r@   s   @r   rB   rB   5   s    
.
* *r   rB   c                   6   ^  \ rS rSrSrSU 4S jjrS rSrU =r$ )AdaptiveDiceLossE   zN
Integration of DiceLoss on both binary
    prediction and thresh prediction.
c                 j   > [         [        U ]  5         SSKJn  U" U5      U l        U" U5      U l        g r$   )r'   rL   r   r(   r    	main_lossthresh_loss)r   r*   r    r,   s      r   r   AdaptiveDiceLoss.__init__K   s+    .0'!##C=r   c                     [        U[        5      (       d   eSU;   d   eSU;   d   eUS   nUS   nUS   nUS   nU R                  X5U5      nU R                  XEU5      nXx-   n	U	[        XxS94$ )Nr0   thresh_binaryr1   r2   rO   rP   )
isinstancer4   rO   rP   )
r   r6   r7   r0   rS   r1   r2   rO   rP   r)   s
             r   r8   AdaptiveDiceLoss.forwardQ   s    $%%%%4$&&&h_-4[V}NN6t4	&&}$?&TIGGGr   rT   r<   r>   r@   s   @r   rL   rL   E   s    
)H Hr   rL   c                   <   ^  \ rS rSrSrSU 4S jjrS rS rSrU =r	$ )AdaptiveInstanceDiceLoss`   z4
InstanceDiceLoss on both binary and thresh_bianry.
c                 $  > [         [        U ]  5         SSKJnJn  U" 5       U l        U" 5       U l        U" 5       U l        U" 5       U l	        [        R                  " [        [        R                  " [        R                  " S5      5      [        R                  " [        R                  " S5      5      [        R                  " [        R                  " S5      5      [        R                  " [        R                  " S5      5      S95      U l        g )Nr%   )InstanceDiceLossr    )mainthreshmain_instancethresh_instance)r'   rX   r   r(   r[   r    rO   main_instance_lossrP   thresh_instance_lossnnParameterDictr4   	Parametertorchonesweights)r   
iou_threshr]   r[   r    r,   s        r   r   !AdaptiveInstanceDiceLoss.__init__e   s    &689!"2"4#:$4$6!''ejjm,<<

1.,,uzz!}5LLA7	)9 :r   c                 `    X!-  [         R                  " [         R                  " U5      5      -   $ r   )re   logsqrt)r   weightr)   s      r   partial_loss%AdaptiveInstanceDiceLoss.partial_lossr   s!    }uyyF);<<<r   c                 4   U R                  US   US   US   5      nU R                  US   US   US   5      nU R                  US   US   US   5      nU R                  US   US   US   5      nU R	                  U R
                  S   U5      U R	                  U R
                  S   U5      -   U R	                  U R
                  S   U5      -   U R	                  U R
                  S   U5      -   n[        UUUUS	9nUR                  U R
                  5        Xx4$ )
Nr0   r1   r2   rS   r\   r]   r^   r_   )rO   rP   r`   ra   )rO   rP   r`   ra   rn   rg   r4   update)	r   r6   r7   rO   rP   r`   ra   r)   metricss	            r   r8    AdaptiveInstanceDiceLoss.forwardu   s:   NN4>5;fN	&&tO'<eDk5QW=Y!44NE$Kv8#88!5;f ?  f!5yA""4<<#9;GH""4<<#@BTUV ""4<<0A#BDXYZ #1!5	7
 	t||$}r   )r`   rO   ra   rP   rg   )g?g333333?)
r   r   r   r   r   r   rn   r8   r   r?   r@   s   @r   rX   rX   `   s    := r   rX   c                   6   ^  \ rS rSrSrSU 4S jjrS rSrU =r$ )
L1DiceLoss   z9
L1Loss on thresh, DiceLoss on thresh_binary and binary.
c                 v   > [         [        U ]  5         [        US9U l        SSKJn  U" 5       U l        X l        g )Nr*   r%   
MaskL1Loss)r'   ru   r   rL   r(   l1_lossrz   l1_scale)r   r*   r|   rz   r,   s       r   r   L1DiceLoss.__init__   s-    j$(*)c2'!| r   c                     U R                  X5      u  p4U R                  US   US   US   5      u  pVX0R                  U-  -   nUR                  " S0 UD6  Xt4$ )Nr]   
thresh_mapthresh_maskr   )r(   r{   r|   rq   )r   r6   r7   r(   rr   r{   	l1_metricr)   s           r   r8   L1DiceLoss.forward   sd    !^^D8	!\\NE,/}1EG ==722##}r   r(   r{   r|   r=   
   r>   r@   s   @r   ru   ru      s    ! r   ru   c                   "    \ rS rSrSrSS jrSrg)FullL1DiceLoss   zw
L1loss on thresh, pixels with topk losses in non-text regions are also counted.
DiceLoss on thresh_binary and binary.
c                     [         R                  R                  U 5        [        US9U l        SSKJn  U" 5       U l        X l        g )Nrx   r%   )BalanceL1Loss)rb   Moduler   rL   r(   r{   r   r|   )r   r*   r|   r   s       r   r   FullL1DiceLoss.__init__   s1    
		4 )c2*$ r   r   Nr   )r   r   r   r   r   r   r   r   r   r   r   r      s    
!r   r   c                   6   ^  \ rS rSrSrSU 4S jjrS rSrU =r$ )L1BalanceCELoss   z
Balanced CrossEntropy Loss on `binary`,
MaskL1Loss on `thresh`,
DiceLoss on `thresh_binary`.
Note: The meaning of inputs can be figured out in `SegDetectorLossBuilder`.
c                    > [         [        U ]  5         SSKJn  SSKJn  SSKJn  U" US9U l        U" 5       U l        U" 5       U l	        X l
        X0l        g Nr%   r&   ry   rE   rx   )r'   r   r   r(   r    r{   rz   rG   rF   bce_lossr|   	bce_scaler   r*   r|   r   r    rz   rF   r,   s          r   r   L1BalanceCELoss.__init__   B    ot-/''G!c*!|/1 "r   c                 H   U R                  US   US   US   5      n[        US9nSU;   as  U R                  US   US   US   5      u  pVU R                  US   US   US   5      nXtS	'   XpR                  U-  -   X0R
                  -  -   nUR                  " S
0 UD6  X4$ UnX4$ )Nr0   r1   r2   r   r]   r   r   rS   rP   r   )r   r4   r{   r(   r|   r   rq   )	r   r6   r7   r   rr   r{   r   r(   r)   s	            r   r8   L1BalanceCELoss.forward   s    ==hteFmL)t!%d8neL>QSXYfSg!hGtO'<eDk5QW=YI%.M"}}w66NN9RRDNN'Y' } D}r   r   r   r(   r{   r|   r=   r      r>   r@   s   @r   r   r      s    
# r   r   c                   6   ^  \ rS rSrSrSU 4S jjrS rSrU =r$ )L1BCEMiningLoss   zl
Basicly the same with L1BalanceCELoss, where the bce loss map is used as
    attention weigts for DiceLoss
c                    > [         [        U ]  5         SSKJn  SSKJn  SSKJn  U" US9U l        U" 5       U l        U" 5       U l	        X l
        X0l        g r   )r'   r   r   r(   r    r{   rz   rG   rF   r   r|   r   r   s          r   r   L1BCEMiningLoss.__init__   r   r   c                    U R                  US   US   US   SS9u  p4U R                  US   US   US   5      u  pVXDR                  5       -
  UR                  5       UR                  5       -
  -  nU R	                  US	   US   US   US
-   S9n[        US9nXxS'   XpR                  U-  -   X0R                  -  -   n	UR                  " S0 UD6  X4$ )Nr0   r1   r2   T)return_originr]   r   r   rS   r%   )rg   r   rP   r   )	r   r{   minmaxr(   r4   r|   r   rq   )
r   r6   r7   r   bce_mapr{   r   r(   rr   r)   s
             r   r8   L1BCEMiningLoss.forward   s     MM$x.%+uV}8< * >!\\$x.%:MuUbOcd[[]*w{{}w{{}/LMNN!5;&M7Q; # 0	 )!*==722X5NN##}r   r   r   r>   r@   s   @r   r   r      s    

# r   r   c                   6   ^  \ rS rSrSrSU 4S jjrS rSrU =r$ )L1LeakyDiceLoss   zK
LeakyDiceLoss on binary,
MaskL1Loss on thresh,
DiceLoss on thresh_binary.
c                    > [         [        U ]  5         SSKJnJn  SSKJn  U" US9U l        U" 5       U l        U" US9U l	        X0l
        g )Nr%   )r    LeakyDiceLossry   )coverage_scalerx   )r'   r   r   r(   r    r   r{   rz   rO   rP   r|   )r   r*   r   r|   r    r   rz   r,   s          r   r   L1LeakyDiceLoss.__init__   s=    ot-/6'&nE!|#, r   c                 
   U R                  US   US   US   5      u  p4U R                  US   US   US   5      nU R                  US   US   US   5      u  pgUR                  " S	0 UDSU0D6  X5-   X`R                  -  -   nX4$ )
Nr0   r1   r2   rS   r]   r   r   rP   r   )rO   rP   r{   rq   r|   )	r   r6   r7   rO   rr   rP   r{   r   r)   s	            r   r8   L1LeakyDiceLoss.forward  s    !^^DNE$KvW	&&tO'<eDk5QW=Y!\\NE,/}1EG<<<&==)@@}r   )r{   r|   rO   rP   )r=   r   r   r>   r@   s   @r   r   r      s    ! r   r   )r   re   torch.nnrb   r   r   r    rB   rL   rX   ru   r   r   r   r   r   r   r   <module>r      s    
  d d<*ryy * *RYY * Hryy H6&ryy &R ,!Z !bii Dbii Bbii r   