a
    yf                     @   s  d Z ddlZddlZddlZddlZddlZddlZddlZddlm	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 ddlmZ ddlmZ dd	lmZmZ dd
lmZmZ ddlmZmZm Z  ddl!m"Z"m#Z#m$Z$ ddl%m&Z&m'Z'm(Z(m)Z)m*Z*m+Z+m,Z,m-Z-m.Z.m/Z/m0Z0m1Z1m2Z2m3Z3 ddl4m5Z5m6Z6m7Z7m8Z8 ddl9m:Z:m;Z;m<Z< ddl=m>Z>m?Z? ddl@mAZA ddlBmCZCmDZDmEZEmFZF dd ZGdd ZHdd ZIG dd dZJG dd dejKjLZMdS )a	  
Export a YOLOv8 PyTorch model to other formats. TensorFlow exports authored by https://github.com/zldrobit.

Format                  | `format=argument`         | Model
---                     | ---                       | ---
PyTorch                 | -                         | yolov8n.pt
TorchScript             | `torchscript`             | yolov8n.torchscript
ONNX                    | `onnx`                    | yolov8n.onnx
OpenVINO                | `openvino`                | yolov8n_openvino_model/
TensorRT                | `engine`                  | yolov8n.engine
CoreML                  | `coreml`                  | yolov8n.mlpackage
TensorFlow SavedModel   | `saved_model`             | yolov8n_saved_model/
TensorFlow GraphDef     | `pb`                      | yolov8n.pb
TensorFlow Lite         | `tflite`                  | yolov8n.tflite
TensorFlow Edge TPU     | `edgetpu`                 | yolov8n_edgetpu.tflite
TensorFlow.js           | `tfjs`                    | yolov8n_web_model/
PaddlePaddle            | `paddle`                  | yolov8n_paddle_model/
NCNN                    | `ncnn`                    | yolov8n_ncnn_model/

Requirements:
    $ pip install "ultralytics[export]"

Python:
    from ultralytics import YOLO
    model = YOLO('yolov8n.pt')
    results = model.export(format='onnx')

CLI:
    $ yolo mode=export model=yolov8n.pt format=onnx

Inference:
    $ yolo predict model=yolov8n.pt                 # PyTorch
                         yolov8n.torchscript        # TorchScript
                         yolov8n.onnx               # ONNX Runtime or OpenCV DNN with dnn=True
                         yolov8n_openvino_model     # OpenVINO
                         yolov8n.engine             # TensorRT
                         yolov8n.mlpackage          # CoreML (macOS-only)
                         yolov8n_saved_model        # TensorFlow SavedModel
                         yolov8n.pb                 # TensorFlow GraphDef
                         yolov8n.tflite             # TensorFlow Lite
                         yolov8n_edgetpu.tflite     # TensorFlow Edge TPU
                         yolov8n_paddle_model       # PaddlePaddle
                         yolov8n_ncnn_model         # NCNN

TensorFlow.js:
    $ cd .. && git clone https://github.com/zldrobit/tfjs-yolov5-example.git && cd tfjs-yolov5-example
    $ npm install
    $ ln -s ../../yolov5/yolov8n_web_model public/yolov8n_web_model
    $ npm start
    N)deepcopy)datetime)Path)	TASK2DATAget_cfg)build_dataloader)YOLODataset)check_cls_datasetcheck_det_dataset)check_class_namesdefault_class_names)C2fDetectRTDETRDecoder)DetectionModelSegmentationModel
WorldModel)ARM64DEFAULT_CFG	IS_JETSONLINUXLOGGERMACOSPYTHON_VERSIONROOTWINDOWS__version__	callbackscolorstrget_default_args	yaml_save)check_imgszcheck_is_path_safecheck_requirementscheck_version)attempt_download_assetget_github_assetssafe_download)	file_sizespaces_in_path)Profile)
TORCH_1_13get_latest_opsetselect_devicesmart_inference_modec                  C   sh   g dg dg dg dg dg dg dg dg d	g d
g dg dg dg} t tg dt|  S )z Ultralytics YOLO export formats.)ZPyTorch-z.ptTT)ZTorchScriptZtorchscript.torchscriptTT)ZONNXonnx.onnxTT)ZOpenVINOopenvino_openvino_modelTF)ZTensorRTengine.engineFT)ZCoreMLcoreml
.mlpackageTF)zTensorFlow SavedModelsaved_model_saved_modelTT)zTensorFlow GraphDefpb.pbTT)zTensorFlow Litetflite.tfliteTF)zTensorFlow Edge TPUedgetpu_edgetpu.tfliteTF)zTensorFlow.jstfjs
_web_modelTF)ZPaddlePaddlepaddle_paddle_modelTT)ZNCNNncnn_ncnn_modelTT)ZFormatArgumentZSuffixZCPUZGPU)dictzipx rL   W/var/www/html/django/DPS/env/lib/python3.9/site-packages/ultralytics/engine/exporter.pyexport_formatsa   s    rN   c                 C   sP   g g  }}| j D ]}||j ||j qtdd tt|t| D S )z,TensorFlow GraphDef model output node names.c                 s   s"   | ]}| d s| dV  qdS )ZNoOpz:0N)
startswith.0rK   rL   rL   rM   	<genexpr>{       zgd_outputs.<locals>.<genexpr>)nodeappendnameextendinputsortedlistset)gd	name_listZ
input_listrT   rL   rL   rM   
gd_outputsu   s
    

r^   c                    s   t   fdd}|S )z*YOLOv8 export decorator, i.e. @try_export.c               
      s    d }znt  "}| i |\}}W d   n1 s80    Y  t| d|jdd| dt|dd ||fW S  ty } z.t| d|jdd	|  |W Y d}~n
