
    hb                         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JrJr  S SKJrJrJr  S	 rS
 rS rS rS rS rg)    N)mock)MODEL)YOLO)get_cfg)Exporter)classifydetectsegment)ASSETSDEFAULT_CFGWEIGHTS_DIRc                      [        S5        g)zETest function callback for evaluating YOLO model performance metrics.zcallback test passedN)printargss    K/home/james-whalen/.local/lib/python3.13/site-packages/tests/test_engine.py	test_funcr      s    	
 !    c                      [        5       n U R                  S[        5        [        U R                  S   ;   d   S5       eU " [	        S5      R
                  S9n[	        U5      " [        5        g)zTTest model exporting functionality by adding a callback and verifying its execution.on_export_startcallback test failedyolo11n.yamlmodelN)r   add_callbackr   	callbacksr   r   r   )exporterfs     r   test_exportr      sZ    zH+Y7**+<==U?UU=tN+112AGFOr   c                     SSSSSS.n [        [        5      nSUl        SUl        [        R
                  " U S9nUR                  S[        5        [        UR                  S   ;   d   S	5       eUR                  5         [        R                  " US
9nUR                  S[        5        [        UR                  S   ;   d   S	5       eU" UR                  S9  [        R                  " SSS/0S9nUR                  S[        5        [        UR                  S   ;   d   S	5       e[        R                  R                  [         S/ 5         U" ["        [$        S9n['        U5      (       d   S5       e SSS5        UR(                  U S'   [        R
                  " U S9n UR                  5         [+        S5      e! , (       d  f       NM= f! [*         a  n[-        SU 35         SnAgSnAff = f)zNTest YOLO object detection training, validation, and prediction functionality.
coco8.yamlr          Fdatar   imgszepochssave	overrideson_train_startr   r   on_val_startr   r&   @   on_predict_startargvsourcer   predictor test failedNresumeExpected exception caught: Resume test failed!)r   r   r%   r&   r	   DetectionTrainerr   r   r   trainDetectionValidatorbestDetectionPredictorr   patchobjectsysr   r   lenlast	Exceptionr   r*   cfgtrainervalpredresultes          r   test_detectrH      s   %WXbghI
+
CCHCI %%	:G)95))*:;;S=SS;MMO 
#
#
-C^Y/n55M7MM5gll $$"b/BCD()4'9::R<RR:			3	+V516{{333{ 
,
 ",,Ih%%	:G
 )
** 
,	+  +A3/0s$   'F6G 6
G
G)G$$G)c                  b   SSSSSSSS.n [        [        5      nSUl        SUl        [        R
                  " U S9nUR                  S[        5        [        UR                  S   ;   d   S	5       eUR                  5         [        R                  " US
9nUR                  S[        5        [        UR                  S   ;   d   S	5       eU" UR                  S9  [        R                  " SSS/0S9nUR                  S[        5        [        UR                  S   ;   d   S	5       eU" [        [        S-  S9n[        U5      (       d   S5       eUR                   U S'   [        R
                  " U S9n UR                  5         [#        S5      e! ["         a  n[%        SU 35         SnAgSnAff = f)zYTest image segmentation training, validation, and prediction pipelines using YOLO models.zcoco8-seg.yamlzyolo11n-seg.yamlr"   r#   F)r%   r   r&   r'   r(   
mask_ratiooverlap_maskr)   r+   r   r   r,   r   r&   r-   r.   zyolo11n-seg.ptr0   r2   r3   r4   Nr5   )r   r   r%   r&   r
   SegmentationTrainerr   r   r   r7   SegmentationValidatorr9   SegmentationPredictorr   r   r>   r?   r@   r   rA   s          r   test_segmentrO   F   s    !#I +
CCHCI ))I>G)95))*:;;S=SS;MMO 
'
'S
1C^Y/n55M7MM5gll ((Gb"X3FGD()4'9::R<RR:{5E'EFFv;;///; ",,Ih))I>G
 )
**	  +A3/0s   1F 
F.F))F.c                     SSSSSS.n [        [        5      nSUl        SUl        [        R
                  " U S9nUR                  S[        5        [        UR                  S   ;   d   S	5       eUR                  5         [        R                  " US
9nUR                  S[        5        [        UR                  S   ;   d   S	5       eU" UR                  S9  [        R                  " SSS/0S9nUR                  S[        5        [        UR                  S   ;   d   S	5       eU" [        UR                  S9n[        U5      (       d   S5       eg)zPTest image classification including training, validation, and prediction phases.
imagenet10zyolo11n-cls.yamlr"   r#   Fr$   r)   r+   r   r   r,   r   r&   r-   r.   r0   r2   N)r   r   r%   r&   r   ClassificationTrainerr   r   r   r7   ClassificationValidatorr9   ClassificationPredictorr   r>   )r*   rB   rC   rD   rE   rF   s         r   test_classifyrU   t   s2   %0BR[\fklI
+
CCHCI ,,yAG)95))*:;;S=SS;MMO 
*
*
4C^Y/n55M7MM5gll ++wR6IJD()4'9::R<RR:w||4Fv;;///;r   c                     ^ S/mU4S jn SSSSS.n[         R                  " US9nUR                  S	U 5        UR                  5         TS
   (       d   S5       eg)z5Test NaN loss detection and recovery during training.Fc                    > U R                   S:X  aQ  U R                  bC  TS   (       d8  U =R                  [        R                  " [	        S5      5      -  sl        STS'   gggg)zHInject NaN into loss during batch processing to test recovery mechanism.r#   Nr   nanT)epochtlosstorchtensorfloat)rC   nan_injecteds    r   
inject_nan%test_nan_recovery.<locals>.inject_nan   sP    ==A'--";LQROMMU\\%,77M"LO ET";r   r!   r   r"      )r%   r   r&   r'   r)   on_train_batch_endr   zNaN injection failedN)r	   r6   r   r7   )r_   r*   rC   r^   s      @r   test_nan_recoveryrc      s\    7L# &WXYI%%	:G-z:MMO?222?r   )r=   unittestr   r[   testsr   ultralyticsr   ultralytics.cfgr   ultralytics.engine.exporterr   ultralytics.models.yolor   r	   r
   ultralytics.utilsr   r   r   r   r   rH   rO   rU   rc    r   r   <module>rl      sH         # 0 = = > >"
%+P++\063r   