
    jb#                     (   d dl Z d dlmZ d dlmZ d dlZd dlZd dlmZm	Z	m
Z
 d dlmZ d dlmZ d dlmZ d dlmZ d d	lmZmZmZmZmZmZ d d
lmZ d dlmZmZmZ d Z d Z!ej"        #                    dej$        ej%        ej&        ddefej'        ej(        ej)        ddedz  fej*        ej+        ej,        dddfej-        ej.        ej/        dddfej0        ej1        ej2        dddfej3        ej4        ej5        dddfg          d             Z6ej"        #                    de
          d             Z7d Z8ej"        #                    d i d!fd"d!id!fd"d#id#fd"eid!fg          d$             Z9ej"        #                    d%d&d'ed!fg          d(             Z:dS ))    NSimpleNamespace)mock)MODELSOURCETASK_MODEL_DATA)YOLO)get_cfg)Exporter)BaseTrainer)classifydetectobbposesegmentsemantic)load_checkpoint)ASSETSDEFAULT_CFGWEIGHTS_DIRc                  $    t          d           dS )zFTest function used as a callback stub to verify callback registration.zcallback test passedN)print)argskwargss     T/var/www/html/Carbon-Document/venv/lib/python3.11/site-packages/tests/test_engine.py	test_funcr      s    	
 !!!!!    c                  
   t                      } |                     dt                     t          | j        d         v s
J d             | t	          d          j                  } t	          |          t                     dS )zTTest model exporting functionality by adding a callback and verifying its execution.on_export_startz'on_export_start callback not registeredyolo26n.yamlmodelN)r   add_callbackr   	callbacksr	   r"   r   )exporterfs     r   test_exportr'      sy    zzH+Y777*+<====?h===tN++1222ADGGFOOOOOr   z:trainer_cls,validator_cls,predictor_cls,data,model,weights
coco8.yamlr    zcoco8-seg.yamlzyolo26n-seg.yamlzyolo26n-seg.pt
imagenet10zyolo26n-cls.yamlz
dota8.yamlzyolo26n-obb.yamlzcoco8-pose.yamlzyolo26n-pose.yamlzcityscapes8.yamlzyolo26n-sem.yamlc                 n   ||dddddd} | |          }|                     dt                     t          |j        d         v s
J d            |                                 t	          t
                    }||_        d|_         ||          }	|	                     d	t                     t          |	j        d	         v s
J d
             |	|j                    |dddgi          }
|
                     dt                     t          |
j        d         v s
J d            |r|n|j        }|dk    rrt          j
                            t          dg           5   |
t          |          }t          |          dk    sJ d|             	 ddd           n# 1 swxY w Y   n2 |
t          |          }t          |          dk    sJ d|             t          j        t"                    5   | i |d|j        i                                           ddd           dS # 1 swxY w Y   dS )zATest YOLO training, validation, and prediction for various tasks.       F)datar"   imgszepochssave
mask_ratiooverlap_mask	overrideson_train_startz&on_train_start callback not registered)r   on_val_startz$on_val_start callback not registeredr!   r.   @   on_predict_startz(on_predict_start callback not registeredr    argv)sourcer"   r   z"Predictor returned no results for Nresume)r#   r   r$   trainr
   r   r-   r.   bestr   patchobjectsysr   lenpytestraisesAssertionErrorlast)trainer_clsvalidator_clspredictor_clsr-   r"   weightsr4   trainercfgvalpred
model_pathresults                r   	test_taskrP   "   s   T  I kI...G)9555)*:;;;;=e;;;MMOOO +

CCHCI
-S
!
!
!C^Y///n55557]555Cgl =Gb"X#6777D()444'9::::<f::: $5JZsFB// 	Q 	QTz:::Fv;;???$P$P$P????	Q 	Q 	Q 	Q 	Q 	Q 	Q 	Q 	Q 	Q 	Q 	Q 	Q 	Q 	Q V:6666{{Q LU L L 
~	&	& M MCCHglCCDDDJJLLLM M M M M M M M M M M M M M M M M Ms$   3FF #F 3*H**H.1H.ztask,weight,datac           
         |dddd|| ddd	}d }d }t          |          }|                    d	|           |                    d