d}~0 0 dS )
zExport a model.prefixNu    export success ✅ .1fzs, saved as 'z' ( MB)u    export failure ❌ zs: )r*   r   infotr(   	Exceptionerror)argskwargsr_   dtfmodeleZ
inner_args
inner_funcrL   rM   
outer_func   s    0,
ztry_export.<locals>.outer_func)r   )rm   rn   rL   rl   rM   
try_export~   s    ro   c                   @   sh  e Zd ZdZeddfddZe d9edddZd:d	d
Z	e
edfddZe
edfddZe
edfddZe
edfddZe
edfddZe
edfddZe
edfddZe
ed fd!d"Ze
ed#fd$d%Ze
ed&fd'd(Ze
ded)fd*d+Ze
ed,fd-d.Zd/d0 Zded1fd2d3Zed4d5d6Zed4d7d8ZdS );Exporterz
    A class for exporting a model.

    Attributes:
        args (SimpleNamespace): Configuration for the exporter.
        callbacks (list, optional): List of callback functions. Defaults to None.
    Nc                 C   sB   t ||| _| jj dv r&dtjd< |p0t | _t|  dS )aH  
        Initializes the Exporter class.

        Args:
            cfg (str, optional): Path to a configuration file. Defaults to DEFAULT_CFG.
            overrides (dict, optional): Configuration overrides. Defaults to None.
            _callbacks (dict, optional): Dictionary of callback functions. Defaults to None.
        >   mlmodelr7   pythonZ&PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATIONN)	r   rf   formatlowerosenvironr   Zget_default_callbacksZadd_integration_callbacks)selfcfgZ	overrides
_callbacksrL   rL   rM   __init__   s
    	
zExporter.__init__returnc           #         s4  |  d t }| jj   dv r*d  dv r6d tt d dd } |vrd	dl}|j |dd
d}|rt	
d  d|d	  d |d	  ntd  d|  fdd|D }t|dkrtd  d| |\}}}	}
}}}}}}}}t|||||f} dkr6| jjdu r6t	
d d| j_t| jjdu rJdn| jj| _t|dsjt |_t|j|_| jjr| jjrt	
d d| j_| jjr|r| jjdkrt	
d d| j_| jjrJ dt| jj|jdd| _| jjr|
rd| j_| jjr:|r$J d| jjdks:J d|rptsPtd n | jjdkrpt	
d! d| j_t|t rt	
d" | jjr| jj!st"j!pt#t$|d#d$ | j_!t	
d%| jj! d& t%j&| jjd'g| jR  '| j}t(t$|d(dpt$|d)dp|j)*d)d*}|j+d+v r2t(|j,}t-|'| j}|. D ]}d|_/qJ|0  |1  |2 }|3 D ]V}t|t4t5fr| jj|_d|_6| jj|_| jj7|_7nt|t8rx|sx|j9|_:qxd}t;dD ]}||}q| jjr|r| jjdkr| |  }}t<j=d,t%j>j?d- t<j=d,t@d- t<j=d,tAd- || _B|| _C|| _Dt|t%jErtt|jFntd.d/ |D | _Gt(| jCj)*d)| jDjHId0d1| _Jt|d2rt|jtKr|jd3 nd*}d4| jJ d5|rd6| nd* }|d7tLM N tOd8d9tPtQ|j|jR| jj| j|jd:| _S|jRd;krL|jCd< jT| jSd=< t	Ud>tVd? d@| dAt|jF dB| jG dCtW|dDdE d*gtX| }|s|r| Y \|d	< }|
r| Z \|d< }|r| [ \|d< }|	r| \ \|d'< }|r
| ] \|dF< }|r| j j|O  _| ^ \|dG< }|s<|rP| j_|dH\|dI< }|rr| j`|d| jjadJ\|dK< }|r| jbt(|dG | jDjH dL dM\|dN< }|r| c \|dO< }|r| d \|dP< }|r| e \|dQ< }dRd |D }t|r&tft(|d< }| jd	 | jd k}|r*d*ndS| j dTtQ| j dU}|rT| jd	 ntf| jdd< IdVd*} |jRdWkr dXkrdY| nd*}!| jjrdZn| jjrd[nd*}"t	Ud\t | dDd]tVd^|jgh  d_|jR d`| da|  dV|" dV|! db|jR d`| da|  dc| dV|" dV| dd |  de |S )fz<Returns list of exported files/dirs after running callbacks.Zon_export_start>   trttensorrtr5   >   Z	mlpackageapple	mlprogramrq   r7   Ziosr7   rG      Nr   g333333?)ncutoffu&   WARNING ⚠️ Invalid export format='z', updating to format=''zInvalid export format='z'. Valid formats are c                    s   g | ]}| kqS rL   rL   rP   fmtrL   rM   
<listcomp>   rS   z%Exporter.__call__.<locals>.<listcomp>uM   WARNING ⚠️ TensorRT requires GPU export, automatically assigning device=00cpunamesuR   WARNING ⚠️ half=True and int8=True are mutually exclusive, setting half=False.FuK   WARNING ⚠️ half=True only compatible with GPU export, i.e. use device=0z>half=True not compatible with dynamic=True, i.e. use only one.   )strideZmin_dimTzHoptimize=True not compatible with format='ncnn', i.e. use optimize=FalsezEoptimize=True not compatible with cuda devices, i.e. use device='cpu'zSEdge TPU export only supported on Linux. See https://coral.ai/docs/edgetpu/compileruF   WARNING ⚠️ Edge TPU export requires batch size 1, setting batch=1.u  WARNING ⚠️ YOLOWorld (original version) export is not supported to any format.
WARNING ⚠️ YOLOWorldv2 models (i.e. 'yolov8s-worldv2.pt') only support export to (torchscript, onnx, openvino, engine, coreml) formats. See https://docs.ultralytics.com/models/yolo-world for details.taskdetectu^   WARNING ⚠️ INT8 export requires a missing 'data' arg for calibration. Using default 'data='.   Zpt_pathZ	yaml_file >   z.ymlz.yamlignore)categoryc                 s   s(   | ] }t t|tjr|jng V  qd S N)tuple
isinstancetorchTensorshaperP   rL   rL   rM   rR     rS   z$Exporter.__call__.<locals>.<genexpr>ZyoloZYOLOrf   datazUltralytics z model ztrained on ZUltralyticsz2AGPL-3.0 License (https://ultralytics.com/license)zhttps://docs.ultralytics.com)descriptionauthordateversionlicensedocsr   r   batchimgszr   Zpose	kpt_shape
zPyTorch:z starting from 'z' with input shape z BCHW and output shape(s) z (r`   ra         keras_model   )r   nmsagnostic_nms   z_full_integer_quant.tflite)tflite_model   	   
      c                 S   s   g | ]}|rt |qS rL   strrP   rL   rL   rM   r   Q  rS   u>   WARNING ⚠️ non-PyTorch val requires square images, 'imgsz=z#' will not work. Use export 'imgsz=z' if val is required. segmentr;   zdata=int8halfz
Export complete (zs)
Results saved to boldz$
Predict:         yolo predict task=z model=z imgsz=z 
Validate:        yolo val task=z data=z$
Visualize:       https://netron.appZon_export_end)irun_callbackstimerf   rs   rt   r   rN   difflibget_close_matchesr   warning
ValueErrorsumanydevicer-   hasattrr   r   r   r   r   typedynamicr!   r   r   optimizer   SystemErrorr   r   r   r   r   r   getattrr   Zzerostor   yamlgetsuffixrV   r   
parametersZrequires_gradevalfloatZfusemodulesr   r   exportZmax_detr   Zforward_splitforwardrangewarningsfilterwarningsjitZTracerWarningUserWarningDeprecationWarningimrj   filer   r   output_shapestemreplacepretty_namerH   r   now	isoformatr   intmaxr   metadatar   rb   r   r(   lenexport_torchscriptexport_engineexport_onnxexport_openvinoexport_coremlexport_saved_model	export_pbexport_tfliter   export_edgetpuexport_tfjsexport_paddleexport_ncnnr   parentresolve)#rw   rj   rc   Zfmtsr   matchesflagsr   r1   xmlr5   r7   r9   r;   r=   r?   rA   rC   rE   Zis_tf_formatr   r   pmy_r   r   ri   r   Zsquaresr   Zpredict_dataqrL   r   rM   __call__   sp   


 




"(




"("
*

*$ 
zExporter.__call__r   c                 C   s   t | d| jj d | jjdkr*tnt| jj}| jj| jj	dkrLdnd }t
|| jjp`d || jj| jd d	|d
}t|}|dk rt | d| d t||ddS )zFBuild and return a dataloader suitable for calibration of INT8 models.z/ collecting INT8 calibration images from 'data=r   classifyr5   r   r   valr   F)r   r   r   ZaugmentZ
batch_sizei,  uD    WARNING ⚠️ >300 images recommended for INT8 calibration, found z images.)r   workers)r   rb   rf   r   rj   r   r	   r
   r   rs   r   splitr   r   r   r   )rw   r_   r   r   datasetr   rL   rL   rM   get_int8_calibration_dataloaderi  s    z(Exporter.get_int8_calibration_dataloaderzTorchScript:c                 C   s   t d| dtj d | jd}tjj| j| j	dd}dt
