
    9i&                        S SK r S SKJrJrJrJrJr  S SKrS SKr	S SK
r
S SKJs  Jr  S SKJr  / SQrSS jrS rS rS rS	 rSS
 jrSS jrS rSS jrSS jr  SS\\\4   S\	R:                  S\\	R<                     S\\	R<                     S\\\\	R<                  \\\\\4   4   4   4
S jjr    SS\\\\\4   S\\	R:                     S\\	R<                     S\\	R<                     S\\\	R<                  4   4
S jjr!g)    N)DictListOptionalTupleUnion	rearrange)
gen_diffuse_maskget_crop_bboxget_roi_without_paddingpatch_aggregation_overlappatch_partition_overlappreprocess_roiresize_on_long_sideroi_to_tensorsmooth_border_mg
whiten_imgc                 6   U R                   S   nU R                   S   nX#:  a<  US-  U-  n[        R                  " U [        X4-  5      U4[        R                  S9nXT4$ US-  U-  n[        R                  " X[        X$-  5      4[        R                  S9nXT4$ )Nr            ?)interpolation)shapecv2resizeintINTER_LINEAR)img	long_side
src_height	src_widthscale_imgs         j/var/www/html/land-doc-ocr/venv/lib/python3.13/site-packages/modelscope/models/cv/skin_retouching/utils.pyr   r      s    1J		!IC*,zz#i'()4**, ; C)+zzS!345**, ;    c           	      `   / nU  H<  nUS   S:X  a    O2UR                  US   S   US   S   US   S   US   S   S.5        M>     [        U5      n/ n[        U5       H  nX   n[        US	   US
   -
  5      n[        US   US   -
  5      nUS
   US	   -   S-  US   US   -   S-  4n	X:  a  UOUn
Sn[	        X-  5      n
Sn[	        U	S   X-  S-  -
  5      [	        U	S   X-  S-  -   5      [	        U	S   X-  S-  -
  5      [	        U	S   X-  S-  -   5      S.nUR                  U5        M     U$ )Nscorebboxr   r         )x1y1x2y2r-   r+   r.   r,   g      ?)r+   r-   r,   r.   )appendlenrangeabsr   )detecting_resultsboxesanno
face_countsuitable_bboxesi	face_bboxface_bbox_widthface_bbox_heightface_bbox_centersquare_bbox_lengthenlarge_ratio	sideScalesquare_bboxs                 r#   r   r   $   s   E!=Bv,q/v,q/v,q/v,q/	
 	 " UJO:H	io	$?@y4@A&t_y>!C&t_y>!CE 2B1S-Yh !CD	  #i&Dq&HHI #i&Dq&HHI #i&Dq&HHI&q)I,JQ,NNO
 	{+3 6 r$   c                     [        US   S5      n[        US   U R                  S   5      n[        US   S5      n[        US   U R                  S   5      nXU2XE24   nUSX#XE/4$ )Nr,   r   r.   r+   r-   r   )maxminr   )r   r(   crop_tcrop_bcrop_lcrop_rrois          r#   r   r   P   sw    dQFdSYYq\*FdQFdSYYq\*F
VmV]*
+CFF333r$   c                 V    [         R                  " U R                  S5      5      S   n U $ )N)r)   r   r   N.)torch
from_numpy	transposer   s    r#   r   r   Y   s%    


3==3
4Y
?CJr$   c                 <    U R                  5       S-  n U S-
  S-  n U $ )N     o@      ?r)   )floatrN   s    r#   r   r   _   s$    
))+
C9/CJr$   c           
      ^   U R                  5       u  pEpgXa-  Xr-  p[        R                  " U X3X3SS4SSS9n / n
[        U5       HQ  n[        U	5       H?  nU S S 2S S 2X-  XS-   -  US-  -   2X,-  X,S-   -  US-  -   24   nU
R	                  U5        MA     MS     [
        R                  " U
SS9nU$ )Nr   constant)padmodevaluer   r)   dim)sizeFrU   r1   r/   rK   cat)imagep1p2paddingBCHWhw
patch_listr8   jpatchoutputs                  r#   r   r   f   s    JA!7AGqEEwA6	E J1XqA!QrU|gk'A A&1u!!;;< =Ee$   YYFMr$   c                 >    U S S 2S S 2X3* 2X3* 24   n [        U SXS9nU$ )Nz$(b h w) c p1 p2 -> b c (h p1) (w p2))re   rf   r   )r]   re   rf   r`   rj   s        r#   r   r   |   s3    !Q('(*::;EuDOFMr$   c                     US-
  n[         R                  " XR                  S S SS9S   R                  SSS5      n X-  nUS-   nU$ )NrQ   r)   bilinearrV   r   r   )r[   interpolater   permute)diffuse_maskmgs     r#   r   r      sW    	cB==hhrl55688?1a8H 		B	cBIr$   c                    SnU(       a  Sn[         R                  " [         R                  XD45      n[         R                  " [         R                  XD45      n[        UR                  5      S:X  a  USS2SS2S4   n[         R
                  " XS5      n[         R                  " XS5      n[         R                  " US5      S-  nUR                  5       n[        R                  " U5      R                  U R                  5      n[        R                  " XU/S	S
9S   nUSS2SS2SS2SS24==   S-  ss'   US-  U-  S-   n[        UR                  5      S:X  d   e[        R                  " XpR                  SS SS9S	   R!                  SSS	5      R#                  5       nU R#                  5       nUS-  nSU-  S-   U-  U-  SU-  U-  -   nUS-  nUR%                  5       nUR'                  5       R)                  5       nU$ )z
image: rgb
   P   r*   Nr'   r   )   rv   rP   r   rX   rJ   g      ?g?rQ      r)   rm   rn   )r   getStructuringElementMORPH_ELLIPSEr0   r   dilateerodeblursqueezerK   rL   todevicestackr[   ro   rp   halfbytecpunumpy)	r]   	skin_maskwhitening_degreeflag_bigKernaldilate_kernalsizenew_kernel1new_kernel2	whiten_mgoutput_preds	            r#   r   r      s    ++-ACK++-ACK
9??q aBh'	

915I		)!4IH-5I!!#I  +..u||<IY9=1E GJ KIaQkd"C"22S8Iy1$$$;;r?55688?1@A9CCG46  **,K%K
Y!"$%	MK$?@K %K""$K//#))+Kr$   c                 B   SnSn[         R                  " X4[         R                  S9n[        U5       Hb  n[        U5       HP  nXB:  a  UX-
  ::  a  XR:  a  UX-
  ::  a  SX4U4'   M%  XB::  a  US-  U-  X4U4'   M8  XAU-
  :  d  MB  X-
  S-  U-  X4U4'   MR     Md     [        U5       H\  n[        U5       HJ  nXR::  a  [	        X4U4   US-  U-  5      X4U4'   M$  XQU-
  :  d  M.  [	        X4U4   X-
  S-  U-  5      X4U4'   ML     M^     [         R
                  " U/U -  5      nU$ )Ni  rv   )r   dtyper   )nponesfloat32r1   rC   dstack)out_channels	mask_sizediffuse_withar8   rh   s         r#   r
   r
      sA   IL
y,BJJ?A9y!A Q,&.232C!0I2Q$"c'L0Q$,./$=C/,>Q$ "  9y!A a1gq3w'=>Q$,./a1g	'<|'KLQ$	 "  			1#$%AHr$   target_sizer]   bboxes	keypointsreturnc           	      $   U u  pEUR                   SS u  pgXW:  a  [        SU SU 35      eXF:  a  [        SU SU 35      eXd:X  a  SnSn	OXF-
  n
U
S-  nX-
  n	Xu:X  a  SnSnOXW-
  nUS-  nX-
  nXX4[        R                  " XXU[        R                  5      S.nUbH  USS2S4==   U-  ss'   USS2S4==   U-  ss'   USS2S4==   U-  ss'   USS2S	4==   U-  ss'   X.S
'   Ub&  USS2S4==   U-  ss'   USS2S4==   U-  ss'   X>S'   U$ )a  Pads the image on the sides to the target_size

Args:
    target_size: (target_height, target_width)
    image:
    bboxes: np.array with shape (num_boxes, 4). Each row: [x_min, y_min, x_max, y_max]
    keypoints: np.array with shape (num_keypoints, 2), each row: [x, y]

Returns:
    {
        "image": padded_image,
        "pads": (x_min_pad, y_min_pad, x_max_pad, y_max_pad),
        "bboxes": shifted_boxes,
        "keypoints": shifted_keypoints
    }

Nr)   z2Target width should bigger than image_widthWe got  z4Target height should bigger than image_heightWe got r   )padsr]   r   r*   r   r   )r   
ValueErrorr   copyMakeBorderBORDER_CONSTANT)r   r]   r   r   target_heighttarget_widthimage_heightimage_width	y_min_pad	y_max_pady_pad	x_min_pad	x_max_padx_padresults                  r#   pad_to_sizer      s   . #.M %BQL! ##/.+@ A 	A # ##0/<.B C 	C $		,QJ	%	"		*QJ	%	 y<5Y9..	0F q!t	!q!t	!q!t	!q!t	!!x!Q$9$!Q$9$'{Mr$   r   c                 F   U u  pEpg0 nUb"  UR                   SS u  pXX-
  2XJU-
  24   US'   UbH  USS2S4==   U-  ss'   USS2S4==   U-  ss'   USS2S4==   U-  ss'   USS2S4==   U-  ss'   X(S'   Ub&  USS2S4==   U-  ss'   USS2S4==   U-  ss'   X8S'   U$ )	a  Crops patch from the center so that sides are equal to pads.

