U
    Hcl                     @   sb  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Zddl	m
Z
 ddlZddlZddlmZ ddlmZ e
e Zejd Zeeejkrejee e
ejee
 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)m*Z*m+Z+m,Z,m-Z-m.Z.m/Z/ ddl0m1Z1 dd Z2e,dfddZ3e,dfddZ4e,dfddZ5e,dfddZ6dde,dfddZ7ddddd d!de,d"fd#d$Z8e,d%fd&d'Z9e,d(fd)d*Z:e,d+fd,d-Z;e,d.fd/d0Z<e= ed1 ed2 d3d4d5d6dddddddd7ddddddd d!fd8d9Z>d:d; Z?d<d= Z@eAd>kr^e? ZBe@eB dS )?a 
  
Export a YOLOv5 PyTorch model to other formats. TensorFlow exports authored by https://github.com/zldrobit

Format                      | `export.py --include`         | Model
---                         | ---                           | ---
PyTorch                     | -                             | yolov5s.pt
TorchScript                 | `torchscript`                 | yolov5s.torchscript
ONNX                        | `onnx`                        | yolov5s.onnx
OpenVINO                    | `openvino`                    | yolov5s_openvino_model/
TensorRT                    | `engine`                      | yolov5s.engine
CoreML                      | `coreml`                      | yolov5s.mlmodel
TensorFlow SavedModel       | `saved_model`                 | yolov5s_saved_model/
TensorFlow GraphDef         | `pb`                          | yolov5s.pb
TensorFlow Lite             | `tflite`                      | yolov5s.tflite
TensorFlow Edge TPU         | `edgetpu`                     | yolov5s_edgetpu.tflite
TensorFlow.js               | `tfjs`                        | yolov5s_web_model/

Requirements:
    $ pip install -r requirements.txt coremltools onnx onnx-simplifier onnxruntime openvino-dev tensorflow-cpu  # CPU
    $ pip install -r requirements.txt coremltools onnx onnx-simplifier onnxruntime-gpu openvino-dev tensorflow  # GPU

Usage:
    $ python path/to/export.py --weights yolov5s.pt --include torchscript onnx openvino engine coreml tflite ...

Inference:
    $ python path/to/detect.py --weights yolov5s.pt                 # PyTorch
                                         yolov5s.torchscript        # TorchScript
                                         yolov5s.onnx               # ONNX Runtime or OpenCV DNN with --dnn
                                         yolov5s.xml                # OpenVINO
                                         yolov5s.engine             # TensorRT
                                         yolov5s.mlmodel            # CoreML (MacOS-only)
                                         yolov5s_saved_model        # TensorFlow SavedModel
                                         yolov5s.pb                 # TensorFlow GraphDef
                                         yolov5s.tflite             # TensorFlow Lite
                                         yolov5s_edgetpu.tflite     # TensorFlow Edge TPU

TensorFlow.js:
    $ cd .. && git clone https://github.com/zldrobit/tfjs-yolov5-example.git && cd tfjs-yolov5-example
    $ npm install
    $ ln -s ../../yolov5/yolov5s_web_model public/yolov5s_web_model
    $ npm start
    N)Path)optimize_for_mobile)Conv)attempt_load)Detect)SiLU)
LoadImages)	LOGGERcheck_datasetcheck_img_sizecheck_requirementscheck_versioncolorstr	file_size
print_argsurl2file)select_devicec                  C   sp   dddgdddgddd	gd
ddgdddgdddgdddgdddgdddgdddgdd d!gg} t j| d"d#d$gd%S )&NZPyTorch-.ptZTorchScripttorchscript.torchscriptONNXonnx.onnxZOpenVINOopenvino_openvino_modelZTensorRTengine.engineZCoreMLcoreml.mlmodelzTensorFlow SavedModelsaved_model_saved_modelzTensorFlow GraphDefpb.pbzTensorFlow Litetflitez.tflitezTensorFlow Edge TPUedgetpuz_edgetpu.tflitezTensorFlow.jstfjs
_web_modelFormatArgumentSuffix)columns)pd	DataFramex r0   </var/www/html/VideoAnalyticsDashboard/src/./yolov5/export.pyexport_formatsL   s    r2   zTorchScript:c           
   
   C   s   zt d| dtj d |d}tjj| |dd}|jtt	| j