| ji}| jjrt | d d	d
lm} ||jt||d n|jt||d |dfS )z YOLOv8 TorchScript model export.r   z starting export with torch ...r0   Fstrictz
config.txtz optimizing for mobile...r   )optimize_for_mobile)Z_extra_filesN)r   rb   r   r   r   with_suffixr   tracerj   r   jsondumpsr   rf   r   Ztorch.utils.mobile_optimizerr  Z_save_for_lite_interpreterr   save)rw   r_   ri   tsZextra_filesr  rL   rL   rM   r   |  s    zExporter.export_torchscriptzONNX:c                 C   s  dg}| j jr,|ddtj r"dnd g7 }t| ddl}| j jpHt }t	
d| d	|j d
| d t| jd}t| jtrddgndg}| j j}|rdddddi}t| jtrddd|d< dddd|d< nt| jtrddd|d< tjj|r| j n| j|r | j n| j|d|ddg||p:dd	 ||}| j jrz.ddl}	t	
| d|	j d |	|}W n: ty }
 z t	| d|
  W Y d}
~
n
d}
~
0 0 | j D ]&\}}|j }|t| |_ |_!q|"|| ||fS )zYOLOv8 ONNX export.onnx>=1.12.0zonnxslim==0.1.34onnxruntimez-gpur   r   Nr   z starting export with onnx z opset r  r2   Zoutput0output1imagesr   heightwidth)r   r   r   anchors)r   r   Zmask_heightZ
mask_widthFT)verboseopset_versionZdo_constant_foldingZinput_namesoutput_namesZdynamic_axesz slimming with onnxslim z simplifier failure: )#rf   simplifyr   cudais_availabler#   r1   Zopsetr,   r   rb   r   r   r   r  r   rj   r   r   r   r   r   r   loadonnxslimslimrd   r   r   itemsZmetadata_propsaddkeyvaluer  )rw   r_   requirementsr1   r  ri   r  r   Z
model_onnxr  rk   kvmetarL   rL   rM   r     sR     

*
zExporter.export_onnxz	OpenVINO:c                    s  t dtrdnd  ddl td| d j d tsNJ d	tj d
 jj	j
