
    hv&                        S SK r S SKrS SKrS SKrS SKJr  S SKr	S SK
Jr  SSKJr  SSKJr  SSKJr  SSKJr  SS	KJr  S SKr " S
 S5      r " S S\5      r\S:X  aa  \" S5      r\R7                  SS9  \R8                  " S5      r\R=                  \5      r\RA                  \S\5      r!\RD                  " S\!5        gg)    N)ImageOnlyTransform   )FaceAnalysis   )get_model_dir)face3d)	get_image)DEFAULT_MP_NAMEc                       \ rS rSr\SS4S jrSS jrS rS rS r	S	 r
S
 rS rS rSS/ SQ4S jr\S 5       r\S 5       rSrg)MaskRenderer   z~/.insightfaceNc                 ^   Xl         X l        X0l        [        X5      n[        R
                  " US5      n[        R                  " U5      (       d   S5       e[        R                  R                  U5      U l
        U R                  R                  U l        [        R
                  " US5      n[        R                  " U5      (       d   S5       e[        R                  R                  R                  U5      nSU l        SU l        SU l        SU l        S	U l        U R                  S
   U R                  S   SsU l        U l        U l        [.        R0                  " U5      nUS S 2S4   U R(                  S
-
  -  US S 2S4'   US S 2S
4   U R*                  S
-
  -  US S 2S
4'   U R*                  US S 2S
4   -
  S
-
  US S 2S
4'   [.        R2                  " U[.        R4                  " UR6                  S   S
45      45      U l        U R                  R                  U l        / SQU l        / SQU l        g )NzBFM.matz/should contains BFM.mat in your model directoryz
BFM_UV.matz2should contains BFM_UV.mat in your model directory)   r   皙?Q??ffffff?r   r      
mask_white	mask_blue
mask_black
mask_green皙?r   r   r   ) mp_namerootinsfar   ospjoinexistsr   morphable_modelMorphabelModelbfmkpt_ind	index_indloadload_uv_coordsuv_size	mask_stxr	mask_styr	mask_etxr	mask_etyrtex_htex_wtex_cnp
zeros_likehstackzerosshapetexcoordX_indmask_image_namesmask_aug_probs)	selfnamer   r   	model_dirbfm_filebfm_uv_file	uv_coordsr6   s	            W/var/www/fran/franai/venv/lib/python3.13/site-packages/insightface/app/mask_renderer.py__init__MaskRenderer.__init__   s   	
!$-	88Iy1zz(##V%VV#))88B))hhy,7zz+&&\(\\&**//>>{K	 .2ll1oQPQ+
TZ==+"1a4DJJN;A"1a4DJJN;Ahq!tn4q8A		8RXXx~~a7H!6L-M"NOXX%%
 W2    c                 (    Xl         X l        X0l        g N)
pre_ctx_idpre_det_threshpre_det_size)r:   ctx_id
det_threshdet_sizes       r@   prepareMaskRenderer.prepare2   s     ($rC   c                 d    SnUS S2S S 24   US S2S S 24'   U[         R                  " X!5      -  nU$ )N      ?r   )r1   dot)r:   shape3DRss       r@   	transformMaskRenderer.transform7   s<     !QAbffQ((rC   c                 
   [         R                  R                  R                  / SQ5      n[        R
                  " U* S-  U* S-  S/5      nUR                  nX-  nU R                  UR                  U5      R                  nU$ )N)r      rW   r   r   )r   meshrT   angle2matrixr1   arrayT)r:   verticeswhR1ts         r@   
preprocessMaskRenderer.preprocess=   sl    [[""//>HHqbAgrQw*+::>>(**b133rC   c                     U R                   R                  XX45      nUR                  5       nX`R                   R                  S S24   $ )Nr   )r$   rT   copyr%   )r:   r\   rS   anglesr`   transformed_verticesprojected_verticess          r@   project_to_2dMaskRenderer.project_to_2dE   sD    #xx11(vI1668!(("2"2BQB"677rC   c                    Uu  pEpgnU R                   R                  XE5      n	U R                   R                  XUU5      n
U R                  U
R                  X25      n
[
        R                  R                  R                  XU5      nU$ rE   )r$   generate_verticesrT   ra   r[   r   rX   to_image)r:   paramsHW	fitted_sp	fitted_epfitted_sfitted_anglesfitted_tfitted_verticesrf   image_verticess               r@   params_to_verticesMaskRenderer.params_to_verticesJ   sz    CI?	hx((44YJ#xx11/]2: <#/C/E/EqL..778LQRSrC   c                     U R                   R                  USS9n[        U5      S:X  a  U$ U R                   R                  X5      $ )Nr   max_numr   )r   getlendraw_on)r:   
face_imagefacess      r@   draw_lmkMaskRenderer.draw_lmkS   s>    

z15u:q=zz!!*44rC   c                    U R                   c_  [        U R                  U R                  SS/S9U l         U R                   R	                  U R
                  U R                  U R                  S9  U R                   R                  USS9n[        U5      S:X  a  g US   R                  S S 2S S24   nU R                  R                  X0R                  S	S
9u  pEpgnXEXgU/$ )N	detectionlandmark_3d_68)r;   r   allowed_modules)rI   rJ   rK   r   rz   r   r   r   )max_iter)r   r   r   r   rL   rF   rG   rH   r|   r}   r   r$   fitr7   )	r:   r   r   landmarkrp   rq   rr   rs   rt   s	            r@   build_paramsMaskRenderer.build_paramsY   s     ::%4<<diiZegwYxyDJJJdoo4CVCVaearars

z15u:q=8**1RaR40BF((,,xYcYcpq,Br?	hxhxHHrC   c                    U R                   S   U R                   S   S4nUu  pEn[        R                  " U R                   S   U R                   S   S4[        R                  S9nUS   US   pUS   US   p[	        XX-  5      [	        XI-  5      p[	        XZ-  5      [	        XK-  5      pX-
  nX-
  n[
        R                  " UUU45      nXX2X24'   U$ )Nr   r   r   )r5   dtyper   )r)   r1   r4   uint8intcv2resize)r:   mask	positionsr)   r^   r]   cuvstxrstyretxretyrstxstyetxetyheightwidths                     r@   generate_mask_uvMaskRenderer.generate_mask_uvh   s    <<?DLLOQ7aXXT\\!_T\\!_a@QlIaLdq\9Q<dqx=#ah-Sqx=#ah-S	zz$0#37CG	rC   FT)r   r   r   r   c           
      2   [        U[        5      (       a  U(       a  SOSn[        X'S9nU R                  X&5      nUR                  u  pnU R                  X9U
5      nS[        R                  R                  R                  XR                  R                  XR                  U R                  R                  X5      -
  S-  nUR                  [        R                  5      nU(       a1  US:H  R                  [        R                  5      nX-  SU-
  U-  -   nU$ U$ )NTF)to_rgbr      )
isinstancestrins_get_imager   r5   rw   r   rX   renderrender_texturer$   full_trianglesr6   astyper1   r   )r:   r   
mask_imagerm   input_is_rgb
auto_blendr   r   uv_mask_imager^   r]   r   rv   outputmask_bdfinals                   r@   render_maskMaskRenderer.render_maskv   s   j#&&)TuF&zAJ--jD  A001=FKK&&55nhhF]F]`mo|o|  C  G  G  V  V  XY  _  _  ad  drxx(s{**2884G&!G)V);;ELrC   c                     [        U S   5      n[        U S   5      n[        U S   5      /n[        U S   5      n[        U S   5      nX-   U-   U-   U-   $ )Nr   r   r   r      )listfloatrm   p0p1p2p3p4s         r@   encode_paramsMaskRenderer.encode_params   s[    &)_&)_F1I&)_&)_uRx{2~rC   c                 |   U SS n[         R                  " U[         R                  S9R                  S5      nU SS n[         R                  " U[         R                  S9R                  S5      nU S   n[	        U SS 5      nU SS n[         R                  " U[         R                  S9R                  S5      nXX4U4$ )	Nr      )r   )r               )r1   rZ   float32reshapetupler   s         r@   decode_paramsMaskRenderer.decode_params   s    Ac]XXb

+33W=C_XXb

+33W=C[6#c?#C_XXb

+33W=rr!!rC   )r7   r$   r&   r   r9   r,   r-   r8   r*   r+   r   rF   rH   rG   r   r0   r.   r/   r6   r)   )r   g      ?   r   )__name__
__module____qualname____firstlineno__r
   rA   rL   rT   ra   rh   rw   r   r   r   r   staticmethodr   r   __static_attributes__ rC   r@   r   r      ss    +2B$ 3H%
8
5I GLZ^j :   	" 	"rC   r   c                   b   ^  \ rS rSr/ SQ/ SQSSSS4U 4S jjrS	 r\S
 5       rS rS r	Sr
U =r$ )MaskAugmentation   r   r   r   gffffff?FrO   c                    > [         [        U ]  XV5        [        5       U l        [        U5      S:  d   e[        U5      [        U5      :X  d   eXl        X l        X0l        X@l	        g )Nr   )
superr   rA   r   rendererr}   
mask_names
mask_probsh_lowh_high)r:   r   r   r   r   always_applyp	__class__s          r@   rA   MaskAugmentation.__init__   sZ     	.|?$:q   :J///$$
rC   c           
         [        U5      S:X  d  [        U5      S:X  d   S5       e[        U5      S:X  a  US   S:  a  U$ USS  nU R                  R                  U5      nU R                  R                  XUSSUS	S
/S9nU$ )N   r   z.make sure the rec dataset includes mask paramsr   g        r   Tr   r   r   )r   r   )r}   r   r   r   )r:   imagehlabel	mask_nameh_posrm   mask_paramss          r@   applyMaskAugmentation.apply   s     6{C3v;#3e5ee3v;ay}ABZFmm11&9))%KVZgjlqsvx{f|)}rC   c                 
    SS/$ )Nr   r   r   r:   s    r@   targets_as_params"MaskAugmentation.targets_as_params   s    ""rC   c                     US   n[         R                  R                  U R                  U R                  S9n[         R                  R                  U R                  U R                  5      nX#US.$ )Nr   )r   )r   r   r   )r1   randomchoicer   r   uniformr   r   )r:   rm   r   r   r   s        r@   get_params_dependent_on_targets0MaskAugmentation.get_params_dependent_on_targets   sV    !II$$T__$H			!!$**dkk: 5IIrC   c                     g)N)r   r   r   r   r   r   s    r@   get_transform_init_args_names.MaskAugmentation.get_transform_init_args_names   s    >rC   )r   r   r   r   r   )r   r   r   r   rA   r   propertyr   r   r   r   __classcell__)r   s   @r@   r   r      sF     O(( # #J? ?rC   r   __main__anteloper   )rK   zTom_Hanks_54745.pngr   zoutput_mask.jpg)#ossysdatetimenumpyr1   os.pathpathr   albumentationsA(albumentations.core.transforms_interfacer   face_analysisr   utilsr   
thirdpartyr   datar	   r   r
   r   r   r   r   toolrL   imreadr   r   rm   r   mask_outimwriter   rC   r@   <module>r     s         G ' !  - # 
Y" Y"v/?) /?d z
#DLL)L$JJ,-Eu%F {F;HKK!8, rC   