| jd}dt|i}|r|t|jt||d	 n|jt||d	 t | d
| dt|dd |W S  tk
r }	 zt | d|	  W 5 d }	~	X Y nX d S )N
z starting export with torch ...r   Fstrict)shapestridenamesz
config.txt)_extra_files export success, saved as  (.1f MB) export failure: )r	   infotorch__version__with_suffixjittracer7   intmaxr8   r9   jsondumpsr   _save_for_lite_interpreterstrsaver   	Exception)
modelimfileoptimizeprefixftsdextra_fileser0   r0   r1   export_torchscript\   s    
"rX   zONNX:c                 C   s  ztt d dd l}td| d|j d |d}	tjj| ||	d||rVtjjj	ntjjj
| dgd	g|rd
dddd
dddnd d
 ||	}
|j|
 |rPzlt d dd l}t| d|j d |j|
||rdt|jind d\}
}|std||
|	 W n8 tk
rN } zt| d|  W 5 d }~X Y nX t| d|	 dt|	dd |	W S  tk
r } zt| d|  W 5 d }~X Y nX d S )N)r   r   r3   z starting export with onnx r4   r   Fimagesoutputbatchheightwidth)r         anchors)r      )rY   rZ   )verboseopset_versiontrainingdo_constant_foldinginput_namesoutput_namesdynamic_axes)zonnx-simplifierz" simplifying with onnx-simplifier )Zdynamic_input_shapeinput_shapeszassert check failedz simplifier failure: r;   r<   r=   r>   r?   )r   r   r	   r@   rB   rC   rA   exportTrainingModeTRAININGEVALloadcheckerZcheck_modelonnxsimsimplifylistr7   AssertionErrorrL   rM   r   )rN   rO   rP   opsettraindynamicrq   rR   r   rS   Z
model_onnxrp   checkrW   r0   r0   r1   export_onnxp   sH    



&"rx   z	OpenVINO:c              
   C   s   zt d dd lm} td| d|j d t|ddtj	 }d|
d	 d
| }tj|dd t| d| dt|dd |W S  tk
r } ztd| d|  W 5 d }~X Y nX d S )N)zopenvino-devr   r3   z starting export with openvino r4   r   r   zmo --input_model r   z --output_dir Tshellr;   r<   r=   r>   r?   )r   openvino.inference_engineinference_enginer	   r@   rB   rK   replaceosseprC   
subprocesscheck_outputr   rM   )rN   rO   rP   rR   ierS   cmdrW   r0   r0   r1   export_openvino   s    "r   zCoreML:c           	   
   C   s   zt d dd l}td| d|j d |d}tjj| |dd}|j	||j
d	|jd
dddgdgd}|| t| d| dt|dd ||fW S  tk
r } z td| d|  W Y dS d }~X Y nX d S )N)coremltoolsr   r3   z" starting export with coremltools r4   r   Fr5   imagegp?)r7   scalebias)inputsr;   r<   r=   r>   r?   NN)r   r   r	   r@   rB   rC   rA   rD   rE   convertZ	ImageTyper7   rL   r   rM   )	rN   rO   rP   rR   ctrS   rT   Zct_modelrW   r0   r0   r1   export_coreml   s    
&
"
r      Fz	TensorRT:c	              
      s  zt d dd l}	|	jd dkrf| jd j}
dd |
D | jd _t| ||d|d| |
| jd _n$t|	jd	d
d t| ||d|d| |d}t	d| d|	j d |j
jdkstd| std| |d}|	|	jj}|r|	jjj|_|	|}| }|d d> |_dt|	jj> }|| |	 |}|t|sftd|  fddt jD } fddt j D }t	| d |D ],}t	| d|j! d|j" d|j#  q|D ],}t	| d|j! d|j" d|j#  q||j$M }t	| d|r(d nd! d"|  |rL|%|	j&j' |( |*}t)|d#}|*|+  W 5 Q R X W 5 Q R X t	| d$| d%t,|d&d' |W S  t-k
r } zt	d| d(|  W 5 d }~X Y nX d S ))N)tensorrtr   7c                 S   s(   g | ] }|d ddddddf qS ).Nra   r0   ).0ar0   r0   r1   
<listcomp>   s     z!export_engine.<locals>.<listcomp>   Fz8.0.0T)hard   r   r3   z starting export with TensorRT r4   cpuzLexport running on CPU but must be on GPU, i.e. `python export.py --device 0`zfailed to export ONNX file: r   ra      zfailed to load ONNX file: c                    s   g | ]}  |qS r0   )Z	get_inputr   inetworkr0   r1   r      s     c                    s   g | ]}  |qS r0   )Z
get_outputr   r   r0   r1   r      s     z Network Description:z	input "z" with shape z and dtype z		output "z building FP       z engine in wbr;   r<   r=   r>   r?   ).r   r   rB   rN   anchor_gridrx   r   rC   r	   r@   devicetypers   existsLoggerINFOZSeverityVERBOSEZmin_severityBuilderZcreate_builder_configZmax_workspace_sizerF   ZNetworkDefinitionCreationFlagZEXPLICIT_BATCHZcreate_networkZ
OnnxParserZparse_from_filerK   RuntimeErrorrange
num_inputsnum_outputsnamer7   dtypeZplatform_has_fast_fp16Zset_flagZBuilderFlagZFP16Zbuild_engineopenwrite	serializer   rM   )rN   rO   rP   ru   halfrq   	workspacerb   rR   trtgridr   rS   loggerbuilderconfigflagparserr   outputsinpoutr   trW   r0   r   r1   export_engine   sV    



