update to 28.04.22 ultralytics/yolov5 (#110)

* update to 28.04.22 ultralytics/yolov5 * add missing arg to train * fix imports * update load * fix docstring * fix docstring * fix train.py * fix fire format for export * fix a log * update key points in readme
fcakyon · Apr 28, 2022 · b9766c8 · b9766c8
1 parent bbdad77
commit b9766c8
Show file tree

Hide file tree

Showing 27 changed files with 356 additions and 239 deletions.
diff --git a/README.md b/README.md
@@ -20,9 +20,12 @@
 
 <div align="center">
 You can finally install <a href="https://github.com/ultralytics/yolov5">YOLOv5 object detector</a> using <a href="https://pypi.org/project/yolov5/">pip</a> and integrate into your project easily.
+
 <img src="https://user-images.githubusercontent.com/26833433/136901921-abcfcd9d-f978-4942-9b97-0e3f202907df.png" width="1000">
+</div>
+
 <br>
-This yolov5 package contains everything from ultralytics/yolov5 <a href="https://github.com/ultralytics/yolov5/commit/ea72b84f5e690cb516642ce2d9ae200145b0af34">at this commit</a> plus:
+This yolov5 package contains everything from ultralytics/yolov5 <a href="https://github.com/ultralytics/yolov5/tree/177da7f348181abdf4820ad26707eb8b3dd4fdc9">at this commit</a> plus:
 <br>
 1. Easy installation via pip: `pip install yolov5`
 <br>
@@ -35,7 +38,7 @@ This yolov5 package contains everything from ultralytics/yolov5 <a href="https:/
 5. <a href="https://neptune.ai/">NeptuneAI</a> logger support (metric, model and dataset logging)
 <br>
 6. Classwise AP logging during experiments
-</div>
+
 
 
 ## <div align="center">Install</div>

diff --git a/tests/test_yolov5.py b/tests/test_yolov5.py
@@ -2,28 +2,26 @@
 
 from tests.test_utils import TestConstants
 
+DEVICE = "cpu"
 
 class TestYolov5(unittest.TestCase):
     def test_load_model(self):
         from yolov5 import YOLOv5
 
         # init model
         model_path = TestConstants.YOLOV5S_MODEL_PATH
-        device = "cpu"
-        yolov5 = YOLOv5(model_path, device, load_on_init=False)
+        yolov5 = YOLOv5(model_path, DEVICE, load_on_init=False)
         yolov5.load_model()
 
         # check if loaded
         self.assertNotEqual(yolov5.model, None)
 
     def test_load_model_on_init(self):
-        from PIL import Image
         from yolov5 import YOLOv5
 
         # init model
         model_path = TestConstants.YOLOV5S_MODEL_PATH
-        device = "cpu"
-        yolov5 = YOLOv5(model_path, device, load_on_init=True)
+        yolov5 = YOLOv5(model_path, DEVICE, load_on_init=True)
 
         # check if loaded
         self.assertNotEqual(yolov5.model, None)
@@ -34,8 +32,7 @@ def test_predict(self):
 
         # init yolov5s model
         model_path = TestConstants.YOLOV5S_MODEL_PATH
-        device = "cpu"
-        yolov5 = YOLOv5(model_path, device, load_on_init=True)
+        yolov5 = YOLOv5(model_path, DEVICE, load_on_init=True)
 
         # prepare image
         image_path = TestConstants.ZIDANE_IMAGE_PATH
@@ -62,8 +59,7 @@ def test_predict(self):
 
         # init yolov5l model
         model_path = TestConstants.YOLOV5L_MODEL_PATH
-        device = "cpu"
-        yolov5 = YOLOv5(model_path, device, load_on_init=True)
+        yolov5 = YOLOv5(model_path, DEVICE, load_on_init=True)
 
         # prepare image
         image_path = TestConstants.BUS_IMAGE_PATH
@@ -89,8 +85,7 @@ def test_predict(self):
 
         # init yolov5s model
         model_path = TestConstants.YOLOV5S_MODEL_PATH
-        device = "cpu"
-        yolov5 = YOLOv5(model_path, device, load_on_init=True)
+        yolov5 = YOLOv5(model_path, DEVICE, load_on_init=True)
 
         # prepare images
         image_path1 = TestConstants.ZIDANE_IMAGE_PATH
@@ -127,7 +122,7 @@ def test_hublike_load_model(self):
 
         # init model
         model_path = TestConstants.YOLOV5S_MODEL_PATH
-        model = yolov5.load(model_path)
+        model = yolov5.load(model_path, device=DEVICE)
 
         # check if loaded
         self.assertNotEqual(model, None)
@@ -138,7 +133,8 @@ def test_hublike_predict(self):
 
         # init yolov5s model
         model_path = TestConstants.YOLOV5S_MODEL_PATH
-        model = yolov5.load(model_path)
+        #model_path = '/home/fatihakyon/dev/obss/ml/drone-tracking/drone-detection/[coco-bb]_33k_yolov5m6_1344input/weights/best.pt'
+        model = yolov5.load(model_path, device=DEVICE)
 
         # prepare image
         image_path = TestConstants.ZIDANE_IMAGE_PATH
@@ -154,7 +150,7 @@ def test_hublike_predict(self):
 
         # init yolov5l model
         model_path = TestConstants.YOLOV5L_MODEL_PATH
-        model = yolov5.load(model_path)
+        model = yolov5.load(model_path, device=DEVICE)
 
         # prepare image
         image_path = TestConstants.BUS_IMAGE_PATH
@@ -169,7 +165,7 @@ def test_hublike_predict(self):
 
         # init yolov5s model
         model_path = TestConstants.YOLOV5S_MODEL_PATH
-        model = yolov5.load(model_path)
+        model = yolov5.load(model_path, device=DEVICE)
 
         # prepare images
         image_path1 = TestConstants.ZIDANE_IMAGE_PATH

diff --git a/yolov5/__init__.py b/yolov5/__init__.py
@@ -1,4 +1,4 @@
 from yolov5.helpers import YOLOv5
 from yolov5.helpers import load_model as load
 
-__version__ = "6.1.1"
+__version__ = "6.1.2"
diff --git a/yolov5/detect.py b/yolov5/detect.py
@@ -17,7 +17,7 @@
                                          yolov5s.onnx               # ONNX Runtime or OpenCV DNN with --dnn
                                          yolov5s.xml                # OpenVINO
                                          yolov5s.engine             # TensorRT
-                                         yolov5s.mlmodel            # CoreML (MacOS-only)
+                                         yolov5s.mlmodel            # CoreML (macOS-only)
                                          yolov5s_saved_model        # TensorFlow SavedModel
                                          yolov5s.pb                 # TensorFlow GraphDef
                                          yolov5s.tflite             # TensorFlow Lite
@@ -34,7 +34,6 @@
 
 FILE = Path(__file__).resolve()
 ROOT = FILE.parents[0]  # YOLOv5 root directory
-
 ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative
 
 from yolov5.models.common import DetectMultiBackend

diff --git a/yolov5/export.py b/yolov5/export.py
@@ -29,7 +29,7 @@
                         yolov5s.onnx               # ONNX Runtime or OpenCV DNN with --dnn
                         yolov5s.xml                # OpenVINO
                         yolov5s.engine             # TensorRT
-                        yolov5s.mlmodel            # CoreML (MacOS-only)
+                        yolov5s.mlmodel            # CoreML (macOS-only)
                         yolov5s_saved_model        # TensorFlow SavedModel
                         yolov5s.pb                 # TensorFlow GraphDef
                         yolov5s.tflite             # TensorFlow Lite
@@ -139,7 +139,13 @@ def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorst
         # Checks
         model_onnx = onnx.load(f)  # load onnx model
         onnx.checker.check_model(model_onnx)  # check onnx model
-        # LOGGER.info(onnx.helper.printable_graph(model_onnx.graph))  # print
+
+        # Metadata
+        d = {'stride': int(max(model.stride)), 'names': model.names}
+        for k, v in d.items():
+            meta = model_onnx.metadata_props.add()
+            meta.key, meta.value = k, str(v)
+        onnx.save(model_onnx, f)
 
         # Simplify
         if simplify:
@@ -161,7 +167,7 @@ def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorst
         LOGGER.info(f'{prefix} export failure: {e}')
 
 
-def export_openvino(model, im, file, prefix=colorstr('OpenVINO:')):
+def export_openvino(model, im, file, half, prefix=colorstr('OpenVINO:')):
     # YOLOv5 OpenVINO export
     try:
         check_requirements(('openvino-dev',))  # requires openvino-dev: https://pypi.org/project/openvino-dev/
@@ -170,7 +176,7 @@ def export_openvino(model, im, file, prefix=colorstr('OpenVINO:')):
         LOGGER.info(f'\n{prefix} starting export with openvino {ie.__version__}...')
         f = str(file).replace('.pt', '_openvino_model' + os.sep)
 
-        cmd = f"mo --input_model {file.with_suffix('.onnx')} --output_dir {f}"
+        cmd = f"mo --input_model {file.with_suffix('.onnx')} --output_dir {f} --data_type {'FP16' if half else 'FP32'}"
         subprocess.check_output(cmd, shell=True)
 
         LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
@@ -179,7 +185,7 @@ def export_openvino(model, im, file, prefix=colorstr('OpenVINO:')):
         LOGGER.info(f'\n{prefix} export failure: {e}')
 
 
-def export_coreml(model, im, file, prefix=colorstr('CoreML:')):
+def export_coreml(model, im, file, int8, half, prefix=colorstr('CoreML:')):
     # YOLOv5 CoreML export
     try:
         check_requirements(('coremltools',))
@@ -190,6 +196,14 @@ def export_coreml(model, im, file, prefix=colorstr('CoreML:')):
 
         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, bias=[0, 0, 0])])