jrbdnjjgjd} fdd}j
jrtjjjdtj }tt|jdj }t d ddl}tjddd}d}tj	j	d trhdtj	 d d  ddd }	|j!d|	 dd|	 dd|	 dd|	 dd|	 dgdgd}|j"||#$|||j%j&|d}
||
| |dfS tjjjd tj }tt|jdj }||| |dfS )!zYOLOv8 OpenVINO export.r3   z
<=2024.0.0z
>=2024.0.0r   Nr   z starting export with openvino r  z2OpenVINO export requires torch>=1.13.0 but torch==z is installed)rX   Zexample_inputc                    s   |  dddg |  dddg |  dddg |  dgdd	g |  jjdd
g |  dd jj D ddg jjdkr|  dddg  jj| |jj	d t
t|jd j dS )z.Set RT info, serialize and save metadata YAML.ZYOLOv8Z
model_infoZ
model_typeTZreverse_input_channelsr   Z	pad_value     o@Zscale_valuesZiou_thresholdc                 S   s   g | ]}| d dqS )r   r   )r   )rQ   r#  rL   rL   rM   r     rS   z?Exporter.export_openvino.<locals>.serialize.<locals>.<listcomp>labelsr   Zfit_to_window_letterboxZresize_type)Zcompress_to_fp16metadata.yamlN)Zset_rt_inforf   Ziourj   r   valuesr   ZruntimeZ
save_modelr   r    r   r   r   )ov_modelr   ovrw   rL   rM   	serialize  s    "z+Exporter.export_openvino.<locals>.serializeZ_int8_openvino_modelz.xmlznncf>=2.8.0r{   c                 S   sX   t | tr| d n| } | jtjks*J d|  tjd }|j	dkrTt
|dS |S )z Quantization transform function.imgz<Input image must be uint8 for the quantization preprocessingr&  r   r   )r   rH   dtyper   Zuint8numpyastypenpfloat32ndimZexpand_dims)Z	data_itemr   rL   rL   rM   transform_fn  s    z.Exporter.export_openvino.<locals>.transform_fnr   .r   z.*z/.*/Addz/.*/Sub*z/.*/Mul*z/.*/Div*z\.dfl.*ZSigmoid)patternstypes)rj   Zcalibration_datasetpresetignored_scoper4   )'r#   r   r3   r   rb   r   r+   r   Zconvert_modelrj   rf   r   r   r   r   r   r   r   r   ru   sepr   r  rV   nncfr2  Zndarrayr   r   joinrZ   Znamed_modulesr   ZIgnoredScopequantizeZDatasetr  ZQuantizationPresetZMIXED)rw   r_   r*  r-  ZfqZfq_ovr<  r5  r:  Zhead_module_nameZquantized_ov_modelri   Zf_ovrL   r+  rM   r     sP    
*






zExporter.export_openvinozPaddlePaddle:c                 C   s   t d ddl}ddlm} td| d|j d t| j	| jj
dtj }|| j|d	| jgd
 tt|d | j |dfS )zYOLOv8 Paddle export.)Zpaddlepaddlex2paddler   N)pytorch2paddler   z starting export with X2Paddle r  rD   r  )modulesave_dirZjit_typeZinput_examplesr(  )r#   r?  Zx2paddle.convertr@  r   rb   r   r   r   r   r   ru   r;  rj   r   r    r   r   )rw   r_   r?  r@  ri   rL   rL   rM   r     s    zExporter.export_paddlezNCNN:c              
      s  t d ddl}td| d|j d tt| j| jj	dt
j }| jd}ttrbd	nd
}| rt|nt| }| st| dt d trdntrdn
trdnd zNtdd\}} fdd|D d }	t|	tsJ dt| d|	  W nV ty\ }
 z<d}d| d  d}	t| d|
 d|	  W Y d}
~
n
d}
~
0 0 td| d|	 dd }tt |rtj|| |d! |d" t| d#|d$  d%|d&  d'|d(  g}d)|d*  d+|d,  d-|d.  d/|d0  g}t|t|g||d1t| jj d2| j j! d3| jj"d4g| j# d5}|j$dd6 t| d7d8%| d9 t&j'|dd: d;d |D }d<d=d>d?g|R D ]}t|j(dd@ qt)|dA | j* t|dfS )Bz;YOLOv8 NCNN export using PNNX https://github.com/pnnx/pnnx.rE   r   Nr   z starting export with NCNN r  rF   r0   zpnnx.exepnnxu    WARNING ⚠️ PNNX not found. Attempting to download binary file from https://github.com/pnnx/pnnx/.
Note PNNX Binary file must be placed in current working directory or in z3. See PNNX repo for full installation instructions.Zmacoswindowszlinux-aarch64linuxz	pnnx/pnnx)repoc                    s   g | ]}  d |v r|qS ).ziprL   rP   systemrL   rM   r   1  rS   z(Exporter.export_ncnn.<locals>.<listcomp>z#Unable to retrieve PNNX repo assetsz+ successfully found latest PNNX asset file Z20240410zpnnx-r/   rG  u.    WARNING ⚠️ PNNX GitHub assets not found: z, using default z/https://github.com/pnnx/pnnx/releases/download//T)delete)srcdsti  z
ncnnparam=zmodel.ncnn.paramzncnnbin=zmodel.ncnn.binzncnnpy=zmodel_ncnn.pyz
pnnxparam=zmodel.pnnx.paramzpnnxbin=zmodel.pnnx.binzpnnxpy=zmodel_pnnx.pyz	pnnxonnx=zmodel.pnnx.onnxzfp16=zdevice=zinputshape="r   ")exist_ok
 running 'r   r   )checkc                 S   s   g | ]}| d d qS )=r   )r   rP   rL   rL   rM   r   Y  rS   z	debug.binzdebug.paramz
debug2.binzdebug2.param
missing_okr(  )+r#   rE   r   rb   r   r   r   r   r   r   ru   r;  r  r   is_filer   r   r   r   r&   r   rd   r'   r"   cwdshutilmovechmodrmtreer   rf   r   r   r   r   r   mkdirr=  
subprocessrununlinkr    r   )rw   r_   rE   ri   Zf_tsrV   rC  releaseZassetsZassetrk   Z	unzip_dirZ	ncnn_argsZ	pnnx_argscmdZ
pnnx_filesZf_debugrL   rH  rM   r     st    "
0

	zExporter.export_ncnnzCoreML:c              
   C   sP  | j j dk}t|rdnd ddl}td| d|j d trNJ d	| j j	d
