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

.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')

requirements.txt

 # 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():