**
$""r   d   ?      ?zTensorFlow SavedModel:c              
      s  zdd l }ddlm} ddlm}m} td| d|j d t	|
dd}t|j^}}}|| j| | j|d	}||f|d}||||||||	}|jj|d|rd n|d}||||||||	}|jj||dd_  |
rj|dd n|fdd}|jd jjd j}||}|| | }| fdd|g|_|| |jj||t|jdr|jj ddn|j  d t| d| dt!|dd |fW S  t"k
r } z td| d|  W Y dS d }~X Y nX d S )Nr   !convert_variables_to_constants_v2)TFDetectTFModelr3   ! starting export with tensorflow r4   r   r!   )cfgrN   ncimgszr_   )r7   
batch_size)r   r   Ftf)Zsave_formatc                    s    | S Nr0   r.   keras_modelr0   r1   <lambda>      z$export_saved_model.<locals>.<lambda>c                    s    | S r   r0   r.   )frozen_funcr0   r1   r     r   z2.6)Zexperimental_custom_gradients)optionsr;   r<   r=   r>   r?   r   )r_   )r_   )#
tensorflow0tensorflow.python.framework.convert_to_constantsr   	models.tfr   r   r	   r@   rB   rK   r}   rr   r7   yamlr   zerospredictkerasInputModelZ	trainablesummaryrL   function
TensorSpecr   r   get_concrete_functionModule__call__r    r   ZSaveOptionsr   rM   )rN   rO   rP   rv   tf_nmsagnostic_nmstopk_per_classtopk_all	iou_thres
conf_thresr   rR   r   r   r   r   rS   r   chr   Ztf_model_r   r   mspectfmrW   r0   )r   r   r1   export_saved_model   sH    


"
r   zTensorFlow GraphDef:c           
   
      s  zdd l }ddlm} td| d|j d |d}| fdd}||	 j
d j j
d j}||}|j  |jj|jt|j|jd	d
 t| d| dt|dd |W S  tk
 r }	 ztd| d|	  W 5 d }	~	X Y nX d S )Nr   r   r3   r   r4   r#   c                    s    | S r   r0   r.   r   r0   r1   r   -  r   zexport_pb.<locals>.<lambda>F)Zgraph_or_graph_defZlogdirr   as_textr;   r<   r=   r>   r?   )r   r   r   r	   r@   rB   rC   r   r   r   r   r7   r   graphas_graph_defioZwrite_graphrK   parentr   r   rM   )
r   rO   rP   rR   r   r   rS   r   r   rW   r0   r   r1   	export_pb$  s    
"
"r   zTensorFlow Lite:c              
      sr  z0dd l }td| d|j d t|j^}}	}
t|dd}|jj	
| }|jjjg|j_|jg|j_|jjjg|_|rddlm tt|d |
d	d
  fdd|_|jjjg|j_g |j_|j|_|j|_d	|_t|dd}| }t|d | t| d| dt!|dd |W S  t"k
rl } ztd| d|  W 5 d }~X Y nX d S )Nr   r3   r   r4   r   z-fp16.tflite)representative_dataset_genru   F)img_sizeautoc                      s
    S r   r0   r0   datasetncalibr   r0   r1   r   I  r   zexport_tflite.<locals>.<lambda>-int8.tfliter   r;   r<   r=   r>   r?   )#r   r	   r@   rB   rr   r7   rK   r}   liteZTFLiteConverterZfrom_keras_modelZOpsSetZTFLITE_BUILTINSZtarget_specZsupported_opsfloat16Zsupported_typesZOptimizeDEFAULTZoptimizationsr   r   r   r
   Zrepresentative_datasetZTFLITE_BUILTINS_INT8uint8Zinference_input_typeZinference_output_typeZexperimental_new_quantizerr   r   r   r   rM   )r   rO   rP   int8datar   rR   r   r   r   r   rS   	converterZtflite_modelrW   r0   r   r1   export_tflite9  s2    "r  z	Edge TPU:c              
   C   sj  z(d}d}t  dks&td| tj|d 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 qltj|ddddj	
  d }td	| d| d t|dd}	t|dd}