ksbJ d| j
|rpdnd}| rt| | j jrt| jddrt| d d| j _g d}d}d}| jjdkr| j jr|t| jj nd}| j}nN| jjdkr*| j jr"t| j| jn| j}n | j jrDt| d | j}tjj| | jdd}	|j|	|jd| jj ||dg||rdndd}
| j j!rdn| j j"rdnd\}}|dk r4d |v rtd! |r|j#j$j%&|
||}
nF|d"kr4ddl'm(  m)} |j*d |d#d$}|j+|d%}|j,|
|d&}
| j jr| jjdkr|rht-t.d'd(d)d* d}n|
/t0| t0|d+ }| j1|
|d,}
| j2}|3d-|
_4|3d.|
_5|3d/|
_6|3d0|
_7|
j89d1d2 |: D  z|
/t0| W nT t;yF } z:t| d3| d4 |d}|
/t0| W Y d}~n
d}~0 0 ||
fS )5zYOLOv8 CoreML export.rq   zcoremltools>=6.0,<=6.2zcoremltools>=7.0r   Nr   z" starting export with coremltools r  zHCoreML export is not supported on Windows, please run on macOS or Linux.r   zDCoreML batch sizes > 1 are not supported. Please retry at 'batch=1'.z.mlmodelr8   Zend2endFuT    WARNING ⚠️ 'nms=True' is not available for end2end models. Forcing 'nms=False'.)        ra  ra  gp?r   r   uQ    WARNING ⚠️ 'nms=True' is only available for Detect models like 'yolov8n.pt'.r  image)r   scalebiasZneuralnetworkr   )inputsclassifier_configZ
convert_to)r   kmeans)   Zlinear)    Nri  rg  zscikit-learnr   i   )modenbitsZweight_threshold)Zglobal_config)configz<3.11zPython T)rV   hardzData/com.apple.CoreML/weightsweights_dirr   r   r   r   c                 S   s   i | ]\}}|t |qS rL   r   )rQ   r"  r#  rL   rL   rM   
<dictcomp>  rS   z*Exporter.export_coreml.<locals>.<dictcomp>u5    WARNING ⚠️ CoreML export to *.mlpackage failed (z), reverting to *.mlmodel export. Known coremltools Python 3.11 and Windows bugs https://github.com/apple/coremltools/issues/1928.)<rf   rs   rt   r#   coremltoolsr   rb   r   r   r   r   r  is_dirrW  rZ  r   r   rj   r   r   ZClassifierConfigrZ   r   r)  IOSDetectModelr   r   r   r  r   convertZ	ImageTyper   r   r   modelsZneural_networkZquantization_utilsZquantize_weightsZcoremltools.optimize.coremlr   r7   ZOpPalettizerConfigZOptimizationConfigZpalettize_weightsr$   r   r  r   _pipeline_coremlr   popZshort_descriptionr   r   r   Zuser_defined_metadataupdater  rd   )rw   r_   rq   ctri   rd  rc  rf  rj   r  Zct_modelbitsrj  ctoZ	op_configrl  ro  r   rk   rL   rL   rM   r   `  s~    
" 
$



$zExporter.export_coremlz	TensorRT:c              
      sF  j jjdksJ d \}}zddlW n& tyT   trHtd ddlY n0 tj	ddd tj	d	d
d t
d| dj	 d tj	dd dk}t| sJ d| jd}jj}jjrjjj|_|}| }tjjd }	|r*|jj|	 n|	|_dtj j!> }
|"|
 |j#oXjj$}|j%ohjj&}' |}|(|st)d|  fddt* j+D } fddt* j,D }|D ],}t
| d|j- d|j. d|j/  q|D ],}t
| d|j- d|j. d|j/  qjj0rj j.}|d dkrTt
1| d |2 }d|d ddf}g |dd fd d!|dd D R }|D ]}|j3|j-|||d" q|4| t
| d#|rd$nd%|rd&nd'  d(|  |r^|5j6j7 |8| j9j:|_;G fd)d*d*j<}|=|djj> t?jd+d,|_@n|rr|5j6jA `BtCD  tEjFG  |r|jHn|jI}| |}tJ|d-\}tKLjM}|NtO|jPd.d/dd0 |N|Q  |N|r|n|R  W d   n1 s0    Y  W d   n1 s40    Y  |dfS )1z=YOLOv8 TensorRT export https://developer.nvidia.com/tensorrt.r   z=export running on CPU but must be on GPU, i.e. use 'device=0'r   Nztensorrt>7.0.0,<=10.1.0z>=7.0.0T)rm  z<=10.1.0z5https://github.com/ultralytics/ultralytics/pull/14239)msgr   z starting export with TensorRT r  r6  r   zfailed to export ONNX file: r6   i   @r   zfailed to load ONNX file: c                    s   g | ]}  |qS rL   )Z	get_inputrQ   inetworkrL   rM   r     rS   z*Exporter.export_engine.<locals>.<listcomp>c                    s   g | ]}  |qS rL   )Z
get_outputr}  r  rL   rM   r     rS   z input "z" with shaper   z	 output "uM    WARNING ⚠️ 'dynamic=True' model requires max batch size, i.e. 'batch=16'ri  r   c                 3   s    | ]}t d  jj| V  qdS )r   N)r   rf   	workspace)rQ   drw   rL   rM   rR     rS   z)Exporter.export_engine.<locals>.<genexpr>)minoptr   z
 building INT8ZFPZ16Z32z engine as c                       sl   e Zd Zdeedd fddZ jdddZedd	d
Ze	dddZ
edddZddddZdS )z0Exporter.export_engine.<locals>.EngineCalibratorr   N)r   cacher|   c                    s:    j |  || _t|| _ jj| _|| _t	|| _
d S r   )IInt8Calibratorrz   r  iter	data_iterCalibrationAlgoTypeZENTROPY_CALIBRATION_2algor   r   r  )rw   r  r   r  r}   rL   rM   rz     s    