|            |j        di | |j        j        }t          |          \  }	}
|
d         dk    s
J d            t          |          } |j        dddi| |j        j        |j        j        cxk    rdk    sn J d            dS )z4Test training resumes from an incomplete checkpoint.   TFr   r+   )	r-   r/   r0   plotsworkersprojectnamer.   exist_okc                 .    | j         dk    r	d| _        d S d S )Nr   T)epochstoprJ   s    r   stop_after_first_epochz6test_resume_incomplete.<locals>.stop_after_first_epoch   s"    =AGLLL r   c                     d | _         d S )Nc                      d S N r`   r   r   <lambda>zDtest_resume_incomplete.<locals>.disable_final_eval.<locals>.<lambda>   s    T r   )
final_evalr[   s    r   disable_final_evalz2test_resume_incomplete.<locals>.disable_final_eval   s    )\r   r5   on_train_epoch_endrY   zcheckpoint should be resumabler;   r,   zresume test failedNr`   )r	   r#   r<   rJ   rE   r   start_epochrY   )taskweightr-   tmp_path
train_argsr\   rc   r"   	last_path_ckptresume_models               r   test_resume_incompletern   y   sB    
 
J     * * * LLE	');<<<	+-CDDDEK*"Ii((GAt=A? 	??LL11d1j111+|/C/INNNNQNNNNNPdNNNNNr   c                      dgfd} ddddd}t          j        |          }|                    d	|            |                                 d
         s
J d            dS )z5Test NaN loss detection and recovery during training.Fc                     | j         dk    rE| j        @d         s:| xj        t          j        t	          d                    z  c_        dd<   dS dS dS dS )zHInject NaN into loss during batch processing to test recovery mechanism.r,   Nr   nanT)rY   tlosstorchtensorfloat)rJ   nan_injecteds    r   
inject_nanz%test_nan_recovery.<locals>.inject_nan   sb    =A'-";LQRO";MMU\%,,777MM"LOOO ";";";";r   r(   r    r+      )r-   r"   r.   r/   r3   on_train_batch_endr   zNaN injection failedN)r   DetectionTrainerr#   r<   )rw   r4   rJ   rv   s      @r   test_nan_recoveryr{      s    7L# # # # # &WXYYI%	:::G-z:::MMOOO?22222?22r   zkwargs,uses_weightsT
pretrainedFc                    t          d          }ddi|_        d|_        d|j        d<   d|j        d<   |j        i  G fd	d
          |                     dd            |                     |dfd           |                     dfd            |j        dddd| d         j        u s
J d            d         j        k    sJ dd          dj                     d         |rndu s
J d            dS )zWTest training reuses loaded checkpoint config while respecting the pretrained argument.r    
checkpointTz/tmp/fake.ptzul://glenn-jocher/m2/exp-14r"   Fr|   c                   0    e Zd Zd fd	Zd fd	Zd ZdS )>test_train_reuses_loaded_checkpoint_model.<locals>.FakeTrainerNc                     || _         || _        d | _        t          d           | _        t
          j        dz  | _        t
          | _        | d<   d S )N)metricsznonexistent-best.ptrJ   )	r4   r$   r"   r   	validatorr   parentr=   rE   )selfr4   
_callbackscaptureds      r   __init__zGtest_train_reuses_loaded_checkpoint_model.<locals>.FakeTrainer.__init__   sP    &DN'DNDJ,T:::DN'<<DIDI"&HYr   Tc                     |d<   |d<   S NrK   rI   r`   )r   rK   rI   verboser   original_models       r   	get_modelzHtest_train_reuses_loaded_checkpoint_model.<locals>.FakeTrainer.get_model   s    !HUO")HY!!r   c                     d S r_   r`   )r   s    r   r<   zDtest_train_reuses_loaded_checkpoint_model.<locals>.FakeTrainer.train   s    4r   )NNNNT)__name__
__module____qualname__r   r   r<   )r   r   s   r   FakeTrainerr      se        	' 	' 	' 	' 	' 	'	" 	" 	" 	" 	" 	" 	"
	 	 	 	 	r   r   z:ultralytics.engine.model.checks.check_pip_update_availablec                      d S r_   r`   r`   r   r   ra   z;test_train_reuses_loaded_checkpoint_model.<locals>.<lambda>   s    ^b r   _smart_loadc                     S r_   r`   )keyr   s    r   ra   z;test_train_reuses_loaded_checkpoint_model.<locals>.<lambda>   s    + r   z(ultralytics.engine.model.load_checkpointc                     ddifS Nr~   Tr`   )pathr   s    r   ra   z;test_train_reuses_loaded_checkpoint_model.<locals>.<lambda>   s    n|T&:; r   r(   r,   )r-   r/   rJ   z%Trainer model does not match originalrK   zConfig mismatch: z != rI   NUnexpected weights loadedr`   )r	   rl   	ckpt_pathr4   r"   setattrr<   yaml)monkeypatchr   uses_weightsr"   r   r   r   s       @@@r   )test_train_reuses_loaded_checkpoint_modelr      s      E%EJ$EO<EOG$)EOL![NH          $ TVbVbccc}.E.E.E.EFFF2;;;;  
 EK6\!66v666I$6668_666E?n11113qxPU3q3q\j\o3q3q111I\#K>>tLLLNiLLLLLr   zpretrained,uses_weights)TT)FFc                    i t          ddi          t                              t                    }d|_        t          |          |_        d|_        dfd		}||_        |                     d
fd           |	                                 d         j
        k    s
J d            d         |rndu s
J d            dS )zOTest .pt models use checkpoint config while respecting the pretrained argument.ncP   )r   z
yolo26n.pt)r|   FNTc                 4    | d<   |d<   t                      S r   r   )rK   rI   r   r   s      r   fake_get_modelzNtest_setup_model_respects_pretrained_arg_for_pt_models.<locals>.fake_get_model   s"    %   r   z*ultralytics.engine.trainer.load_checkpointc                     ddifS r   r`   )r   checkpoint_models    r   ra   zHtest_setup_model_respects_pretrained_arg_for_pt_models.<locals>.<lambda>   s    DTWceiVjCk r   rK   zCheckpoint config was not usedrI   r   r   )r   r?   __new__r   r"   r   r;   r   r   setup_modelr   )r   r|   r   rJ   r   r   r   s        @@r   6test_setup_model_respects_pretrained_arg_for_pt_modelsr      s    H&T2J777nn[))G GM"j999GLGN! ! ! ! ! !
 'G46k6k6k6k   E?.33335U333I|#M#3#3NNNPkNNNNNr   );r@   typesr   unittestr   rB   rs   testsr   r   r   ultralyticsr	   ultralytics.cfgr
   ultralytics.engine.exporterr   ultralytics.engine.trainerr   ultralytics.models.yolor   r   r   r   r   r   ultralytics.nn.tasksr   ultralytics.utilsr   r   r   r   r'   markparametrizerz   DetectionValidatorDetectionPredictorSegmentationTrainerSegmentationValidatorSegmentationPredictorClassificationTrainerClassificationValidatorClassificationPredictor
OBBTrainerOBBValidatorOBBPredictorPoseTrainerPoseValidatorPosePredictorSemanticSegmentationTrainerSemanticSegmentationValidatorSemanticSegmentationPredictorrP   rn   r{   r   r   r`   r   r   <module>r      s$   


 ! ! ! ! ! !         0 0 0 0 0 0 0 0 0 0       # # # # # # 0 0 0 0 0 0 2 2 2 2 2 2 R R R R R R R R R R R R R R R R 0 0 0 0 0 0 > > > > > > > > > >" " "
   @ #%%	
 '))**	
 *,,	
 
)3+;\K]_cd		4-t/ACTVikop022	
7#& &N-M -MO& &N-M` +_== e  e >= eF3 3 3" $Z<&-u0Eu/MQ]_dPegkOlm 'j 'j	 'jT 2\>TY[_S`4abbl l cbl l lr   