
    h'                     V    S SK Jr  S SKrS SKrS SKrS SKrSSKJr  S/r	 " S S5      r
g)    )divisionN   )
face_alignArcFaceONNXc                   <    \ rS rSrS
S jrS rS rS rS rS r	S	r
g)r      Nc                 ,   Uc   eXl         X l        SU l        SnSn[        R                  " U R                   5      nUR
                  n[        UR                  S S 5       H  u  pxUR                  R                  S5      (       d   UR                  R                  S5      (       a  SnUR                  R                  S5      (       d"  UR                  R                  S5      (       d  M  SnM     U(       a  U(       a  S	n	S
n
OSn	Sn
Xl
        Xl        U R                  c&  [        R                  " U R                   S 5      U l        U R                  R                  5       S   nUR                  nUR                  n[!        USS S S S2   5      U l        Xl        U R                  R'                  5       n/ nU H  nUR)                  UR                  5        M      Xl        Xl        [/        U R,                  5      S:X  d   eUS   R                  U l        g )NrecognitionF   Sub_minusTMul_mulg              ?g     _@r   r         )
model_filesessiontasknameonnxloadgraph	enumeratenodename
startswith
input_mean	input_stdonnxruntimeInferenceSession
get_inputsshapetuple
input_sizeinput_shapeget_outputsappend
input_nameoutput_nameslenoutput_shape)selfr   r   find_subfind_mulmodelr   nidr   r   r   	input_cfgr&   r)   outputsr*   outs                    \/var/www/fran/franai/venv/lib/python3.13/site-packages/insightface/model_zoo/arcface_onnx.py__init__ArcFaceONNX.__init__   s   %%%$%		$//*"5::bq>2ICyy##E**dii.B.B8.L.Lyy##E**dii.B.B6.J.J 3 JIJI$"<<&77NDLLL++-a0	oo^^
Aa 02 67&,,**,C) $(4$$%q(((#AJ,,    c                 J    US:  a  U R                   R                  S/5        g g )Nr   CPUExecutionProvider)r   set_providers)r-   ctx_idkwargss      r5   prepareArcFaceONNX.prepare=   s$    !8LL&&(>'?@ r8   c                     [         R                  " XR                  U R                  S   S9nU R	                  U5      R                  5       Ul        UR                  $ )Nr   )landmark
image_size)r   	norm_cropkpsr%   get_featflatten	embedding)r-   imgfaceaimgs       r5   getArcFaceONNX.getA   sE    ##C((tWXGYZt,446~~r8   c                     SSK Jn  UR                  5       nUR                  5       n[        R                  " X5      U" U5      U" U5      -  -  nU$ )Nr   )norm)numpy.linalgrN   ravelnpdot)r-   feat1feat2rN   sims        r5   compute_simArcFaceONNX.compute_simF   s@    %ffU"d5kDK&?@
r8   c                 V   [        U[        5      (       d  U/nU R                  n[        R                  R                  USU R                  -  UU R                  U R                  U R                  4SS9nU R                  R                  U R                  U R                  U05      S   nU$ )Nr   T)swapRBr   )
isinstancelistr%   cv2dnnblobFromImagesr   r   r   runr*   r)   )r-   imgsr%   blobnet_outs        r5   rE   ArcFaceONNX.get_featM   s    $%%6D__
ww%%dC$..,@*'+&Ybf & h,,""4#4#4t6MNqQr8   c                     XR                   -
  U R                  -  nU R                  R                  U R                  U R
                  U05      S   nU$ )Nr   )r   r   r   r_   r*   r)   )r-   
batch_datara   rb   s       r5   forwardArcFaceONNX.forwardW   sF    __,>,,""4#4#4t6MNqQr8   )
r   r)   r&   r%   r   r   r*   r,   r   r   )NN)__name__
__module____qualname____firstlineno__r6   r>   rK   rV   rE   rf   __static_attributes__ r8   r5   r   r      s"    '-RA
r8   )
__future__r   numpyrQ   r\   r   r    utilsr   __all__r   rm   r8   r5   <module>rr      s1       
    
G Gr8   