z9Exporter.export_engine.<locals>.EngineCalibrator.__init__r{   c                 S   s   | j S )z%Get the calibration algorithm to use.)r  r  rL   rL   rM   get_algorithm  s    z>Exporter.export_engine.<locals>.EngineCalibrator.get_algorithmc                 S   s
   | j pdS )z*Get the batch size to use for calibration.r   )r   r  rL   rL   rM   get_batch_size  s    z?Exporter.export_engine.<locals>.EngineCalibrator.get_batch_sizec                 S   sV   z<t | jd d }|jjdkr*|dn|}t| gW S  tyP   Y dS 0 dS )zOGet the next batch to use for calibration, as a list of device memory pointers.r.  r&  r   r  N)nextr  r   r   r   r   Zdata_ptrStopIteration)rw   r   Zim0srL   rL   rM   	get_batch  s    z:Exporter.export_engine.<locals>.EngineCalibrator.get_batchc                 S   s$   | j  r | j jdkr | j  S dS )zSUse existing cache instead of calibrating again, otherwise, implicitly return None..cacheN)r  existsr   
read_bytesr  rL   rL   rM   read_calibration_cache  s    zGExporter.export_engine.<locals>.EngineCalibrator.read_calibration_cachec                 S   s   | j |}dS )z Write calibration cache to disk.N)r  write_bytes)rw   r  r   rL   rL   rM   write_calibration_cache  s    zHExporter.export_engine.<locals>.EngineCalibrator.write_calibration_cache)r   )__name__
__module____qualname__r   r   rz   r  r  r  rZ   r  bytesr  r  rL   r  rL   rM   EngineCalibrator  s    
r  r  )r  r   r  wbr   little)	byteordersigned)Sr   r   r   r   r~   ImportErrorr   r#   r$   r   r   rb   r   r   r   r  r   r  LoggerINFOrf   r  ZSeverityVERBOSEZmin_severityBuilderZcreate_builder_configr  Zset_memory_pool_limitZMemoryPoolType	WORKSPACEZmax_workspace_sizeZNetworkDefinitionCreationFlagZEXPLICIT_BATCHZcreate_networkZplatform_has_fast_fp16r   Zplatform_has_fast_int8r   Z
OnnxParserZparse_from_fileRuntimeErrorr   Z
num_inputsZnum_outputsrV   r   r/  r   r   Zcreate_optimization_profile	set_shapeZadd_optimization_profileZset_flagZBuilderFlagr  Zset_calibration_profileZProfilingVerbosityZDETAILEDZprofiling_verbosityr  r  r   r   Zint8_calibratorZFP16rj   gcZcollectr   r  Zempty_cacheZbuild_serialized_networkZbuild_engineopenr	  r
  r   writer   to_bytesencoder-  )rw   r_   f_onnxr   Zis_trt10ri   loggerbuilderrl  r  flagr   r   parserre  outputsinpoutr   ZprofileZ	min_shapeZ	max_shaper  buildr5   rc   r$  rL   )r  rw   r}   rM   r     s    

**
,
2

*


XzExporter.export_enginezTensorFlow SavedModel:c                 C   s  t j }zddl}W nL tyb   tr,dntr4dn
|r<dnd}d}td| |  ddl}Y n0 td	d
dddddtr|dndd|rdndf
dd t	
d| d|j d t|jddddd ddl}tt| j| jjd}| rt| td}| st| dddd  d| j_|  \}	}
d}| jjr|d! }d"}| jjr|  d#d$ | |D }t |d  }t!"t||# $t!j% d%|g d&gggg d'ggggg}nd(}t	
| d)|j d |j&|	t|d|| jjd*|ddd+	 t'|d, | j( | jjrr|j)dd- |*d.D ]&}|+|,|j-d/d0|j  q0|*d1D ]}|)  qb|*d2D ]&}d3t|v r|) n| .| q|t||j/j0|ddd4fS )5z$YOLOv8 TensorFlow SavedModel export.r   Nz-macosz-aarch64r   z-cpuz>=2.0.0
tensorflowZkerasZtf_keraszsng4onnx>=1.0.1zonnx_graphsurgeon>=0.3.26r  zonnx2tf>1.17.5,<=1.22.3zonnxslim>=0.1.31ztflite_support<=0.4.3tflite_supportzflatbuffers>=23.5.26,<100zonnxruntime-gpur  z---extra-index-url https://pypi.ngc.nvidia.com)Zcmdsr   ! starting export with tensorflow r  Tz6https://github.com/ultralytics/ultralytics/issues/5161)rV   r  r|  r:   z6calibration_image_sample_data_20x128x128x3_float32.npyrG  )ZunziprK  z&tmp_tflite_int8_calibration_images.npyrb   c                 S   s    g | ]}|d   ddddqS )r.  r   r   r   r   )Zpermute)rQ   r   rL   rL   rM   r   l  rS   z/Exporter.export_saved_model.<locals>.<listcomp>r  )r   r   r   )   r  r  re   z% starting TFLite export with onnx2tf z
per-tensor)	Zinput_onnx_file_pathZoutput_folder_pathZnot_use_onnxsim	verbosityZoutput_integer_quantized_tfliteZ
quant_typeZ!custom_input_op_name_np_data_pathZdisable_group_convolutionZenable_batchmatmul_unfoldr(  rS  z*_dynamic_range_quant.tfliteZ_dynamic_range_quantZ_int8z%*_integer_quant_with_int16_act.tflitez*.tflitezquant_with_int16_act.tflite)tagsoptions)1r   r  r  r  r  r   r   r#   r   r   rb   r   r$   onnx2tfr   r   r   r   r   rr  rW  rZ  r  r%   rf   r  r   r   r   r[  r  catr   r2  r  r0  r1  r3  rt  r    r   r^  rglobrename	with_namer   _add_tflite_metadatar9   r  )rw   r_   r  tfr   r   r  ri   Zonnx2tf_filer  r   Znp_datatmp_filer  r  r   rL   rL   rM   r   3  s    






