Update `export.py` with Detect, Validate usages (#6280)

export.py

@@ -82,6 +82,7 @@ def export_torchscript(model, im, file, optimize, prefix=colorstr('TorchScript:'
             ts.save(str(f), _extra_files=extra_files)
 
         LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
+        return f
     except Exception as e:
         LOGGER.info(f'{prefix} export failure: {e}')
 
@@ -125,7 +126,7 @@ def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorst
             except Exception as e:
                 LOGGER.info(f'{prefix} simplifier failure: {e}')
         LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
-        LOGGER.info(f"{prefix} run --dynamic ONNX model inference with: 'python detect.py --weights {f}'")
+        return f
     except Exception as e:
         LOGGER.info(f'{prefix} export failure: {e}')
 
@@ -143,13 +144,13 @@ def export_openvino(model, im, file, prefix=colorstr('OpenVINO:')):
         subprocess.check_output(cmd, shell=True)
 
         LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
+        return f
     except Exception as e:
         LOGGER.info(f'\n{prefix} export failure: {e}')
 
 
 def export_coreml(model, im, file, prefix=colorstr('CoreML:')):
     # YOLOv5 CoreML export
-    ct_model = None
     try:
         check_requirements(('coremltools',))
         import coremltools as ct
@@ -162,10 +163,10 @@ def export_coreml(model, im, file, prefix=colorstr('CoreML:')):
         ct_model.save(f)
 
         LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
+        return ct_model, f
     except Exception as e:
         LOGGER.info(f'\n{prefix} export failure: {e}')
-
+        return None, None
-    return ct_model
 
 
 def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=False, prefix=colorstr('TensorRT:')):
@@ -216,7 +217,7 @@ def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=F
         with builder.build_engine(network, config) as engine, open(f, 'wb') as t:
             t.write(engine.serialize())
         LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
-
+        return f
     except Exception as e:
         LOGGER.info(f'\n{prefix} export failure: {e}')
 
@@ -225,7 +226,6 @@ def export_saved_model(model, im, file, dynamic,
                        tf_nms=False, agnostic_nms=False, topk_per_class=100, topk_all=100, iou_thres=0.45,
                        conf_thres=0.25, prefix=colorstr('TensorFlow SavedModel:')):
     # YOLOv5 TensorFlow SavedModel export
-    keras_model = None
     try:
         import tensorflow as tf
         from tensorflow import keras
@@ -247,10 +247,10 @@ def export_saved_model(model, im, file, dynamic,
         keras_model.save(f, save_format='tf')
 
         LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
+        return keras_model, f
     except Exception as e:
         LOGGER.info(f'\n{prefix} export failure: {e}')
-
+        return None, None
-    return keras_model
 
 
 def export_pb(keras_model, im, file, prefix=colorstr('TensorFlow GraphDef:')):
@@ -269,6 +269,7 @@ def export_pb(keras_model, im, file, prefix=colorstr('TensorFlow GraphDef:')):
         tf.io.write_graph(graph_or_graph_def=frozen_func.graph, logdir=str(f.parent), name=f.name, as_text=False)
 
         LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
+        return f
     except Exception as e:
         LOGGER.info(f'\n{prefix} export failure: {e}')
 
@@ -300,7 +301,7 @@ def export_tflite(keras_model, im, file, int8, data, ncalib, prefix=colorstr('Te
         tflite_model = converter.convert()
         open(f, "wb").write(tflite_model)
         LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
-
+        return f
     except Exception as e:
         LOGGER.info(f'\n{prefix} export failure: {e}')
 
@@ -328,6 +329,7 @@ def export_edgetpu(keras_model, im, file, prefix=colorstr('Edge TPU:')):
         subprocess.run(cmd, shell=True, check=True)
 
         LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
+        return f
     except Exception as e:
         LOGGER.info(f'\n{prefix} export failure: {e}')
 
@@ -364,6 +366,7 @@ def export_tfjs(keras_model, im, file, prefix=colorstr('TensorFlow.js:')):
             j.write(subst)
 
         LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
+        return f
     except Exception as e:
         LOGGER.info(f'\n{prefix} export failure: {e}')
 
@@ -431,15 +434,15 @@ def run(data=ROOT / 'data/coco128.yaml',  # 'dataset.yaml path'
 
     # Exports
     if 'torchscript' in include:
-        export_torchscript(model, im, file, optimize)
+        f = export_torchscript(model, im, file, optimize)
     if 'engine' in include:  # TensorRT required before ONNX
-        export_engine(model, im, file, train, half, simplify, workspace, verbose)
+        f = export_engine(model, im, file, train, half, simplify, workspace, verbose)
     if ('onnx' in include) or ('openvino' in include):  # OpenVINO requires ONNX
-        export_onnx(model, im, file, opset, train, dynamic, simplify)
+        f = export_onnx(model, im, file, opset, train, dynamic, simplify)
     if 'openvino' in include:
-        export_openvino(model, im, file)
+        f = export_openvino(model, im, file)
     if 'coreml' in include:
-        export_coreml(model, im, file)
+        _, f = export_coreml(model, im, file)
 
     # TensorFlow Exports
     if any(tf_exports):
@@ -447,22 +450,26 @@ def run(data=ROOT / 'data/coco128.yaml',  # 'dataset.yaml path'
         if int8 or edgetpu:  # TFLite --int8 bug https://github.com/ultralytics/yolov5/issues/5707
             check_requirements(('flatbuffers==1.12',))  # required before `import tensorflow`
         assert not (tflite and tfjs), 'TFLite and TF.js models must be exported separately, please pass only one type.'
-        model = export_saved_model(model, im, file, dynamic, tf_nms=nms or agnostic_nms or tfjs,
+        model, f = export_saved_model(model, im, file, dynamic, tf_nms=nms or agnostic_nms or tfjs,
-                                   agnostic_nms=agnostic_nms or tfjs, topk_per_class=topk_per_class, topk_all=topk_all,
+                                      agnostic_nms=agnostic_nms or tfjs, topk_per_class=topk_per_class,
-                                   conf_thres=conf_thres, iou_thres=iou_thres)  # keras model
+                                      topk_all=topk_all,
+                                      conf_thres=conf_thres, iou_thres=iou_thres)  # keras model
         if pb or tfjs:  # pb prerequisite to tfjs
-            export_pb(model, im, file)
+            f = export_pb(model, im, file)
         if tflite or edgetpu:
-            export_tflite(model, im, file, int8=int8 or edgetpu, data=data, ncalib=100)
+            f = export_tflite(model, im, file, int8=int8 or edgetpu, data=data, ncalib=100)
         if edgetpu:
-            export_edgetpu(model, im, file)
+            f = export_edgetpu(model, im, file)
         if tfjs:
-            export_tfjs(model, im, file)
+            f = export_tfjs(model, im, file)
 
     # Finish
     LOGGER.info(f'\nExport complete ({time.time() - t:.2f}s)'
                 f"\nResults saved to {colorstr('bold', file.parent.resolve())}"
-                f'\nVisualize with https://netron.app')
+                f"\nVisualize with https://netron.app"
+                f"\nDetect with `python detect.py --weights {f}`"
+                f" or `model = torch.hub.load('ultralytics/yolov5', 'custom', '{f}')"
+                f"\nValidate with `python val.py --weights {f}`")
 
 
 def parse_opt():
@@ -490,7 +497,7 @@ def parse_opt():
     parser.add_argument('--conf-thres', type=float, default=0.25, help='TF.js NMS: confidence threshold')
     parser.add_argument('--include', nargs='+',
                         default=['torchscript', 'onnx'],
-                        help='available formats are (torchscript, onnx, engine, coreml, saved_model, pb, tflite, tfjs)')
+                        help='torchscript, onnx, openvino, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs')
     opt = parser.parse_args()
     print_args(FILE.stem, opt)
     return opt