+        bits, mode = (8, 'kmeans_lut') if int8 else (16, 'linear') if half else (32, None)
+        if bits < 32:
+            if platform.system() == 'Darwin':  # quantization only supported on macOS
+                with warnings.catch_warnings():
+                    warnings.filterwarnings("ignore", category=DeprecationWarning)  # suppress numpy==1.20 float warning
+                    ct_model = ct.models.neural_network.quantization_utils.quantize_weights(ct_model, bits, mode)
+            else:
+                print(f'{prefix} quantization only supported on macOS, skipping...')
         ct_model.save(f)
 
         LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
@@ -202,7 +216,9 @@ def export_coreml(model, im, file, prefix=colorstr('CoreML:')):
 def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=False, prefix=colorstr('TensorRT:')):
     # YOLOv5 TensorRT export https://developer.nvidia.com/tensorrt
     try:
-        check_requirements(('tensorrt',))
+        assert im.device.type != 'cpu', 'export running on CPU but must be on GPU, i.e. `python export.py --device 0`'
+        if platform.system() == 'Linux':
+            check_requirements(('nvidia-tensorrt',), cmds=('-U --index-url https://pypi.ngc.nvidia.com',))
         import tensorrt as trt
 
         if trt.__version__[0] == '7':  # TensorRT 7 handling https://github.com/ultralytics/yolov5/issues/6012