$
$$zExporter.export_saved_modelzTensorFlow GraphDef:c                    s   ddl }ddlm} td| d|j d | jd}| fdd	}|	|
 jd j jd j}||}|j  |jj|jt|j|jd
d |dfS )zYYOLOv8 TensorFlow GraphDef *.pb export https://github.com/leimao/Frozen_Graph_TensorFlow.r   N)!convert_variables_to_constants_v2r   r  r  r<   c                    s    | S r   rL   rJ   r   rL   rM   <lambda>  rS   z$Exporter.export_pb.<locals>.<lambda>F)Zgraph_or_graph_defZlogdirrV   Zas_text)r  Z0tensorflow.python.framework.convert_to_constantsr  r   rb   r   r   r  functionZget_concrete_functionZ
TensorSpecre  r   r/  graphas_graph_defioZwrite_graphr   r   rV   )rw   r   r_   r  r  ri   r   Zfrozen_funcrL   r   rM   r     s    "
zExporter.export_pbzTensorFlow Lite:c                 C   s   ddl }td| d|j d tt| j| jjd}| j	j
rX|| jj d }n.| j	jrt|| jj d }n|| jj d	 }t|dfS )
zYOLOv8 TensorFlow Lite export.r   Nr   r  r  r:   z_int8.tflitez_float16.tflitez_float32.tflite)r  r   rb   r   r   r   r   r   r   rf   r   r   r   )rw   r   r   r   r_   r  r9   ri   rL   rL   rM   r     s    zExporter.export_tflitez	Edge TPU:c           	      C   s0  t | d d}d}ts*J d| tj|tjtjddjdkrt d| d	|  tjd
ddjdk}dD ]$}tj|r|n
|ddddd qttj|ddddj	
  d }t d| d| d t|dd}dt|j d| d}t | d| d tj|dd | | |dfS )zCYOLOv8 Edge TPU export https://coral.ai/docs/edgetpu/models-intro/.uY    WARNING ⚠️ Edge TPU known bug https://github.com/ultralytics/ultralytics/issues/1185zedgetpu_compiler --versionz'https://coral.ai/docs/edgetpu/compiler/z$export only supported on Linux. See T)stdoutstderrshellr   r   z< export requires Edge TPU compiler. Attempting install from zsudo --version >/dev/nullr  )zOcurl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -zecho "deb https://packages.cloud.google.com/apt coral-edgetpu-stable main" | sudo tee /etc/apt/sources.list.d/coral-edgetpu.listzsudo apt-get updatez%sudo apt-get install edgetpu-compilerzsudo r   )r  rQ  )r  capture_outputrQ  r   z( starting export with Edge TPU compiler r  r>   r@   z(edgetpu_compiler -s -d -k 10 --out_dir "" "rN  rP  r   N)r   r   r   r\  r]  DEVNULL
returncoderb   r   r  decoder   r   r   r   r  )	rw   r   r_   r`  Zhelp_urlZsudocverri   rL   rL   rM   r     s"    " 
zExporter.export_edgetpuzTensorFlow.js:c                 C   s  t d trt d ddl}ddl}td| d|j d t| j	| jj
d}t| jd	}|  }t|d
}||  W d   n1 s0    Y  dt|}td| d|  | jjrdn| jjrdnd}	t||}
t|R}d|	 d| d|
 d| d	}t| d| d tj|dd W d   n1 sT0    Y  W d   n1 st0    Y  d|v rt| d| d tt|d | j |dfS )zYOLOv8 TensorFlow.js export.tensorflowjsznumpy==1.23.5r   Nr   z# starting export with tensorflowjs r  rB   r<   rb,z output node names: z--quantize_float16z--quantize_uint8r   z6tensorflowjs_converter --input_format=tf_frozen_model z --output_node_names=z "r  rN  rP  r   Tr  r   uG    WARNING ⚠️ your model may not work correctly with spaces in path 'r   r(  )r#   r   r  r  r   rb   r   r   r   r   r   r  ZGraphr  r  ParseFromStringreadr=  r^   rf   r   r   r)   r\  r]  r   r    r   r   )rw   r_   r  rA   ri   Zf_pbr\   r   r  ZquantizationZfpb_Zf_r`  rL   rL   rM   r     s>    ,N
zExporter.export_tfjsc                 C   s  ddl }zddlm} ddlm} W n* tyN   ddlm} ddlm} Y n0 | }| jd |_| jd |_	| jd |_
| jd |_t|jd	 }t|d
 }|t| j W d   n1 s0    Y  | }|j|_|jj|_| }	d|	_d|	_| |	_| |	j_|jj|	jj_|jj|	j_ | }
d|
_d|
_|g|
_!| j"j#dkrt| }d|_d|_|g|_!|$ }|	g|_%| j"j#dkr|
|gn|
g|_&|g|_'|(d}|)|*||j+j, |- }|j+.t|}|/| |0t|g |1  |2  dS )z\Add metadata to *.tflite models per https://www.tensorflow.org/lite/models/convert/metadata.r   N)metadata_schema_py_generated)r   r   r   r   r   ztemp_meta.txtwrb  zInput image to be detected.outputzCCoordinates of detected objects, class labels, and confidence scorer   zMask protos)3flatbuffersZ tensorflow_lite_support.metadatar  Z'tensorflow_lite_support.metadata.pythonr   r  r  ZModelMetadataTrV   r   r   r   r   r   r  r  r   ZAssociatedFileTZAssociatedFileTypeZTENSOR_AXIS_LABELSr   ZTensorMetadataTr   ZContentTcontentZImagePropertiesTZcontentPropertiesZColorSpaceTypeRGBZ
colorSpaceZContentPropertiesZImagePropertiesZcontentPropertiesTypeZassociatedFilesrj   r   ZSubGraphMetadataTZinputTensorMetadataZoutputTensorMetadataZsubgraphMetadatar  ZFinishZPackZMetadataPopulatorZMETADATA_FILE_IDENTIFIEROutputZwith_model_fileZload_metadata_bufferZload_associated_filesZpopulater^  )rw   r   r  Zschemar   Z
model_metar  ri   Z
label_fileZ
input_metar  Zoutput2ZsubgraphbZmetadata_bufZ	populatorrL   rL   rM   r    s\    .



