
    ]i                         S r SSKJr  SSKr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SKJrJr  SSKJr  SS	KJr  \R*                  " \5      rS
 rSS jrSS jrg)zTest time evaluation, either using the original SDR from [Vincent et al. 2006]
or the newest SDR definition from the MDX 2021 competition (this one will
be reported as `nsdr` for `new sdr`).
    )futuresN)LogProgress   )apply_model)convert_audio
save_audio)distrib)DummyPoolExecutorc                 P   U R                  5       S:X  d   eUR                  5       S:X  d   eSn[        R                  " [        R                  " U 5      SS9n[        R                  " [        R                  " X-
  5      SS9nX2-  nXB-  nS[        R                  " X4-  5      -  nU$ )z
Compute the SDR according to the MDX challenge definition.
Adapted from AIcrowd/music-demixing-challenge-starter-kit (MIT license)
   gHz>)      dim
   )r   thsumsquarelog10)
references	estimatesdeltanumdenscoress         L/mnt/rpi/tmp/demucs-venv-sys/lib/python3.13/site-packages/demucs/evaluate.pynew_sdrr      s    
 >>q   ==?aE
&&:&F
3C
&&:12
?CLCLC"((39%%FM    c           
         U R                  SS5      R                  5       n UR                  SS5      R                  5       n[        U R                  5       S    UR                  5       S    5      S   nU(       d  S U4$ U R	                  5       n UR	                  5       n[
        R                  R                  XSUUSSS9S S nXe4$ )Nr   r   r   F)compute_permutationwindowhopframewise_filtersbsseval_sources_version)	transposedoubler   cpunumpymusevalmetricsbss_eval)r   r   winr"   compute_sdr
new_scoresr   s          r   
eval_trackr0   .   s    %%a+224J##Aq)002I)$/1FGJJZ%%'
OO%	)) %#$) * + ,/B0 !!r   c                    U R                   nU R                  S-  nUR                  SSS9  U R                  S-  nUR                  SSS9  UR                  R                  c,  [
        R                  " UR                  R
                  S/SS9nO+[
        R                  " UR                  R                  S/SS9nUR                  R                  nS	nU R                  n[        S
UR                  -  5      n	[        S
UR                  -  5      n
[        [        R                  [        U5      [        R                   5      n[#        [$        XR&                  R(                  SS9n/ nUR                  R*                  (       a  [,        R.                  O[0        nU" UR                  R*                  5       nU GH  nUR2                  U   n[4        R6                  " UR8                  5      R;                  5       R=                  5       nUR?                  5       S:X  a  US   nURA                  U RB                  5      nURE                  SS9nUURE                  5       -
  URG                  5       -  n[I        UXhR                  URJ                  5      n[M        UUS   UR                  RN                  UR                  RP                  UR                  RR                  S9S   nUURG                  5       -  URE                  5       -   nURA                  U5      n[4        RT                  " URV                   Vs/ s H>  n[4        R6                  " URX                  U   R8                  5      R;                  5       PM@     sn5      nUR?                  5       S:X  a	  USS2S4   nURA                  U5      n[I        UUUR                  URJ                  5      nUR                  RZ                  (       aw  U R                  S-  UR\                  -  nUR                  SSS9  [_        URV                  U5       H1  u  nn[a        URc                  5       UUS-   -  UR                  5        M3     URe                  UR\                  URg                  [h        UUXUS945        GM     [#        [$        XR&                  R(                  SS9n0 nU H  u  nnURk                  5       nUu  nn0 UU'   [m        URV                  5       H  u  nnS[=        UU   5      /0UU   U'   M     Uc  M[  Uu  nnn n![m        URV                  5       Hh  u  nnUU   Ro                  5       U U   Ro                  5       UU   Ro                  5       U!U   Ro                  5       S.n"UU   U   Rq                  U"5        Mj     M     0 n#[        [        R                   5       H)  n$U#Rq                  [        Rr                  " UU$5      5        M+     0 n%[u        [w        U#Ry                  5       5      5      URV                  S      n&U& GH  n'Sn(Sn)URV                   H  n*U#R{                  5        Vs/ s H"  n[|        R~                  " U#U   U*   U'   5      PM$     n+n[|        RD                  " U+5      n,[|        R                  " U+5      n-U,U%U'R                  5       S-   U*-   '   U-U%U'R                  5       S-   S-   U*-   '   U(U,[        URV                  5      -  -  n(U)U-[        URV                  5      -  -  n)M     U(U%U'R                  5       '   U)U%U'R                  5       S-   '   GM     U%sSSS5        $ s  snf s  snf ! , (       d  f       g= f)z
Evaluate model using museval.
compute_sdr=False means using only the MDX definition of the SDR, which
is much faster to evaluate.
resultsT)exist_okparentszresults/testNtest)subsetsis_wavFr(   g      ?Eval)updatesnamer   r   r   )shiftssplitoverlapr   wavz.mp3)r-   r"   r.   z
Eval (BSS)nsdr)SDRSIRISRSAR__med)Bargsfoldermkdirr5   nonhqmusdbDBdsetmusdb_sampleratemodelint
samplerateranger	   ranklen
world_sizer   loggermisc
num_printsworkersr   ProcessPoolExecutorr
   tracksr   
from_numpyaudiotfloatr   todevicemeanstdr   audio_channelsr   r;   r<   r=   stacksourcestargetssaver:   zipr   r(   appendsubmitr0   result	enumeratetolistupdatesharenextitervalueskeysnp	nanmedianmedianlower).solverr.   rF   
output_dirjson_foldertest_setsrc_rateeval_devicerN   r-   r"   indexespendingspoolindextrackmixrefr   r:   r   rG   estimaterZ   
track_namependingr   nsdrsidxtargetsdrisrsirsarrr   
all_trackssrcrk   metric_namesmetric_nameavgavg_of_medianssourcemediansra   rv   s.                                                 r   evaluater   C   s?    ;;D*JdD1--.0KtT2 yy88DIIOOfXdK88DIIOOfXeLyy))HKLLE
b5###
$C
b5###
$CGLL#h-1C1CDG&'993G3G%'GH*.))*;*;7&&ARD	dii	 DEOOE*E--,..0668CwwyA~$i&&'C((q(/C#swwy0CX/?/?AUAUVC#E3t9+/99+;+;499??,0II,=,=??@BI "CGGI-
:I![1IJO--X-$u}}T2889;;=-XZJ~~1$'40
#{3J&z8','7'79M9MOJyy~~.;dD9&)%--&CND(x||~v/GIYIYZ 'D OOUZZJ	sQ\ *5 *^ _ `; @ vx9M9M$02#+Jnn&G#MFE!#F:(7V.4uU3Z7H6I-Jz"6*  8!'-$c3#,U]]#;KC"3x0"3x0"3x0"3x0	F :&v.55f= $< $," 
++,CgmmFC89 - D!2!2!456u}}Q7GH'KCN-- ",!24!2 LLE!26!:;!GH!2  4 www'7+=A{((*S069:FL{((*V3c9FBCtc%--000&3u}}+="== ( +.F;$$&'3AF;$$&/0 ( U 
!	 $ YZ4 
!	 s5   1E<]-A]2F]D3])]
0C]
]
])T)F)__doc__
concurrentr   loggingdora.logr   r)   rt   rJ   r*   torchr   applyr   r\   r   r    r	   utilsr
   	getLogger__name__rU   r   r0   r    r   r   <module>r      sO   
          ,  $ 
		8	$ "*kr   