Args:
    image:
    pads: (x_min_pad, y_min_pad, x_max_pad, y_max_pad)
    bboxes: np.array with shape (num_boxes, 4). Each row: [x_min, y_min, x_max, y_max]
    keypoints: np.array with shape (num_keypoints, 2), each row: [x, y]

Returns: cropped image

{
        "image": cropped_image,
        "bboxes": shifted_boxes,
        "keypoints": shifted_keypoints
    }

Nr)   r]   r   r   r*   r   r   )r   )r   r]   r   r   r   r   r   r   r   heightwidths              r#   unpad_from_sizer     s    . 26.I)FBQ&*< < ))*; ;!< =w q!t	!q!t	!q!t	!q!t	!!x!Q$9$!Q$9$'{Mr$   )i   )    )F)r*   )NN)NNN)"timetypingr   r   r   r   r   r   r   r   rK   torch.nn.functionalnn
functionalr[   einopsr	   __all__r   r   r   r   r   r   r   r   r   r
   r   arrayndarraystrr   r    r$   r#   <module>r      sU    5 5 
     $)X4,$N: $(&*	HsCxH88H RZZ H 

#	H
 
#uRZZsCc'9!::;
;<HZ !%#'&*	.
S#s"
#.BHH. RZZ . 

#	.
 
#rzz/.r$   