d|
 }tj|ddd t| d|	 dt|	dd |	W S  tk
rd } ztd	| d|  W 5 d }~X Y nX d S )Nzedgetpu_compiler --versionz'https://coral.ai/docs/edgetpu/compiler/Linuxz$export only supported on Linux. See z >/dev/nullTry   r   r3   z< export requires Edge TPU compiler. Attempting install from zsudo --version >/dev/null)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  )rz   rw   )rz   capture_outputrw   r   z( starting export with Edge TPU compiler r4   r   z-int8_edgetpu.tfliter   zedgetpu_compiler -s r;   r<   r=   r>   r?   )platformsystemrs   r   run
returncoder	   r@   r}   stdoutdecodesplitrK   r   rM   )r   rO   rP   rR   r   Zhelp_urlZsudocverrS   Zf_tflrW   r0   r0   r1   export_edgetpuY  s&    " 
"r  zTensorFlow.js:c              
   C   s  zt d dd l}dd l}td| d|j d t|dd}|d}|d	 }d
| d| }	t	j
|	dd t| }
t|d}|dd|
}|| W 5 Q R X t| d| dt|dd |W S  tk
r } ztd| d|  W 5 d }~X Y nX d S )N)tensorflowjsr   r3   z# starting export with tensorflowjs r4   r   r'   r#   z/model.jsonzvtensorflowjs_converter --input_format=tf_frozen_model --output_node_names="Identity,Identity_1,Identity_2,Identity_3"  Try   wz{"outputs": {"Identity.?.?": {"name": "Identity.?.?"}, "Identity.?.?": {"name": "Identity.?.?"}, "Identity.?.?": {"name": "Identity.?.?"}, "Identity.?.?": {"name": "Identity.?.?"}}}z{"outputs": {"Identity": {"name": "Identity"}, "Identity_1": {"name": "Identity_1"}, "Identity_2": {"name": "Identity_2"}, "Identity_3": {"name": "Identity_3"}}}r;   r<   r=   r>   r?   )r   rer  r	   r@   rB   rK   r}   rC   r   r	  r   readsubr   r   rM   )r   rO   rP   rR   r  r&   rS   Zf_pbZf_jsonr   rH   jZsubstrW   r0   r0   r1   export_tfjsv  s,    

"r  data/coco128.yaml
yolov5s.pt)  r  ra   r   )r   r   r   c           .         s6  t   }dd D tt d dd  }fdd|D }t|tksbtd d| |\
}}}}}}}} }!}"tt|drt	|n|}#t
|}|jd	kr|rtd
t||ddd}$|$j|$j }%}&|t|dkrdnd9 }dkrdn|}|%t|&ks"td|% dt|& tt|$j  fdd|D }tj|df| |}'|rr|' |$  }'}$|r|$ n|$  |$ D ]\\}(})t|)trt|)jtjrt |)_n,t|)tr||)_||)_ t!|)dr|)j"|)_#qt$dD ]}*|$|'}+qt|+d j%},t&'dt(d d|# d|, dt)|#dd	 dgd }-t*j+dtj,j-d  |rvt.|$|'|#|	|-d< |rt/|$|'|#||||||-d< |s|rt0|$|'|#|||||-d< |rt1|$|'|#|-d< |rt2|$|'|#\}*|-d!< t3||| |!|"fr|
s
|!rt4d" | r&|"r&td#t5|$|'|#||p>|p>|"|pF|"||||d$
\}$|-d%< |sh|"rxt6|$|'|#|-d&< | s|!rt7|$|'|#|
p|!| d'd(|-d)< |!rt8|$|'|#|-d*< |"rt9|$|'|#|-d+< d,d |-D }-t3|-r2t&'d-t   | d.d/t(d0|#j:;  d1|-d2  d3|-d2  d4|-d2  d5 |-S )6Nc                 S   s   g | ]}|  qS r0   )lowerr   r/   r0   r0   r1   r     s     zrun.<locals>.<listcomp>r)   ra   c                    s   g | ]}| kqS r0   r0   r  )includer0   r1   r     s     zERROR: Invalid --include z , valid --include arguments are )zhttp:/zhttps:/r   z;--half only compatible with GPU export, i.e. use --device 0T)map_locationinplacefuser^   r   r   zModel class count z != len(names) c                    s   g | ]}t | qS r0   )r   r  )gsr0   r1   r     s     r_   forward_exportr   r3   zPyTorch:z starting from z with output shape r<   r=   r>   r  
   ignore)actioncategoryr   )zflatbuffers==1.12zOTFLite and TF.js models must be exported separately, please pass only one type.)r   r   r   r   r   r         r   )r   r  r         	   c                 S   s   g | ]}|rt |qS r0   )rK   r  r0   r0   r1   r     s      z