zExporter._add_tflite_metadatazCoreML Pipeline:c                 C   s  ddl }t| d|j d t| jj\}}}}| }t|j	j
\}	}
trddlm} |d||f}|d|i}||	j j}||
j j}n&| jd | jd	 d
 f}| jd d
f}| jd }|j	jd jjj|j	jd jjj }}|\}}t||ksJ t| d| ||	jjjdd< ||
jjjdd< |jj||d}|jj }d|_t dD ]T}|j!j	j
| " }|j	j#  |j	j| $| |j	j
#  |j	j
| $| qdd|j	j
d _d|j	j
d	 _|d
g}t dD ]z}|j	j
| jj}|j%j&#  d|j%j&d _'d|j%j&d _(|j%j&#  || |j%j&d	 _'|| |j%j&d	 _(|jdd= q|j)}|	j|_*|
j|_+d|_,d|_-d|_.d|_/d|_0d|_1d|j2_3|j4j56|7  |j|}|jj8j9d|jj:;d||fd|jj:< fd|jj:< fgddgd}|=| |=| |j>j	jd $|j!j	jd "  |j>j	j
d $|j!j	j
d "  |j>j	j
d	 $|j!j	j
d	 "  d|j>_|j>j	jj?@tA|j0tA|j1d |jj|j>|d}d|jBd< d|j0 d|jBd< d|j1 d|jBd< d|jCd< d|jCd< t| d  |S )!zYOLOv8 CoreML pipeline.r   Nz$ starting pipeline with coremltools r  )Imager  rb  r   r   r   r   z names found for nc=rn  r   
confidenceZcoordinatesr   iouThresholdconfidenceThresholdg?g      ?Tr   )Zinput_featuresZoutput_features)zIoU thresholdzConfidence thresholdzInput imagez,(optional) IoU threshold override (default: )z3(optional) Confidence threshold override (default: u?   Boxes × Class confidence (see user-defined metadata "classes")u7   Boxes × [x, y, width, height] (relative to image size)z pipeline success)Drq  r   rb   r   rZ   r   r   Zget_specr  r   r  r   ZPILr  newZpredictrV   r   r   rX   r   Z	imageTyper  r  r   ZmultiArrayTyperu  ZMLModelprotoZ	Model_pb2ZModelZspecificationVersionr   _specZSerializeToStringr  r  Z
shapeRangeZ
sizeRangesZ
lowerBoundZ
upperBoundZnonMaximumSuppressionZconfidenceInputFeatureNameZcoordinatesInputFeatureNameZconfidenceOutputFeatureNameZcoordinatesOutputFeatureNameZiouThresholdInputFeatureNameZ#confidenceThresholdInputFeatureNamer  r  ZpickTopZperClassZstringClassLabelsZvectorrW   r)  pipelineZPipelineZ	datatypesZArrayDoubleZ	add_modelspecZuserDefinedrx  r   Zinput_descriptionZoutput_description)rw   rj   ro  r_   ry  r   hr  r  Zout0Zout1r  r.  r  Z
out0_shapeZ
out1_shaper   nxnyncZnms_specr~  Zdecoder_outputZoutput_sizesZma_typer   Z	nms_modelr  rL   rL   rM   rv  1  s    
&$

"""


zExporter._pipeline_coreml)eventc                 C   s   | j | | dS )zAppends the given callback.N)r   rU   rw   r  callbackrL   rL   rM   add_callback  s    zExporter.add_callbackc                 C   s    | j |g D ]}||  qdS )z(Execute all callbacks for a given event.N)r   r   r  rL   rL   rM   r     s    zExporter.run_callbacks)N)r   )r  r  r  __doc__r   rz   r.   r   r   r  ro   r   r   r   r   r   r   r   r   r   r   r   r   r   r  rv  r  r   rL   rL   rL   rM   rp      sF    ?
8GCL [#@jrp   c                       s(   e Zd ZdZ fddZdd Z  ZS )rs  z;Wrap an Ultralytics YOLO model for Apple iOS CoreML export.c                    sf   t    |j\}}}}|| _t|j| _||kr>d| | _n$t	d| d| d| d| g| _dS )zHInitialize the IOSDetectModel class with a YOLO model and example image.g      ?N)
superrz   r   rj   r   r   r  	normalizer   Ztensor)rw   rj   r   r   r  r  	__class__rL   rM   rz     s    
zIOSDetectModel.__init__c                 C   s6   |  |d ddd| jfd\}}||| j fS )zRNormalize predictions of object detection model with input size-dependent factors.r   r   r   )rj   Z	transposer   r  r  )rw   rK   ZxywhclsrL   rL   rM   r     s    (zIOSDetectModel.forward)r  r  r  r  rz   r   __classcell__rL   rL   r  rM   rs    s   rs  )Nr  r  r	  ru   rW  r\  r   r   copyr   r   pathlibr   r0  r2  r   Zultralytics.cfgr   r   Zultralytics.datar   Zultralytics.data.datasetr   Zultralytics.data.utilsr	   r
   Zultralytics.nn.autobackendr   r   Zultralytics.nn.modulesr   r   r   Zultralytics.nn.tasksr   r   r   Zultralytics.utilsr   r   r   r   r   r   r   r   r   r   r   r   r   r    Zultralytics.utils.checksr!   r"   r#   r$   Zultralytics.utils.downloadsr%   r&   r'   Zultralytics.utils.filesr(   r)   Zultralytics.utils.opsr*   Zultralytics.utils.torch_utilsr+   r,   r-   r.   rN   r^   ro   rp   nnModulers  rL   rL   rL   rM   <module>   sL   3@	        