From 92b52424d468feb48c51c3dde173d5d2c606a44b Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Sun, 18 Sep 2022 17:34:34 +0200
Subject: [PATCH] TensorFlow macOS AutoUpdate (#9471)

* TensorFlow macOS AutoUpdate

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 export.py        | 11 ++++++++---
 requirements.txt |  2 +-
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/export.py b/export.py
index ac9b13db8ec0..ae292afe06f6 100644
--- a/export.py
+++ b/export.py
@@ -72,6 +72,8 @@
                            check_yaml, colorstr, file_size, get_default_args, print_args, url2file, yaml_save)
 from utils.torch_utils import select_device, smart_inference_mode
 
+MACOS = platform.system() == 'Darwin'  # macOS environment
+
 
 def export_formats():
     # YOLOv5 export formats
@@ -224,7 +226,7 @@ def export_coreml(model, im, file, int8, half, prefix=colorstr('CoreML:')):
     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
+        if MACOS:  # 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)
@@ -310,8 +312,11 @@ def export_saved_model(model,
                        keras=False,
                        prefix=colorstr('TensorFlow SavedModel:')):
     # YOLOv5 TensorFlow SavedModel export
-    check_requirements('tensorflow' if torch.cuda.is_available() else 'tensorflow-cpu')
-    import tensorflow as tf
+    try:
+        import tensorflow as tf
+    except Exception:
+        check_requirements(f"tensorflow{'' if torch.cuda.is_available() else '-macos' if MACOS else '-cpu'}")
+        import tensorflow as tf
     from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
 
     from models.tf import TFModel
diff --git a/requirements.txt b/requirements.txt
index 44fe1ce697b7..835346f218a4 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -30,7 +30,7 @@ seaborn>=0.11.0
 # nvidia-pyindex  # TensorRT export
 # nvidia-tensorrt  # TensorRT export
 # scikit-learn==0.19.2  # CoreML quantization
-# tensorflow>=2.4.1  # TFLite export (or tensorflow-cpu, tensorflow-aarch64)
+# tensorflow>=2.4.1  # TF exports (-cpu, -aarch64, -macos)
 # tensorflowjs>=3.9.0  # TF.js export
 # openvino-dev  # OpenVINO export