Export complete (z.2fzs)
Results saved to boldz-
Detect:          python detect.py --weights r   zJ
PyTorch Hub:     model = torch.hub.load('ultralytics/yolov5', 'custom', 'z,')
Validate:        python val.py --weights z$
Visualize:       https://netron.app)<timetupler2   sumlenrs   r   rK   
startswithr   r   r   r   r   r9   rF   rG   r8   rA   r   tor   ru   evalnamed_modules
isinstancer   actnnr   r   r   Zonnx_dynamichasattrr#  forwardr   r7   r	   r@   r   r   warningsfilterwarningsrD   TracerWarningrX   r   rx   r   r   anyr   r   r   r  r  r  r   resolve).r  weightsr   r   r   r  r   r   ru   rQ   r   rv   rq   rt   rb   r   nmsr   r   r   r   r   r   formatsflagsrD   r   xmlr   r   r    r"   r$   r%   r&   rP   rN   r   r9   rO   kr   r   yr7   rS   r0   )r"  r  r1   r	    s    $&
.
   
Lr	  c               	   C   s  t  } | jdttd dd | jddttd dd	 | jd
dddtddgdd	 | jdtddd | jdddd | jdddd | jdddd | jdddd | jdddd | jd dd!d | jd"dd#d | jd$dd%d | jd&td'd(d | jd)dd*d | jd+td,d-d | jd.dd/d | jd0dd1d | jd2td3d4d | jd5td3d6d | jd7td8d9d | jd:td;d<d | jd=dd>d?gd@dA |  }tt	j
| |S )BNz--datar  zdataset.yaml path)r   defaulthelpz	--weights+r  zmodel.pt path(s))nargsr   rG  rH  z--imgszz--imgz
--img-sizer  zimage (h, w)z--batch-sizera   z
batch sizez--devicer   z%cuda device, i.e. 0 or 0,1,2,3 or cpu)rG  rH  z--half
store_truezFP16 half-precision export)r&  rH  z	--inplacez set YOLOv5 Detect() inplace=Truez--trainzmodel.train() modez
--optimizez TorchScript: optimize for mobilez--int8zCoreML/TF INT8 quantizationz	--dynamiczONNX/TF: dynamic axesz
--simplifyzONNX: simplify modelz--opsetr   zONNX: opset versionz	--verbosezTensorRT: verbose logz--workspacer   zTensorRT: workspace size (GB)z--nmszTF: add NMS to modelz--agnostic-nmszTF: add agnostic NMS to modelz--topk-per-classr   z!TF.js NMS: topk per class to keepz
--topk-allz'TF.js NMS: topk for all classes to keepz--iou-thresr   zTF.js NMS: IoU thresholdz--conf-thresr   zTF.js NMS: confidence thresholdz	--includer   r   zStorchscript, onnx, openvino, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs)rJ  rG  rH  )argparseArgumentParseradd_argumentrK   ROOTrF   float
parse_argsr   FILEstem)r   optr0   r0   r1   	parse_opt	  s:    rU  c                 C   s4   t | jtr| jn| jgD ]| _tf t|  qd S r   )r6  r@  rr   r	  vars)rT  r0   r0   r1   main(  s     rW  __main__)C__doc__rL  rH   r~   r  r   sysr.  r;  pathlibr   pandasr,   rA   torch.nnr8  Ztorch.utils.mobile_optimizerr   __file__r?  rR  parentsrO  rK   pathappendrelpathcwdmodels.commonr   models.experimentalr   models.yolor   Zutils.activationsr   utils.datasetsr   utils.generalr	   r
   r   r   r   r   r   r   r   utils.torch_utilsr   r2   rX   rx   r   r   r   r   r   r  r  r  no_gradr	  rU  rW  __name__rT  r0   r0   r0   r1   <module>   s   +
,,7      
, %m
