a
    yfF                     @   s   d dl 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
 d dlmZmZmZ d dlmZmZmZ d	d
 Zdd Zdd Zdd Zdd ZdS )    N)mock)MODEL)YOLO)get_cfg)Exporter)classifydetectsegment)ASSETSDEFAULT_CFGWEIGHTS_DIRc                  G   s   t d dS )zETest function callback for evaluating YOLO model performance metrics.zcallback test passedN)printargs r   M/var/www/html/django/DPS/env/lib/python3.9/site-packages/tests/test_engine.py	test_func   s    r   c                  C   sH   t  } | dt t| jd v s(J d| tdjd}t|t dS )zTTests the model exporting function by adding a callback and asserting its execution.Zon_export_startcallback test failedyolo11n.yamlmodelN)r   add_callbackr   	callbacksr   r   r
   )Zexporterfr   r   r   test_export   s
    r   c               
   C   s~  dddddd} t t}d|_d|_tj| d}|dt t|jd v sRJ d	|	  tj
|d
}|dt t|jd v sJ d	||jd tjdddgid}|dt t|jd v sJ d	tjtdg , |ttd}t|sJ dW d   n1 s0    Y  |j| d< tj| d}z|	  W n6 typ } ztd|  W Y d}~dS d}~0 0 td dS )zNTest YOLO object detection training, validation, and prediction functionality.z
coco8.yamlr          Fdatar   imgszZ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   ZDetectionTrainerr   r   r   trainZDetectionValidatorbestZDetectionPredictorr   patchobjectsysr
   r   lenlast	Exceptionr   r"   cfgtrainervalpredresulter   r   r   test_detect   s4    0
r=   c               
   C   sP  dddddd} t t}d|_d|_tj| d}|dt t|jd v sRJ d	|	  tj
|d
}|dt t|jd v sJ d	||jd tjdddgid}|dt t|jd v sJ d	|ttd d}t|sJ d|j| d< tj| d}z|	  W n6 tyB } ztd|  W Y d}~dS d}~0 0 td dS )zZTests image segmentation training, validation, and prediction pipelines using YOLO models.zcoco8-seg.yamlzyolo11n-seg.yamlr   r   Fr   r!   r#   r   r   r$   r   r   r%   r&   zyolo11n-seg.ptr(   r*   r+   r,   Nr-   )r   r   r   r   r	   ZSegmentationTrainerr   r   r   r.   ZSegmentationValidatorr/   ZSegmentationPredictorr
   r   r3   r4   r5   r   r6   r   r   r   test_segmentC   s2    
r>   c                  C   s   dddddd} t t}d|_d|_tj| d}|dt t|jd v sRJ d	|	  tj
|d
}|dt t|jd v sJ d	||jd tjdddgid}|dt t|jd v sJ d	|t|jd}t|sJ ddS )zPTest image classification including training, validation, and prediction phases.Z
imagenet10zyolo11n-cls.yamlr   r   Fr   r!   r#   r   r   r$   r   r   r%   r&   r(   r*   N)r   r   r   r   r   ZClassificationTrainerr   r   r   r.   ZClassificationValidatorr/   ZClassificationPredictorr
   r3   )r"   r7   r8   r9   r:   r;   r   r   r   test_classifyj   s"    r?   )r2   Zunittestr   testsr   Zultralyticsr   Zultralytics.cfgr   Zultralytics.engine.exporterr   Zultralytics.models.yolor   r   r	   Zultralytics.utilsr
   r   r   r   r   r=   r>   r?   r   r   r   r   <module>   s   	''