Add TensorFlow formats to `export.py` (#4479)

* Initial commit

* Remove unused export_torchscript return

* ROOT variable

* Add prefix to fcn arg

* fix ROOT

* check_yaml into run()

* interim fixes

* imgsz=(320, 320)

* Hardcode tf_raw_resize False

* Finish opt elimination

* Update representative_dataset_gen()

* Update export.py with TF methods

* SiLU and GraphDef fixes

* file_size() directory handling feature

* export fixes

* add lambda: to representative_dataset

* Detect training False default

* Fuse false for TF models

* Embed agnostic NMS arguments

* Remove lambda

* TensorFlow.js export success

* Add pb to Usage

* Add *_tfjs_model/ to ignore files

* prepend YOLOv5 to function headers

* Remove end --- comments

* parameterize tfjs export pb file

* update run() data default /ROOT

* update --include help

* update imports

* return ct_model

* Consolidate TFLite export

* pb prerequisite to tfjs

* TF modules CamelCase

* Remove exports from tf.py and cleanup

* pass agnostic NMS arguments

* CI

* CI

* ignore *_web_model/

* Add tensorflow to CI dependencies

* CI tensorflow-cpu

* Update requirements.txt

* Remove tensorflow check_requirement

* CI coreml tfjs

* export only onnx torchscript

* reorder exports torchscript first

.dockerignore

@@ -22,6 +22,7 @@ data/samples/*
 **/*.h5
 **/*.pb
 *_saved_model/
+*_web_model/
 
 # Below Copied From .gitignore -----------------------------------------------------------------------------------------
 # Below Copied From .gitignore -----------------------------------------------------------------------------------------

.github/workflows/ci-testing.yml

@@ -48,7 +48,7 @@ jobs:
         run: |
           python -m pip install --upgrade pip
           pip install -qr requirements.txt -f https://download.pytorch.org/whl/cpu/torch_stable.html
-          pip install -q onnx onnx-simplifier coremltools  # for export
+          pip install -q onnx tensorflow-cpu  # for export
           python --version
           pip --version
           pip list
@@ -75,6 +75,7 @@ jobs:
           python val.py --img 128 --batch 16 --weights runs/train/exp/weights/last.pt --device $di
 
           python hubconf.py  # hub
-          python models/yolo.py --cfg ${{ matrix.model }}.yaml  # inspect
-          python export.py --img 128 --batch 1 --weights ${{ matrix.model }}.pt --include onnx torchscript  # export
+          python models/yolo.py --cfg ${{ matrix.model }}.yaml  # build PyTorch model
+          python models/tf.py --weights ${{ matrix.model }}.pt  # build TensorFlow model
+          python export.py --img 128 --batch 1 --weights ${{ matrix.model }}.pt --include torchscript onnx  # export
         shell: bash

.gitignore

@@ -52,6 +52,7 @@ VOC/
 *.tflite
 *.h5
 *_saved_model/
+*_web_model/
 darknet53.conv.74
 yolov3-tiny.conv.15
 

detect.py

@@ -253,7 +253,7 @@ def run(weights='yolov5s.pt',  # model.pt path(s)
 
 def parse_opt():
     parser = argparse.ArgumentParser()
-    parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model.pt path(s)')
+    parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model path(s)')
     parser.add_argument('--source', type=str, default='data/images', help='file/dir/URL/glob, 0 for webcam')
     parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640], help='inference size h,w')
     parser.add_argument('--conf-thres', type=float, default=0.25, help='confidence threshold')

export.py

@@ -1,12 +1,28 @@
 # YOLOv5 🚀 by Ultralytics, GPL-3.0 license
 """
-Export a PyTorch model to TorchScript, ONNX, CoreML formats
+Export a YOLOv5 PyTorch model to TorchScript, ONNX, CoreML, TensorFlow (saved_model, pb, TFLite, TF.js,) formats
+TensorFlow exports authored by https://github.com/zldrobit
 
 Usage:
-    $ python path/to/export.py --weights yolov5s.pt --img 640 --batch 1
+    $ python path/to/export.py --weights yolov5s.pt --include torchscript onnx coreml saved_model pb tflite tfjs
+
+Inference:
+    $ python path/to/detect.py --weights yolov5s.pt
+                                         yolov5s.onnx  (must export with --dynamic)
+                                         yolov5s_saved_model
+                                         yolov5s.pb
+                                         yolov5s.tflite
+
+TensorFlow.js:
+    $ # Edit yolov5s_web_model/model.json to sort Identity* in ascending order
+    $ 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
 """
 
 import argparse
+import subprocess
 import sys
 import time
 from pathlib import Path
@@ -16,40 +32,42 @@ import torch.nn as nn
 from torch.utils.mobile_optimizer import optimize_for_mobile
 
 FILE = Path(__file__).resolve()
-sys.path.append(FILE.parents[0].as_posix())  # add yolov5/ to path
+ROOT = FILE.parents[0]  # yolov5/ dir
+sys.path.append(ROOT.as_posix())  # add yolov5/ to path
 
 from models.common import Conv
-from models.yolo import Detect
 from models.experimental import attempt_load
-from utils.activations import Hardswish, SiLU
-from utils.general import colorstr, check_img_size, check_requirements, file_size, set_logging
+from models.yolo import Detect
+from utils.activations import SiLU
+from utils.datasets import LoadImages
+from utils.general import colorstr, check_dataset, check_img_size, check_requirements, file_size, set_logging
 from utils.torch_utils import select_device
 
 
-def export_torchscript(model, img, file, optimize):
-    # TorchScript model export
-    prefix = colorstr('TorchScript:')
+def export_torchscript(model, im, file, optimize, prefix=colorstr('TorchScript:')):
+    # YOLOv5 TorchScript model export
     try:
         print(f'\n{prefix} starting export with torch {torch.__version__}...')
         f = file.with_suffix('.torchscript.pt')
-        ts = torch.jit.trace(model, img, strict=False)
+
+        ts = torch.jit.trace(model, im, strict=False)
         (optimize_for_mobile(ts) if optimize else ts).save(f)
+
         print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
-        return ts
     except Exception as e:
         print(f'{prefix} export failure: {e}')
 
 
-def export_onnx(model, img, file, opset, train, dynamic, simplify):
-    # ONNX model export
-    prefix = colorstr('ONNX:')
+def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorstr('ONNX:')):
+    # YOLOv5 ONNX export
     try:
         check_requirements(('onnx',))
         import onnx
 
         print(f'\n{prefix} starting export with onnx {onnx.__version__}...')
         f = file.with_suffix('.onnx')
-        torch.onnx.export(model, img, f, verbose=False, opset_version=opset,
+
+        torch.onnx.export(model, im, f, verbose=False, opset_version=opset,
                           training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.EVAL,
                           do_constant_folding=not train,
                           input_names=['images'],
@@ -73,7 +91,7 @@ def export_onnx(model, img, file, opset, train, dynamic, simplify):
                 model_onnx, check = onnxsim.simplify(
                     model_onnx,
                     dynamic_input_shape=dynamic,
-                    input_shapes={'images': list(img.shape)} if dynamic else None)
+                    input_shapes={'images': list(im.shape)} if dynamic else None)
                 assert check, 'assert check failed'
                 onnx.save(model_onnx, f)
             except Exception as e:
@@ -84,26 +102,131 @@ def export_onnx(model, img, file, opset, train, dynamic, simplify):
         print(f'{prefix} export failure: {e}')
 
 
-def export_coreml(model, img, file):
-    # CoreML model export
-    prefix = colorstr('CoreML:')
+def export_coreml(model, im, file, prefix=colorstr('CoreML:')):
+    # YOLOv5 CoreML export
+    ct_model = None
     try:
         check_requirements(('coremltools',))
         import coremltools as ct
 
         print(f'\n{prefix} starting export with coremltools {ct.__version__}...')
         f = file.with_suffix('.mlmodel')
+
         model.train()  # CoreML exports should be placed in model.train() mode
-        ts = torch.jit.trace(model, img, strict=False)  # TorchScript model
-        model = ct.convert(ts, inputs=[ct.ImageType('image', shape=img.shape, scale=1 / 255.0, bias=[0, 0, 0])])
-        model.save(f)
+        ts = torch.jit.trace(model, im, strict=False)  # TorchScript model
+        ct_model = ct.convert(ts, inputs=[ct.ImageType('image', shape=im.shape, scale=1 / 255.0, bias=[0, 0, 0])])
+        ct_model.save(f)
+
         print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
     except Exception as e:
         print(f'\n{prefix} export failure: {e}')
 
+    return ct_model
+
+
+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 saved_model:')):
+    # YOLOv5 TensorFlow saved_model export
+    keras_model = None
+    try:
+        import tensorflow as tf
+        from tensorflow import keras
+        from models.tf import TFModel, TFDetect
+
+        print(f'\n{prefix} starting export with tensorflow {tf.__version__}...')
+        f = str(file).replace('.pt', '_saved_model')
+        batch_size, ch, *imgsz = list(im.shape)  # BCHW
 
-def run(weights='./yolov5s.pt',  # weights path
-        img_size=(640, 640),  # image (height, width)
+        tf_model = TFModel(cfg=model.yaml, model=model, nc=model.nc, imgsz=imgsz)
+        im = tf.zeros((batch_size, *imgsz, 3))  # BHWC order for TensorFlow
+        y = tf_model.predict(im, tf_nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres)
+        inputs = keras.Input(shape=(*imgsz, 3), batch_size=None if dynamic else batch_size)
+        outputs = tf_model.predict(inputs, tf_nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres)
+        keras_model = keras.Model(inputs=inputs, outputs=outputs)
+        keras_model.summary()
+        keras_model.save(f, save_format='tf')
+
+        print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
+    except Exception as e:
+        print(f'\n{prefix} export failure: {e}')
+
+    return keras_model
+
+
+def export_pb(keras_model, im, file, prefix=colorstr('TensorFlow GraphDef:')):
+    # YOLOv5 TensorFlow GraphDef *.pb export https://github.com/leimao/Frozen_Graph_TensorFlow
+    try:
+        import tensorflow as tf
+        from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
+
+        print(f'\n{prefix} starting export with tensorflow {tf.__version__}...')
+        f = file.with_suffix('.pb')
+
+        m = tf.function(lambda x: keras_model(x))  # full model
+        m = m.get_concrete_function(tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype))
+        frozen_func = convert_variables_to_constants_v2(m)
+        frozen_func.graph.as_graph_def()
+        tf.io.write_graph(graph_or_graph_def=frozen_func.graph, logdir=str(f.parent), name=f.name, as_text=False)
+
+        print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
+    except Exception as e:
+        print(f'\n{prefix} export failure: {e}')
+
+
+def export_tflite(keras_model, im, file, tfl_int8, data, ncalib, prefix=colorstr('TensorFlow Lite:')):
+    # YOLOv5 TensorFlow Lite export
+    try:
+        import tensorflow as tf
+        from models.tf import representative_dataset_gen
+
+        print(f'\n{prefix} starting export with tensorflow {tf.__version__}...')
+        batch_size, ch, *imgsz = list(im.shape)  # BCHW
+        f = file.with_suffix('.tflite')
+
+        converter = tf.lite.TFLiteConverter.from_keras_model(keras_model)
+        converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS]
+        converter.optimizations = [tf.lite.Optimize.DEFAULT]
+        if tfl_int8:
+            dataset = LoadImages(check_dataset(data)['train'], img_size=imgsz, auto=False)  # representative data
+            converter.representative_dataset = lambda: representative_dataset_gen(dataset, ncalib)
+            converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
+            converter.inference_input_type = tf.uint8  # or tf.int8
+            converter.inference_output_type = tf.uint8  # or tf.int8
+            converter.experimental_new_quantizer = False
+            f = str(file).replace('.pt', '-int8.tflite')
+
+        tflite_model = converter.convert()
+        open(f, "wb").write(tflite_model)
+        print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
+
+    except Exception as e:
+        print(f'\n{prefix} export failure: {e}')
+
+
+def export_tfjs(keras_model, im, file, prefix=colorstr('TensorFlow.js:')):
+    # YOLOv5 TensorFlow.js export
+    try:
+        check_requirements(('tensorflowjs',))
+        import tensorflowjs as tfjs
+
+        print(f'\n{prefix} starting export with tensorflowjs {tfjs.__version__}...')
+        f = str(file).replace('.pt', '_web_model')  # js dir
+        f_pb = file.with_suffix('.pb')  # *.pb path
+
+        cmd = f"tensorflowjs_converter --input_format=tf_frozen_model " \
+              f"--output_node_names='Identity,Identity_1,Identity_2,Identity_3' {f_pb} {f}"
+        subprocess.run(cmd, shell=True)
+
+        print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
+    except Exception as e:
+        print(f'\n{prefix} export failure: {e}')
+
+
+@torch.no_grad()
+def run(data=ROOT / 'data/coco128.yaml',  # 'dataset.yaml path'
+        weights=ROOT / 'yolov5s.pt',  # weights path
+        imgsz=(640, 640),  # image (height, width)
         batch_size=1,  # batch size
         device='cpu',  # cuda device, i.e. 0 or 0,1,2,3 or cpu
         include=('torchscript', 'onnx', 'coreml'),  # include formats
@@ -117,29 +240,28 @@ def run(weights='./yolov5s.pt',  # weights path
         ):
     t = time.time()
     include = [x.lower() for x in include]
-    img_size *= 2 if len(img_size) == 1 else 1  # expand
+    tf_exports = list(x in include for x in ('saved_model', 'pb', 'tflite', 'tfjs'))  # TensorFlow exports
+    imgsz *= 2 if len(imgsz) == 1 else 1  # expand
     file = Path(weights)
 
     # Load PyTorch model
     device = select_device(device)
     assert not (device.type == 'cpu' and half), '--half only compatible with GPU export, i.e. use --device 0'
-    model = attempt_load(weights, map_location=device)  # load FP32 model
-    names = model.names
+    model = attempt_load(weights, map_location=device, inplace=True, fuse=not any(tf_exports))  # load FP32 model
+    nc, names = model.nc, model.names  # number of classes, class names
 
     # Input
     gs = int(max(model.stride))  # grid size (max stride)
-    img_size = [check_img_size(x, gs) for x in img_size]  # verify img_size are gs-multiples
-    img = torch.zeros(batch_size, 3, *img_size).to(device)  # image size(1,3,320,192) iDetection
+    imgsz = [check_img_size(x, gs) for x in imgsz]  # verify img_size are gs-multiples
+    im = torch.zeros(batch_size, 3, *imgsz).to(device)  # image size(1,3,320,192) BCHW iDetection
 
     # Update model
     if half:
-        img, model = img.half(), model.half()  # to FP16
+        im, model = im.half(), model.half()  # to FP16
     model.train() if train else model.eval()  # training mode = no Detect() layer grid construction
     for k, m in model.named_modules():
         if isinstance(m, Conv):  # assign export-friendly activations
-            if isinstance(m.act, nn.Hardswish):
-                m.act = Hardswish()
-            elif isinstance(m.act, nn.SiLU):
+            if isinstance(m.act, nn.SiLU):
                 m.act = SiLU()
         elif isinstance(m, Detect):
             m.inplace = inplace
@@ -147,16 +269,28 @@ def run(weights='./yolov5s.pt',  # weights path
             # m.forward = m.forward_export  # assign forward (optional)
 
     for _ in range(2):
-        y = model(img)  # dry runs
+        y = model(im)  # dry runs
     print(f"\n{colorstr('PyTorch:')} starting from {weights} ({file_size(weights):.1f} MB)")
 
     # Exports
     if 'torchscript' in include:
-        export_torchscript(model, img, file, optimize)
+        export_torchscript(model, im, file, optimize)
     if 'onnx' in include:
-        export_onnx(model, img, file, opset, train, dynamic, simplify)
+        export_onnx(model, im, file, opset, train, dynamic, simplify)
     if 'coreml' in include:
-        export_coreml(model, img, file)
+        export_coreml(model, im, file)
+
+    # TensorFlow Exports
+    if any(tf_exports):
+        pb, tflite, tfjs = tf_exports[1:]
+        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=tfjs, agnostic_nms=tfjs)  # keras model
+        if pb or tfjs:  # pb prerequisite to tfjs
+            export_pb(model, im, file)
+        if tflite:
+            export_tflite(model, im, file, tfl_int8=False, data=data, ncalib=100)
+        if tfjs:
+            export_tfjs(model, im, file)
 
     # Finish
     print(f'\nExport complete ({time.time() - t:.2f}s)'
@@ -166,18 +300,21 @@ def run(weights='./yolov5s.pt',  # weights path
 
 def parse_opt():
     parser = argparse.ArgumentParser()
-    parser.add_argument('--weights', type=str, default='./yolov5s.pt', help='weights path')
-    parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image (height, width)')
+    parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path')
+    parser.add_argument('--weights', type=str, default=ROOT / 'yolov5s.pt', help='weights path')
+    parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640, 640], help='image (h, w)')
     parser.add_argument('--batch-size', type=int, default=1, help='batch size')
     parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
-    parser.add_argument('--include', nargs='+', default=['torchscript', 'onnx', 'coreml'], help='include formats')
     parser.add_argument('--half', action='store_true', help='FP16 half-precision export')
     parser.add_argument('--inplace', action='store_true', help='set YOLOv5 Detect() inplace=True')
     parser.add_argument('--train', action='store_true', help='model.train() mode')
     parser.add_argument('--optimize', action='store_true', help='TorchScript: optimize for mobile')
-    parser.add_argument('--dynamic', action='store_true', help='ONNX: dynamic axes')
+    parser.add_argument('--dynamic', action='store_true', help='ONNX/TF: dynamic axes')
     parser.add_argument('--simplify', action='store_true', help='ONNX: simplify model')
     parser.add_argument('--opset', type=int, default=13, help='ONNX: opset version')
+    parser.add_argument('--include', nargs='+',
+                        default=['torchscript', 'onnx'],
+                        help='available formats are (torchscript, onnx, coreml, saved_model, pb, tflite, tfjs)')
     opt = parser.parse_args()
     return opt
 

models/tf.py

@@ -1,67 +1,44 @@
 # YOLOv5 🚀 by Ultralytics, GPL-3.0 license
 """
-TensorFlow/Keras and TFLite versions of YOLOv5
+TensorFlow, Keras and TFLite versions of YOLOv5
 Authored by https://github.com/zldrobit in PR https://github.com/ultralytics/yolov5/pull/1127
 
 Usage:
-    $ python models/tf.py --weights yolov5s.pt --cfg yolov5s.yaml
-
-Export int8 TFLite models:
-    $ python models/tf.py --weights yolov5s.pt --cfg models/yolov5s.yaml --tfl-int8 \
-        --source path/to/images/ --ncalib 100
-
-Detection:
-    $ python detect.py --weights yolov5s.pb          --img 320
-    $ python detect.py --weights yolov5s_saved_model --img 320
-    $ python detect.py --weights yolov5s-fp16.tflite --img 320
-    $ python detect.py --weights yolov5s-int8.tflite --img 320 --tfl-int8
-
-For TensorFlow.js:
-    $ python models/tf.py --weights yolov5s.pt --cfg models/yolov5s.yaml --img 320 --tf-nms --agnostic-nms
-    $ pip install tensorflowjs
-    $ tensorflowjs_converter \
-          --input_format=tf_frozen_model \
-          --output_node_names='Identity,Identity_1,Identity_2,Identity_3' \
-          yolov5s.pb \
-          web_model
-    $ # Edit web_model/model.json to sort Identity* in ascending order
-    $ cd .. && git clone https://github.com/zldrobit/tfjs-yolov5-example.git && cd tfjs-yolov5-example
-    $ npm install
-    $ ln -s ../../yolov5/web_model public/web_model
-    $ npm start
+    $ python models/tf.py --weights yolov5s.pt
+
+Export:
+    $ python path/to/export.py --weights yolov5s.pt --include saved_model pb tflite tfjs
 """
 
 import argparse
 import logging
-import os
 import sys
-import traceback
 from copy import deepcopy
 from pathlib import Path
 
-sys.path.append('./')  # to run '$ python *.py' files in subdirectories
+FILE = Path(__file__).resolve()
+ROOT = FILE.parents[1]  # yolov5/ dir
+sys.path.append(ROOT.as_posix())  # add yolov5/ to path
 
 import numpy as np
 import tensorflow as tf
 import torch
 import torch.nn as nn
-import yaml
 from tensorflow import keras
-from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
 
 from models.common import Conv, Bottleneck, SPP, DWConv, Focus, BottleneckCSP, Concat, autopad, C3
 from models.experimental import MixConv2d, CrossConv, attempt_load
 from models.yolo import Detect
-from utils.datasets import LoadImages
-from utils.general import check_dataset, check_yaml, make_divisible
+from utils.general import colorstr, make_divisible, set_logging
+from utils.activations import SiLU
 
-logger = logging.getLogger(__name__)
+LOGGER = logging.getLogger(__name__)
 
 
-class tf_BN(keras.layers.Layer):
+class TFBN(keras.layers.Layer):
     # TensorFlow BatchNormalization wrapper
     def __init__(self, w=None):
-        super(tf_BN, self).__init__()
+        super(TFBN, self).__init__()
         self.bn = keras.layers.BatchNormalization(
             beta_initializer=keras.initializers.Constant(w.bias.numpy()),
             gamma_initializer=keras.initializers.Constant(w.weight.numpy()),
@@ -73,20 +50,20 @@ class tf_BN(keras.layers.Layer):
         return self.bn(inputs)
 
 
-class tf_Pad(keras.layers.Layer):
+class TFPad(keras.layers.Layer):
     def __init__(self, pad):
-        super(tf_Pad, self).__init__()
+        super(TFPad, self).__init__()
         self.pad = tf.constant([[0, 0], [pad, pad], [pad, pad], [0, 0]])
 
     def call(self, inputs):
         return tf.pad(inputs, self.pad, mode='constant', constant_values=0)
 
 
-class tf_Conv(keras.layers.Layer):
+class TFConv(keras.layers.Layer):
     # Standard convolution
     def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True, w=None):
         # ch_in, ch_out, weights, kernel, stride, padding, groups
-        super(tf_Conv, self).__init__()
+        super(TFConv, self).__init__()
         assert g == 1, "TF v2.2 Conv2D does not support 'groups' argument"
         assert isinstance(k, int), "Convolution with multiple kernels are not allowed."
         # TensorFlow convolution padding is inconsistent with PyTorch (e.g. k=3 s=2 'SAME' padding)
@@ -95,27 +72,29 @@ class tf_Conv(keras.layers.Layer):
         conv = keras.layers.Conv2D(
             c2, k, s, 'SAME' if s == 1 else 'VALID', use_bias=False,
             kernel_initializer=keras.initializers.Constant(w.conv.weight.permute(2, 3, 1, 0).numpy()))
-        self.conv = conv if s == 1 else keras.Sequential([tf_Pad(autopad(k, p)), conv])
-        self.bn = tf_BN(w.bn) if hasattr(w, 'bn') else tf.identity
+        self.conv = conv if s == 1 else keras.Sequential([TFPad(autopad(k, p)), conv])
+        self.bn = TFBN(w.bn) if hasattr(w, 'bn') else tf.identity
 
         # YOLOv5 activations
         if isinstance(w.act, nn.LeakyReLU):
             self.act = (lambda x: keras.activations.relu(x, alpha=0.1)) if act else tf.identity
         elif isinstance(w.act, nn.Hardswish):
             self.act = (lambda x: x * tf.nn.relu6(x + 3) * 0.166666667) if act else tf.identity
-        elif isinstance(w.act, nn.SiLU):
+        elif isinstance(w.act, (nn.SiLU, SiLU)):
             self.act = (lambda x: keras.activations.swish(x)) if act else tf.identity
+        else:
+            raise Exception(f'no matching TensorFlow activation found for {w.act}')
 
     def call(self, inputs):
         return self.act(self.bn(self.conv(inputs)))
 
 
-class tf_Focus(keras.layers.Layer):
+class TFFocus(keras.layers.Layer):
     # Focus wh information into c-space
     def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True, w=None):
         # ch_in, ch_out, kernel, stride, padding, groups
-        super(tf_Focus, self).__init__()
-        self.conv = tf_Conv(c1 * 4, c2, k, s, p, g, act, w.conv)
+        super(TFFocus, self).__init__()
+        self.conv = TFConv(c1 * 4, c2, k, s, p, g, act, w.conv)
 
     def call(self, inputs):  # x(b,w,h,c) -> y(b,w/2,h/2,4c)
         # inputs = inputs / 255.  # normalize 0-255 to 0-1
@@ -125,23 +104,23 @@ class tf_Focus(keras.layers.Layer):
                                     inputs[:, 1::2, 1::2, :]], 3))
 
 
-class tf_Bottleneck(keras.layers.Layer):
+class TFBottleneck(keras.layers.Layer):
     # Standard bottleneck
     def __init__(self, c1, c2, shortcut=True, g=1, e=0.5, w=None):  # ch_in, ch_out, shortcut, groups, expansion
-        super(tf_Bottleneck, self).__init__()
+        super(TFBottleneck, self).__init__()
         c_ = int(c2 * e)  # hidden channels
-        self.cv1 = tf_Conv(c1, c_, 1, 1, w=w.cv1)
-        self.cv2 = tf_Conv(c_, c2, 3, 1, g=g, w=w.cv2)
+        self.cv1 = TFConv(c1, c_, 1, 1, w=w.cv1)
+        self.cv2 = TFConv(c_, c2, 3, 1, g=g, w=w.cv2)
         self.add = shortcut and c1 == c2
 
     def call(self, inputs):
         return inputs + self.cv2(self.cv1(inputs)) if self.add else self.cv2(self.cv1(inputs))
 
 
-class tf_Conv2d(keras.layers.Layer):
+class TFConv2d(keras.layers.Layer):
     # Substitution for PyTorch nn.Conv2D
     def __init__(self, c1, c2, k, s=1, g=1, bias=True, w=None):
-        super(tf_Conv2d, self).__init__()
+        super(TFConv2d, self).__init__()
         assert g == 1, "TF v2.2 Conv2D does not support 'groups' argument"
         self.conv = keras.layers.Conv2D(
             c2, k, s, 'VALID', use_bias=bias,
@@ -152,19 +131,19 @@ class tf_Conv2d(keras.layers.Layer):
         return self.conv(inputs)
 
 
-class tf_BottleneckCSP(keras.layers.Layer):
+class TFBottleneckCSP(keras.layers.Layer):
     # CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks
     def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5, w=None):
         # ch_in, ch_out, number, shortcut, groups, expansion
-        super(tf_BottleneckCSP, self).__init__()
+        super(TFBottleneckCSP, self).__init__()
         c_ = int(c2 * e)  # hidden channels
-        self.cv1 = tf_Conv(c1, c_, 1, 1, w=w.cv1)
-        self.cv2 = tf_Conv2d(c1, c_, 1, 1, bias=False, w=w.cv2)
-        self.cv3 = tf_Conv2d(c_, c_, 1, 1, bias=False, w=w.cv3)
-        self.cv4 = tf_Conv(2 * c_, c2, 1, 1, w=w.cv4)
-        self.bn = tf_BN(w.bn)
+        self.cv1 = TFConv(c1, c_, 1, 1, w=w.cv1)
+        self.cv2 = TFConv2d(c1, c_, 1, 1, bias=False, w=w.cv2)
+        self.cv3 = TFConv2d(c_, c_, 1, 1, bias=False, w=w.cv3)
+        self.cv4 = TFConv(2 * c_, c2, 1, 1, w=w.cv4)
+        self.bn = TFBN(w.bn)
         self.act = lambda x: keras.activations.relu(x, alpha=0.1)
-        self.m = keras.Sequential([tf_Bottleneck(c_, c_, shortcut, g, e=1.0, w=w.m[j]) for j in range(n)])
+        self.m = keras.Sequential([TFBottleneck(c_, c_, shortcut, g, e=1.0, w=w.m[j]) for j in range(n)])
 
     def call(self, inputs):
         y1 = self.cv3(self.m(self.cv1(inputs)))
@@ -172,28 +151,28 @@ class tf_BottleneckCSP(keras.layers.Layer):
         return self.cv4(self.act(self.bn(tf.concat((y1, y2), axis=3))))
 
 
-class tf_C3(keras.layers.Layer):
+class TFC3(keras.layers.Layer):
     # CSP Bottleneck with 3 convolutions
     def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5, w=None):
         # ch_in, ch_out, number, shortcut, groups, expansion
-        super(tf_C3, self).__init__()
+        super(TFC3, self).__init__()
         c_ = int(c2 * e)  # hidden channels
-        self.cv1 = tf_Conv(c1, c_, 1, 1, w=w.cv1)
-        self.cv2 = tf_Conv(c1, c_, 1, 1, w=w.cv2)
-        self.cv3 = tf_Conv(2 * c_, c2, 1, 1, w=w.cv3)
-        self.m = keras.Sequential([tf_Bottleneck(c_, c_, shortcut, g, e=1.0, w=w.m[j]) for j in range(n)])
+        self.cv1 = TFConv(c1, c_, 1, 1, w=w.cv1)
+        self.cv2 = TFConv(c1, c_, 1, 1, w=w.cv2)
+        self.cv3 = TFConv(2 * c_, c2, 1, 1, w=w.cv3)
+        self.m = keras.Sequential([TFBottleneck(c_, c_, shortcut, g, e=1.0, w=w.m[j]) for j in range(n)])
 
     def call(self, inputs):
         return self.cv3(tf.concat((self.m(self.cv1(inputs)), self.cv2(inputs)), axis=3))
 
 
-class tf_SPP(keras.layers.Layer):
+class TFSPP(keras.layers.Layer):
     # Spatial pyramid pooling layer used in YOLOv3-SPP
     def __init__(self, c1, c2, k=(5, 9, 13), w=None):
-        super(tf_SPP, self).__init__()
+        super(TFSPP, self).__init__()
         c_ = c1 // 2  # hidden channels
-        self.cv1 = tf_Conv(c1, c_, 1, 1, w=w.cv1)
-        self.cv2 = tf_Conv(c_ * (len(k) + 1), c2, 1, 1, w=w.cv2)
+        self.cv1 = TFConv(c1, c_, 1, 1, w=w.cv1)
+        self.cv2 = TFConv(c_ * (len(k) + 1), c2, 1, 1, w=w.cv2)
         self.m = [keras.layers.MaxPool2D(pool_size=x, strides=1, padding='SAME') for x in k]
 
     def call(self, inputs):
@@ -201,9 +180,9 @@ class tf_SPP(keras.layers.Layer):
         return self.cv2(tf.concat([x] + [m(x) for m in self.m], 3))
 
 
-class tf_Detect(keras.layers.Layer):
-    def __init__(self, nc=80, anchors=(), ch=(), w=None):  # detection layer
-        super(tf_Detect, self).__init__()
+class TFDetect(keras.layers.Layer):
+    def __init__(self, nc=80, anchors=(), ch=(), imgsz=(640, 640), w=None):  # detection layer
+        super(TFDetect, self).__init__()
         self.stride = tf.convert_to_tensor(w.stride.numpy(), dtype=tf.float32)
         self.nc = nc  # number of classes
         self.no = nc + 5  # number of outputs per anchor
@@ -213,22 +192,20 @@ class tf_Detect(keras.layers.Layer):
         self.anchors = tf.convert_to_tensor(w.anchors.numpy(), dtype=tf.float32)
         self.anchor_grid = tf.reshape(tf.convert_to_tensor(w.anchor_grid.numpy(), dtype=tf.float32),
                                       [self.nl, 1, -1, 1, 2])
-        self.m = [tf_Conv2d(x, self.no * self.na, 1, w=w.m[i]) for i, x in enumerate(ch)]
-        self.export = False  # onnx export
-        self.training = True  # set to False after building model
+        self.m = [TFConv2d(x, self.no * self.na, 1, w=w.m[i]) for i, x in enumerate(ch)]
+        self.training = False  # set to False after building model
+        self.imgsz = imgsz
         for i in range(self.nl):
-            ny, nx = opt.img_size[0] // self.stride[i], opt.img_size[1] // self.stride[i]
+            ny, nx = self.imgsz[0] // self.stride[i], self.imgsz[1] // self.stride[i]
             self.grid[i] = self._make_grid(nx, ny)
 
     def call(self, inputs):
-        # x = x.copy()  # for profiling
         z = []  # inference output
-        self.training |= self.export
         x = []
         for i in range(self.nl):
             x.append(self.m[i](inputs[i]))
             # x(bs,20,20,255) to x(bs,3,20,20,85)
-            ny, nx = opt.img_size[0] // self.stride[i], opt.img_size[1] // self.stride[i]
+            ny, nx = self.imgsz[0] // self.stride[i], self.imgsz[1] // self.stride[i]
             x[i] = tf.transpose(tf.reshape(x[i], [-1, ny * nx, self.na, self.no]), [0, 2, 1, 3])
 
             if not self.training:  # inference
@@ -236,8 +213,8 @@ class tf_Detect(keras.layers.Layer):
                 xy = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i]  # xy
                 wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]
                 # Normalize xywh to 0-1 to reduce calibration error
-                xy /= tf.constant([[opt.img_size[1], opt.img_size[0]]], dtype=tf.float32)
-                wh /= tf.constant([[opt.img_size[1], opt.img_size[0]]], dtype=tf.float32)
+                xy /= tf.constant([[self.imgsz[1], self.imgsz[0]]], dtype=tf.float32)
+                wh /= tf.constant([[self.imgsz[1], self.imgsz[0]]], dtype=tf.float32)
                 y = tf.concat([xy, wh, y[..., 4:]], -1)
                 z.append(tf.reshape(y, [-1, 3 * ny * nx, self.no]))
 
@@ -251,25 +228,23 @@ class tf_Detect(keras.layers.Layer):
         return tf.cast(tf.reshape(tf.stack([xv, yv], 2), [1, 1, ny * nx, 2]), dtype=tf.float32)
 
 
-class tf_Upsample(keras.layers.Layer):
-    def __init__(self, size, scale_factor, mode, w=None):
-        super(tf_Upsample, self).__init__()
+class TFUpsample(keras.layers.Layer):
+    def __init__(self, size, scale_factor, mode, w=None):  # warning: all arguments needed including 'w'
+        super(TFUpsample, self).__init__()
         assert scale_factor == 2, "scale_factor must be 2"
+        self.upsample = lambda x: tf.image.resize(x, (x.shape[1] * 2, x.shape[2] * 2), method=mode)
         # self.upsample = keras.layers.UpSampling2D(size=scale_factor, interpolation=mode)
-        if opt.tf_raw_resize:
-            # with default arguments: align_corners=False, half_pixel_centers=False
-            self.upsample = lambda x: tf.raw_ops.ResizeNearestNeighbor(images=x,
-                                                                       size=(x.shape[1] * 2, x.shape[2] * 2))
-        else:
-            self.upsample = lambda x: tf.image.resize(x, (x.shape[1] * 2, x.shape[2] * 2), method=mode)
+        # with default arguments: align_corners=False, half_pixel_centers=False
+        # self.upsample = lambda x: tf.raw_ops.ResizeNearestNeighbor(images=x,
+        #                                                            size=(x.shape[1] * 2, x.shape[2] * 2))
 
     def call(self, inputs):
         return self.upsample(inputs)
 
 
-class tf_Concat(keras.layers.Layer):
+class TFConcat(keras.layers.Layer):
     def __init__(self, dimension=1, w=None):
-        super(tf_Concat, self).__init__()
+        super(TFConcat, self).__init__()
         assert dimension == 1, "convert only NCHW to NHWC concat"
         self.d = 3
 
@@ -277,8 +252,8 @@ class tf_Concat(keras.layers.Layer):
         return tf.concat(inputs, self.d)
 
 
-def parse_model(d, ch, model):  # model_dict, input_channels(3)
-    logger.info('\n%3s%18s%3s%10s  %-40s%-30s' % ('', 'from', 'n', 'params', 'module', 'arguments'))
+def parse_model(d, ch, model, imgsz):  # model_dict, input_channels(3)
+    LOGGER.info('\n%3s%18s%3s%10s  %-40s%-30s' % ('', 'from', 'n', 'params', 'module', 'arguments'))
     anchors, nc, gd, gw = d['anchors'], d['nc'], d['depth_multiple'], d['width_multiple']
     na = (len(anchors[0]) // 2) if isinstance(anchors, list) else anchors  # number of anchors
     no = na * (nc + 5)  # number of outputs = anchors * (classes + 5)
@@ -310,10 +285,11 @@ def parse_model(d, ch, model):  # model_dict, input_channels(3)
             args.append([ch[x + 1] for x in f])
             if isinstance(args[1], int):  # number of anchors
                 args[1] = [list(range(args[1] * 2))] * len(f)
+            args.append(imgsz)
         else:
             c2 = ch[f]
 
-        tf_m = eval('tf_' + m_str.replace('nn.', ''))
+        tf_m = eval('TF' + m_str.replace('nn.', ''))
         m_ = keras.Sequential([tf_m(*args, w=model.model[i][j]) for j in range(n)]) if n > 1 \
             else tf_m(*args, w=model.model[i])  # module
 
@@ -321,16 +297,16 @@ def parse_model(d, ch, model):  # model_dict, input_channels(3)
         t = str(m)[8:-2].replace('__main__.', '')  # module type
         np = sum([x.numel() for x in torch_m_.parameters()])  # number params
         m_.i, m_.f, m_.type, m_.np = i, f, t, np  # attach index, 'from' index, type, number params
-        logger.info('%3s%18s%3s%10.0f  %-40s%-30s' % (i, f, n, np, t, args))  # print
+        LOGGER.info('%3s%18s%3s%10.0f  %-40s%-30s' % (i, f, n, np, t, args))  # print
         save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1)  # append to savelist
         layers.append(m_)
         ch.append(c2)
     return keras.Sequential(layers), sorted(save)
 
 
-class tf_Model():
-    def __init__(self, cfg='yolov5s.yaml', ch=3, nc=None, model=None):  # model, input channels, number of classes
-        super(tf_Model, self).__init__()
+class TFModel:
+    def __init__(self, cfg='yolov5s.yaml', ch=3, nc=None, model=None, imgsz=(640, 640)):  # model, channels, classes
+        super(TFModel, self).__init__()
         if isinstance(cfg, dict):
             self.yaml = cfg  # model dict
         else:  # is *.yaml
@@ -343,9 +319,10 @@ class tf_Model():
         if nc and nc != self.yaml['nc']:
             print('Overriding %s nc=%g with nc=%g' % (cfg, self.yaml['nc'], nc))
             self.yaml['nc'] = nc  # override yaml value
-        self.model, self.savelist = parse_model(deepcopy(self.yaml), ch=[ch], model=model)  # model, savelist, ch_out
+        self.model, self.savelist = parse_model(deepcopy(self.yaml), ch=[ch], model=model, imgsz=imgsz)
 
-    def predict(self, inputs, profile=False):
+    def predict(self, inputs, tf_nms=False, agnostic_nms=False, topk_per_class=100, topk_all=100, iou_thres=0.45,
+                conf_thres=0.25):
         y = []  # outputs
         x = inputs
         for i, m in enumerate(self.model.layers):
@@ -356,18 +333,18 @@ class tf_Model():
             y.append(x if m.i in self.savelist else None)  # save output
 
         # Add TensorFlow NMS
-        if opt.tf_nms:
-            boxes = xywh2xyxy(x[0][..., :4])
+        if tf_nms:
+            boxes = self._xywh2xyxy(x[0][..., :4])
             probs = x[0][:, :, 4:5]
             classes = x[0][:, :, 5:]
             scores = probs * classes
-            if opt.agnostic_nms:
-                nms = agnostic_nms_layer()((boxes, classes, scores))
+            if agnostic_nms:
+                nms = AgnosticNMS()((boxes, classes, scores), topk_all, iou_thres, conf_thres)
                 return nms, x[1]
             else:
                 boxes = tf.expand_dims(boxes, 2)
                 nms = tf.image.combined_non_max_suppression(
-                    boxes, scores, opt.topk_per_class, opt.topk_all, opt.iou_thres, opt.score_thres, clip_boxes=False)
+                    boxes, scores, topk_per_class, topk_all, iou_thres, conf_thres, clip_boxes=False)
                 return nms, x[1]
 
         return x[0]  # output only first tensor [1,6300,85] = [xywh, conf, class0, class1, ...]
@@ -377,182 +354,94 @@ class tf_Model():
         # cls = tf.reshape(tf.cast(tf.argmax(x[..., 5:], axis=1), tf.float32), (-1, 1))  # x(6300,1)  classes
         # return tf.concat([conf, cls, xywh], 1)
 
+    @staticmethod
+    def _xywh2xyxy(xywh):
+        # Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
+        x, y, w, h = tf.split(xywh, num_or_size_splits=4, axis=-1)
+        return tf.concat([x - w / 2, y - h / 2, x + w / 2, y + h / 2], axis=-1)
+
 
-class agnostic_nms_layer(keras.layers.Layer):
-    # wrap map_fn to avoid TypeSpec related error https://stackoverflow.com/a/65809989/3036450
-    def call(self, input):
-        return tf.map_fn(agnostic_nms, input,
+class AgnosticNMS(keras.layers.Layer):
+    # TF Agnostic NMS
+    def call(self, input, topk_all, iou_thres, conf_thres):
+        # wrap map_fn to avoid TypeSpec related error https://stackoverflow.com/a/65809989/3036450
+        return tf.map_fn(self._nms, input,
                          fn_output_signature=(tf.float32, tf.float32, tf.float32, tf.int32),
                          name='agnostic_nms')
 
-
-def agnostic_nms(x):
-    boxes, classes, scores = x
-    class_inds = tf.cast(tf.argmax(classes, axis=-1), tf.float32)
-    scores_inp = tf.reduce_max(scores, -1)
-    selected_inds = tf.image.non_max_suppression(
-        boxes, scores_inp, max_output_size=opt.topk_all, iou_threshold=opt.iou_thres, score_threshold=opt.score_thres)
-    selected_boxes = tf.gather(boxes, selected_inds)
-    padded_boxes = tf.pad(selected_boxes,
-                          paddings=[[0, opt.topk_all - tf.shape(selected_boxes)[0]], [0, 0]],
-                          mode="CONSTANT", constant_values=0.0)
-    selected_scores = tf.gather(scores_inp, selected_inds)
-    padded_scores = tf.pad(selected_scores,
-                           paddings=[[0, opt.topk_all - tf.shape(selected_boxes)[0]]],
-                           mode="CONSTANT", constant_values=-1.0)
-    selected_classes = tf.gather(class_inds, selected_inds)
-    padded_classes = tf.pad(selected_classes,
-                            paddings=[[0, opt.topk_all - tf.shape(selected_boxes)[0]]],
-                            mode="CONSTANT", constant_values=-1.0)
-    valid_detections = tf.shape(selected_inds)[0]
-    return padded_boxes, padded_scores, padded_classes, valid_detections
-
-
-def xywh2xyxy(xywh):
-    # Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
-    x, y, w, h = tf.split(xywh, num_or_size_splits=4, axis=-1)
-    return tf.concat([x - w / 2, y - h / 2, x + w / 2, y + h / 2], axis=-1)
-
-
-def representative_dataset_gen():
-    # Representative dataset for use with converter.representative_dataset
-    n = 0
-    for path, img, im0s, vid_cap in dataset:
-        # Get sample input data as a numpy array in a method of your choosing.
-        n += 1
+    @staticmethod
+    def _nms(x, topk_all=100, iou_thres=0.45, conf_thres=0.25):  # agnostic NMS
+        boxes, classes, scores = x
+        class_inds = tf.cast(tf.argmax(classes, axis=-1), tf.float32)
+        scores_inp = tf.reduce_max(scores, -1)
+        selected_inds = tf.image.non_max_suppression(
+            boxes, scores_inp, max_output_size=topk_all, iou_threshold=iou_thres, score_threshold=conf_thres)
+        selected_boxes = tf.gather(boxes, selected_inds)
+        padded_boxes = tf.pad(selected_boxes,
+                              paddings=[[0, topk_all - tf.shape(selected_boxes)[0]], [0, 0]],
+                              mode="CONSTANT", constant_values=0.0)
+        selected_scores = tf.gather(scores_inp, selected_inds)
+        padded_scores = tf.pad(selected_scores,
+                               paddings=[[0, topk_all - tf.shape(selected_boxes)[0]]],
+                               mode="CONSTANT", constant_values=-1.0)
+        selected_classes = tf.gather(class_inds, selected_inds)
+        padded_classes = tf.pad(selected_classes,
+                                paddings=[[0, topk_all - tf.shape(selected_boxes)[0]]],
+                                mode="CONSTANT", constant_values=-1.0)
+        valid_detections = tf.shape(selected_inds)[0]
+        return padded_boxes, padded_scores, padded_classes, valid_detections
+
+
+def representative_dataset_gen(dataset, ncalib=100):
+    # Representative dataset generator for use with converter.representative_dataset, returns a generator of np arrays
+    for n, (path, img, im0s, vid_cap) in enumerate(dataset):
         input = np.transpose(img, [1, 2, 0])
         input = np.expand_dims(input, axis=0).astype(np.float32)
         input /= 255.0
         yield [input]
-        if n >= opt.ncalib:
+        if n >= ncalib:
             break
 
 
-if __name__ == "__main__":
+def run(weights=ROOT / 'yolov5s.pt',  # weights path
+        imgsz=(640, 640),  # inference size h,w
+        batch_size=1,  # batch size
+        dynamic=False,  # dynamic batch size
+        ):
+    # PyTorch model
+    im = torch.zeros((batch_size, 3, *imgsz))  # BCHW image
+    model = attempt_load(weights, map_location=torch.device('cpu'), inplace=True, fuse=False)
+    y = model(im)  # inference
+    model.info()
+
+    # TensorFlow model
+    im = tf.zeros((batch_size, *imgsz, 3))  # BHWC image
+    tf_model = TFModel(cfg=model.yaml, model=model, nc=model.nc, imgsz=imgsz)
+    y = tf_model.predict(im)  # inference
+
+    # Keras model
+    im = keras.Input(shape=(*imgsz, 3), batch_size=None if dynamic else batch_size)
+    keras_model = keras.Model(inputs=im, outputs=tf_model.predict(im))
+    keras_model.summary()
+
+
+def parse_opt():
     parser = argparse.ArgumentParser()
-    parser.add_argument('--cfg', type=str, default='yolov5s.yaml', help='cfg path')
-    parser.add_argument('--weights', type=str, default='yolov5s.pt', help='weights path')
-    parser.add_argument('--img-size', nargs='+', type=int, default=[320, 320], help='image size')  # height, width
+    parser.add_argument('--weights', type=str, default=ROOT / 'yolov5s.pt', help='weights path')
+    parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640], help='inference size h,w')
     parser.add_argument('--batch-size', type=int, default=1, help='batch size')
-    parser.add_argument('--dynamic-batch-size', action='store_true', help='dynamic batch size')
-    parser.add_argument('--source', type=str, default='../data/coco128.yaml', help='dir of images or data.yaml file')
-    parser.add_argument('--ncalib', type=int, default=100, help='number of calibration images')
-    parser.add_argument('--tfl-int8', action='store_true', dest='tfl_int8', help='export TFLite int8 model')
-    parser.add_argument('--tf-nms', action='store_true', dest='tf_nms', help='TF NMS (without TFLite export)')
-    parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
-    parser.add_argument('--tf-raw-resize', action='store_true', dest='tf_raw_resize',
-                        help='use tf.raw_ops.ResizeNearestNeighbor for resize')
-    parser.add_argument('--topk-per-class', type=int, default=100, help='topk per class to keep in NMS')
-    parser.add_argument('--topk-all', type=int, default=100, help='topk for all classes to keep in NMS')
-    parser.add_argument('--iou-thres', type=float, default=0.5, help='IOU threshold for NMS')
-    parser.add_argument('--score-thres', type=float, default=0.4, help='score threshold for NMS')
+    parser.add_argument('--dynamic', action='store_true', help='dynamic batch size')
     opt = parser.parse_args()
-    opt.cfg = check_yaml(opt.cfg)  # check YAML
-    opt.img_size *= 2 if len(opt.img_size) == 1 else 1  # expand
-    print(opt)
-
-    # Input
-    img = torch.zeros((opt.batch_size, 3, *opt.img_size))  # image size(1,3,320,192) iDetection
-
-    # Load PyTorch model
-    model = attempt_load(opt.weights, map_location=torch.device('cpu'), inplace=True, fuse=False)
-    model.model[-1].export = False  # set Detect() layer export=True
-    y = model(img)  # dry run
-    nc = y[0].shape[-1] - 5
-
-    # TensorFlow saved_model export
-    try:
-        print('\nStarting TensorFlow saved_model export with TensorFlow %s...' % tf.__version__)
-        tf_model = tf_Model(opt.cfg, model=model, nc=nc)
-        img = tf.zeros((opt.batch_size, *opt.img_size, 3))  # NHWC Input for TensorFlow
-
-        m = tf_model.model.layers[-1]
-        assert isinstance(m, tf_Detect), "the last layer must be Detect"
-        m.training = False
-        y = tf_model.predict(img)
-
-        inputs = keras.Input(shape=(*opt.img_size, 3), batch_size=None if opt.dynamic_batch_size else opt.batch_size)
-        keras_model = keras.Model(inputs=inputs, outputs=tf_model.predict(inputs))
-        keras_model.summary()
-        path = opt.weights.replace('.pt', '_saved_model')  # filename
-        keras_model.save(path, save_format='tf')
-        print('TensorFlow saved_model export success, saved as %s' % path)
-    except Exception as e:
-        print('TensorFlow saved_model export failure: %s' % e)
-        traceback.print_exc(file=sys.stdout)
-
-    # TensorFlow GraphDef export
-    try:
-        print('\nStarting TensorFlow GraphDef export with TensorFlow %s...' % tf.__version__)
-
-        # https://github.com/leimao/Frozen_Graph_TensorFlow
-        full_model = tf.function(lambda x: keras_model(x))
-        full_model = full_model.get_concrete_function(
-            tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype))
-
-        frozen_func = convert_variables_to_constants_v2(full_model)
-        frozen_func.graph.as_graph_def()
-        f = opt.weights.replace('.pt', '.pb')  # filename
-        tf.io.write_graph(graph_or_graph_def=frozen_func.graph,
-                          logdir=os.path.dirname(f),
-                          name=os.path.basename(f),
-                          as_text=False)
-
-        print('TensorFlow GraphDef export success, saved as %s' % f)
-    except Exception as e:
-        print('TensorFlow GraphDef export failure: %s' % e)
-        traceback.print_exc(file=sys.stdout)
-
-    # TFLite model export
-    if not opt.tf_nms:
-        try:
-            print('\nStarting TFLite export with TensorFlow %s...' % tf.__version__)
-
-            # fp32 TFLite model export ---------------------------------------------------------------------------------
-            # converter = tf.lite.TFLiteConverter.from_keras_model(keras_model)
-            # converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS]
-            # converter.allow_custom_ops = False
-            # converter.experimental_new_converter = True
-            # tflite_model = converter.convert()
-            # f = opt.weights.replace('.pt', '.tflite')  # filename
-            # open(f, "wb").write(tflite_model)
-
-            # fp16 TFLite model export ---------------------------------------------------------------------------------
-            converter = tf.lite.TFLiteConverter.from_keras_model(keras_model)
-            converter.optimizations = [tf.lite.Optimize.DEFAULT]
-            # converter.representative_dataset = representative_dataset_gen
-            # converter.target_spec.supported_types = [tf.float16]
-            converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS]
-            converter.allow_custom_ops = False
-            converter.experimental_new_converter = True
-            tflite_model = converter.convert()
-            f = opt.weights.replace('.pt', '-fp16.tflite')  # filename
-            open(f, "wb").write(tflite_model)
-            print('\nTFLite export success, saved as %s' % f)
-
-            # int8 TFLite model export ---------------------------------------------------------------------------------
-            if opt.tfl_int8:
-                # Representative Dataset
-                if opt.source.endswith('.yaml'):
-                    with open(check_yaml(opt.source)) as f:
-                        data = yaml.load(f, Loader=yaml.FullLoader)  # data dict
-                        check_dataset(data)  # check
-                    opt.source = data['train']
-                dataset = LoadImages(opt.source, img_size=opt.img_size, auto=False)
-                converter = tf.lite.TFLiteConverter.from_keras_model(keras_model)
-                converter.optimizations = [tf.lite.Optimize.DEFAULT]
-                converter.representative_dataset = representative_dataset_gen
-                converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
-                converter.inference_input_type = tf.uint8  # or tf.int8
-                converter.inference_output_type = tf.uint8  # or tf.int8
-                converter.allow_custom_ops = False
-                converter.experimental_new_converter = True
-                converter.experimental_new_quantizer = False
-                tflite_model = converter.convert()
-                f = opt.weights.replace('.pt', '-int8.tflite')  # filename
-                open(f, "wb").write(tflite_model)
-                print('\nTFLite (int8) export success, saved as %s' % f)
-
-        except Exception as e:
-            print('\nTFLite export failure: %s' % e)
-            traceback.print_exc(file=sys.stdout)
+    opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1  # expand
+    return opt
+
+
+def main(opt):
+    set_logging()
+    print(colorstr('tf.py: ') + ', '.join(f'{k}={v}' for k, v in vars(opt).items()))
+    run(**vars(opt))
+
+
+if __name__ == "__main__":
+    opt = parse_opt()
+    main(opt)

requirements.txt

@@ -1,6 +1,6 @@
 # pip install -r requirements.txt
 
-# base ----------------------------------------
+# Base ----------------------------------------
 matplotlib>=3.2.2
 numpy>=1.18.5
 opencv-python>=4.1.2
@@ -11,21 +11,23 @@ torch>=1.7.0
 torchvision>=0.8.1
 tqdm>=4.41.0
 
-# logging -------------------------------------
+# Logging -------------------------------------
 tensorboard>=2.4.1
 # wandb
 
-# plotting ------------------------------------
+# Plotting ------------------------------------
 seaborn>=0.11.0
 pandas
 
-# export --------------------------------------
-# coremltools>=4.1
-# onnx>=1.9.0
-# scikit-learn==0.19.2  # for coreml quantization
-# tensorflow==2.4.1  # for TFLite export
+# Export --------------------------------------
+# coremltools>=4.1  # CoreML export
+# onnx>=1.9.0  # ONNX export
+# onnx-simplifier>=0.3.6  # ONNX simplifier
+# scikit-learn==0.19.2  # CoreML quantization
+# tensorflow>=2.4.1  # TFLite export
+# tensorflowjs>=3.9.0  # TF.js export
 
-# extras --------------------------------------
+# Extras --------------------------------------
 # Cython  # for pycocotools https://github.com/cocodataset/cocoapi/issues/172
 # pycocotools>=2.0  # COCO mAP
 # albumentations>=1.0.3

utils/general.py

@@ -161,9 +161,15 @@ def emojis(str=''):
     return str.encode().decode('ascii', 'ignore') if platform.system() == 'Windows' else str
 
 
-def file_size(file):
-    # Return file size in MB
-    return Path(file).stat().st_size / 1e6
+def file_size(path):
+    # Return file/dir size (MB)
+    path = Path(path)
+    if path.is_file():
+        return path.stat().st_size / 1E6
+    elif path.is_dir():
+        return sum(f.stat().st_size for f in path.glob('**/*') if f.is_file()) / 1E6
+    else:
+        return 0.0
 
 
 def check_online():