@@ -216,7 +232,6 @@ def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=F
         onnx = file.with_suffix('.onnx')
 
         LOGGER.info(f'\n{prefix} starting export with TensorRT {trt.__version__}...')
-        assert im.device.type != 'cpu', 'export running on CPU but must be on GPU, i.e. `python export.py --device 0`'
         assert onnx.exists(), f'failed to export ONNX file: {onnx}'
         f = file.with_suffix('.engine')  # TensorRT engine file
         logger = trt.Logger(trt.Logger.INFO)
@@ -328,7 +343,7 @@ def export_pb(keras_model, im, file, prefix=colorstr('TensorFlow GraphDef:')):
         LOGGER.info(f'\n{prefix} export failure: {e}')
 
 
-def export_tflite(keras_model, im, file, int8, data, ncalib, prefix=colorstr('TensorFlow Lite:')):
+def export_tflite(keras_model, im, file, int8, data, nms, agnostic_nms, prefix=colorstr('TensorFlow Lite:')):
     # YOLOv5 TensorFlow Lite export
     try:
         import tensorflow as tf
@@ -344,13 +359,15 @@ def export_tflite(keras_model, im, file, int8, data, ncalib, prefix=colorstr('Te
         if int8:
             from models.tf import representative_dataset_gen
             dataset = LoadImages(check_dataset(data)['train'], img_size=imgsz, auto=False)  # representative data
-            converter.representative_dataset = lambda: representative_dataset_gen(dataset, ncalib)
+            converter.representative_dataset = lambda: representative_dataset_gen(dataset, ncalib=100)
             converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
             converter.target_spec.supported_types = []
             converter.inference_input_type = tf.uint8  # or tf.int8
             converter.inference_output_type = tf.uint8  # or tf.int8
             converter.experimental_new_quantizer = True
             f = str(file).replace('.pt', '-int8.tflite')
+        if nms or agnostic_nms:
+            converter.target_spec.supported_ops.append(tf.lite.OpsSet.SELECT_TF_OPS)
 
         tflite_model = converter.convert()
         open(f, "wb").write(tflite_model)
@@ -380,7 +397,7 @@ def export_edgetpu(keras_model, im, file, prefix=colorstr('Edge TPU:')):
         f = str(file).replace('.pt', '-int8_edgetpu.tflite')  # Edge TPU model
         f_tfl = str(file).replace('.pt', '-int8.tflite')  # TFLite model
 
-        cmd = f"edgetpu_compiler -s {f_tfl}"
+        cmd = f"edgetpu_compiler -s -o {file.parent} {f_tfl}"
         subprocess.run(cmd, shell=True, check=True)
 
         LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
@@ -392,7 +409,7 @@ def export_edgetpu(keras_model, im, file, prefix=colorstr('Edge TPU:')):
 def export_tfjs(keras_model, im, file, prefix=colorstr('TensorFlow.js:')):
     # YOLOv5 TensorFlow.js export
     try:
-        #check_requirements(('tensorflowjs',))
+        check_requirements(('tensorflowjs',))
         import re
 
         import tensorflowjs as tfjs
@@ -406,7 +423,8 @@ def export_tfjs(keras_model, im, file, prefix=colorstr('TensorFlow.js:')):
               f'--output_node_names="Identity,Identity_1,Identity_2,Identity_3" {f_pb} {f}'
         subprocess.run(cmd, shell=True)
 
-        json = open(f_json).read()
+        with open(f_json) as j:
+            json = j.read()
         with open(f_json, 'w') as j:  # sort JSON Identity_* in ascending order
             subst = re.sub(
                 r'{"outputs": {"Identity.?.?": {"name": "Identity.?.?"}, '
@@ -464,13 +482,13 @@ def run(
 
     # 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'
+    if half:
+        assert device.type != 'cpu' or coreml or xml, '--half only compatible with GPU export, i.e. use --device 0'
     model = attempt_load(weights, map_location=device, inplace=True, fuse=True)  # load FP32 model
     nc, names = model.nc, model.names  # number of classes, class names
 
     # Checks
     imgsz *= 2 if len(imgsz) == 1 else 1  # expand
-    opset = 12 if ('openvino' in include) else opset  # OpenVINO requires opset <= 12
     assert nc == len(names), f'Model class count {nc} != len(names) {len(names)}'
 
     # Input
@@ -479,19 +497,14 @@ def run(
     im = torch.zeros(batch_size, 3, *imgsz).to(device)  # image size(1,3,320,192) BCHW iDetection
 
     # Update model
-    if half:
+    if half and not (coreml or xml):
         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.SiLU):
-        #         m.act = SiLU()
         if isinstance(m, Detect):
             m.inplace = inplace
             m.onnx_dynamic = dynamic
             m.export = True
-            if hasattr(m, 'forward_export'):
-                m.forward = m.forward_export  # assign custom forward (optional)
 
     for _ in range(2):
         y = model(im)  # dry runs
@@ -508,9 +521,9 @@ def run(
     if onnx or xml:  # OpenVINO requires ONNX
         f[2] = export_onnx(model, im, file, opset, train, dynamic, simplify)
     if xml:  # OpenVINO
-        f[3] = export_openvino(model, im, file)
+        f[3] = export_openvino(model, im, file, half)
     if coreml:
-        _, f[4] = export_coreml(model, im, file)
+        _, f[4] = export_coreml(model, im, file, int8, half)
 
     # TensorFlow Exports
     if any((saved_model, pb, tflite, edgetpu, tfjs)):
@@ -530,7 +543,7 @@ def run(
         if pb or tfjs:  # pb prerequisite to tfjs
             f[6] = export_pb(model, im, file)
         if tflite or edgetpu:
-            f[7] = export_tflite(model, im, file, int8=int8 or edgetpu, data=data, ncalib=100)
+            f[7] = export_tflite(model, im, file, int8=int8 or edgetpu, data=data, nms=nms, agnostic_nms=agnostic_nms)
         if edgetpu:
             f[8] = export_edgetpu(model, im, file)
         if tfjs:
@@ -541,9 +554,9 @@ def run(
     if any(f):
         LOGGER.info(f'\nExport complete ({time.time() - t:.2f}s)'
                     f"\nResults saved to {colorstr('bold', file.parent.resolve())}"
-                    f"\nDetect:          python detect.py --weights {f[-1]}"
-                    f"\nPyTorch Hub:     model = torch.hub.load('ultralytics/yolov5', 'custom', '{f[-1]}')"
-                    f"\nValidate:        python val.py --weights {f[-1]}"
+                    f"\nDetect:          yolov5 detect --weights {f[-1]}"
+                    f"\n\PIP Package:    model = yolov5.load('{f[-1]}')"
+                    f"\nValidate:        yolov5 val --weights {f[-1]}"
                     f"\nVisualize:       https://netron.app")
     return f  # return list of exported files/dirs