
    y h                     T    d dl Z d dlmZ d dlmc mZ  G d dej                        Zy)    Nc                   ^     e Zd Z fdZd Zd ZddZd Zd Zd Z	d Z
d	 Zdd
Zd Z xZS )PSS_Lossc                 d    t         t        |           d| _        t	        d|z   dz         | _        y )Ngư>zself._loss)superr   __init__epseval	criterion)selfcls_loss	__class__s     {/home/developers/rajanand/mypropertyqr-fmb-refixing-v2/venv/lib/python3.12/site-packages/easyocr/DBNet/decoders/pss_loss.pyr   zPSS_Loss.__init__   s-    h&(g07:;    c                     t        j                  ||z  |z        }t        j                  ||z        t        j                  ||z        z   | j                  z   }dd|z  |z  z
  }|dkD  rt        ||       |S )N          @)torchsumr	   print)r   predgtmintersectionunionlosss          r   	dice_losszPSS_Loss.dice_loss   sn    yyb+		$q&!EIIbdO3dhh>3|#E))!8,&r   c                    |dk(  |dk(  z  }|dk(  |dk(  z  }|j                         j                         j                         }|j                         j                         j                         }|dk(  s||dz  k  r| j                  |||      S t	        |dz        }||   }||   }	t        j                  |	d      \  }
}|
d | }|j                         }|j                  j                  d       |j                         }|j                         }|j                  j                  d       |j                         }t        j                  ||f      }t        j                  ||f      }t        j                  ||z        }t        j                  |      t        j                  |      z   | j                  z   }dd|z  |z  z
  }|S )	Nr   r   g      @   T)
descendingg      ?r   )floatr   itemr   intr   sortclonedatafill_detachcatr	   )r   r   r   r   	pos_index	neg_indexpos_numneg_numpos_predneg_predneg_sort_sampled_neg_predpos_gtneg_gttpredtgtr   r   r   s                       r   dice_ohnm_losszPSS_Loss.dice_ohnm_loss   s   1Wa(	1Wa(	//#'')..0//#'')..0a<7WS[0>>$A..'!)nGIHIH**X$?KHa'1^^%FKKc"]]_F%++-FKKa ]]_FIIx)9:;E))VV,-C 99US[1LIIe$uyy}4txx?Es\)E11Dr   c                    |dk(  j                         }|dk(  j                         }||z  t        j                  d|j                  z
  |      z  d|z
  |z  t        j                  |j                  |      z  z   }t	        j
                  |||d      }	t        j                  |	|z        | j                  |j                         z   z  }
|
dz  }
|
S )Nr   r   noneweight	reduction
   )r!   r   powr&   Fbinary_cross_entropyr   r	   )r   r   r   r   alphagammapos_maskneg_maskmasklr   s              r   
focal_losszPSS_Loss.focal_loss.   s    !G??$!G??$X~IIa		k5)*+,U7+YYtyy%0+11 ""4DFKyy1~txx/0
r   c                    |j                         \  }}}t        j                  |      dk(  sJ ||j                            }||j                            }|dk(  j	                         }	|	j                         d|	z
  j                         j                         z  d|	z
  j                         |	j                         j                         z  z   }||	j                  d      z  }t        j                  |||d      }
|
S )Nr   r   r   )r<   )
sizer   maxbytesqueezer!   r   r"   r?   r@   )r   r   r   r   nhw	pos_neg_p	pos_neg_trC   r   s              r   wbce_orig_losszPSS_Loss.wbce_orig_loss9   s    ))+1a		""#"N	qvvxL	N++-NNH11388::\  "X\\^%8%8%::;q!"%%iAOr   c                 D   |dk(  j                         |z  }|dk(  j                         |z  }||j                         z  |j                         z  |z   }t        j                  |||d      }t	        j                  |      |j                         | j
                  z   z  }|S )Nr   r   r9   r:   )r!   r   r?   r@   r   r	   )	r   r   r   r   rC   rD   rE   rF   r   s	            r   	wbce_losszPSS_Loss.wbce_lossE   s    !G??$Q&!G??$Q& (,,.(8<<>9HD""4DFKyy|QUUWTXX-.r   c                 x    t        j                  |||d      }||j                         | j                  z   z  }|S )Nr   r:   )r?   r@   r   r	   )r   r   r   r   rF   r   s         r   bce_losszPSS_Loss.bce_lossO   s6    ""4AG!%%'$(("#r   c                 V    | j                  |||      | j                  |||      z   dz  S Nr   )r   rV   r   r   r   r   s       r   dice_bce_losszPSS_Loss.dice_bce_lossT   s,    tR+dmmD"a.HHCOOr   c                 V    | j                  |||      | j                  |||      z   dz  S rX   )r7   rV   rY   s       r   dice_ohnm_bce_losszPSS_Loss.dice_ohnm_bce_lossW   s/    ##D"a04==r13MMQTTTr   c           	      f   |dk(  r| j                  |||      }|S |dk(  rb| j                  ||d d d dd d d d f   |      }|d d dd d d d f   }t        j                  |d      \  }}|| j                  |||      z  }|S |dk(  r| j                  |d d d dd d d d f   |d d d dd d d d f   |      }	|d d dd d d d f   }t        j                  |d      \  }}|	| j                  |||      z  }	| j                  |d d dd d d d f   |d d dd d d d f   |      }
|	|
fS t	        d|      S )Nshrinkpss   r   both   zgt_type [%s] is not implemented)get_lossr   rJ   r   NotImplementedError)r   r   r   rE   gt_typer   g_gg_pr1   pss_lossshrink_losss              r   forwardzPSS_Loss.forwardZ   sJ   h==r40DK==r!RaRA+=DQ1aZ.CYYtQ'FCDNN3T22DK}}T!RaRA+%61bqb!Q;NHQ1aZ.CYYtQ'FCsC66H..Q1aZ "Q1aZ.$8K[((&'H'RRr   c                     t        j                  d      }t        |j                  d            D ]2  }|| j	                  |d d |d d d d f   |d d |d d d d f   |      z  }4 |S )Ng        r   )r   tensorrangerI   r   )r   r   r   rE   r   inds         r   rc   zPSS_Loss.get_losso   sd    ||B1& 	OCDNN431#5r!S!Q,7GNND	Or   )g      ?g333333?)r^   )__name__
__module____qualname__r   r   r7   rG   rR   rT   rV   rZ   r\   rj   rc   __classcell__)r   s   @r   r   r      s=    <
4	

PUS*r   r   )r   torch.nnnntorch.nn.functional
functionalr?   Moduler    r   r   <module>ry      s#       nryy nr   