From eb6a8d0c31cdb9255ae8e496090d58f8e22a50b7 Mon Sep 17 00:00:00 2001
From: Francesco Mattioli <francesco.mttl@gmail.com>
Date: Mon, 21 Oct 2024 11:48:37 +0200
Subject: [PATCH] updated mct version, added imx500 dependency

---
 ultralytics/engine/exporter.py | 40 +++++++++++++++-------------------
 ultralytics/nn/autobackend.py  |  2 +-
 2 files changed, 19 insertions(+), 23 deletions(-)

diff --git a/ultralytics/engine/exporter.py b/ultralytics/engine/exporter.py
index a42e0bb431..eb41f841ae 100644
--- a/ultralytics/engine/exporter.py
+++ b/ultralytics/engine/exporter.py
@@ -1054,11 +1054,10 @@ class Exporter:
 
     @try_export
     def export_mct(self, prefix=colorstr("Sony MCT:")):
-        check_requirements(["mct-nightly", "sony-custom-layers[torch]"])
+        check_requirements(["model_compression_toolkit==2.1.0", "sony-custom-layers[torch]"])
         import model_compression_toolkit as mct
         import onnx
-        from model_compression_toolkit.core import BitWidthConfig
-        from model_compression_toolkit.core.common.network_editors import NodeNameScopeFilter
+        import subprocess
         from model_compression_toolkit.core.pytorch.pytorch_device_config import get_working_device, set_working_device
         from sony_custom_layers.pytorch.object_detection.nms import multiclass_nms
 
@@ -1109,25 +1108,12 @@ class Exporter:
                 yield [img]
 
         tpc = mct.get_target_platform_capabilities(
-            fw_name="pytorch", target_platform_name="imx500", target_platform_version="v3"
-        )
-
-        # Configure MCT manually for specific layers
-        bit_cfg = BitWidthConfig()
-        bit_cfg.set_manual_activation_bit_width(
-            [
-                NodeNameScopeFilter("mul"),
-                NodeNameScopeFilter("sub"),
-                NodeNameScopeFilter("add_6"),
-                NodeNameScopeFilter("cat_17"),
-            ],
-            16,
+            fw_name="pytorch", target_platform_name="imx500", target_platform_version="v1"
         )
 
         config = mct.core.CoreConfig(
             mixed_precision_config=mct.core.MixedPrecisionQuantizationConfig(num_of_images=10),
-            quantization_config=mct.core.QuantizationConfig(concat_threshold_update=True),
-            bit_width_config=bit_cfg,
+            quantization_config=mct.core.QuantizationConfig(concat_threshold_update=True)
         )
 
         resource_utilization = mct.core.ResourceUtilization(weights_memory=3146176 * 0.76)
@@ -1141,7 +1127,6 @@ class Exporter:
                 core_config=config,
                 target_platform_capabilities=tpc,
             )
-            print("Quantized model is ready")
 
         else:
             gptq_config = mct.gptq.get_pytorch_gptq_config(n_epochs=1000, use_hessian_based_weights=False)
@@ -1157,8 +1142,6 @@ class Exporter:
                 target_platform_capabilities=tpc,
             )
 
-            print("Quantized-PTQ model is ready")
-
         if self.args.nms:
             # Define PostProcess params
             score_threshold = 0.001
@@ -1171,7 +1154,7 @@ class Exporter:
                 iou_threshold=iou_threshold,
                 max_detections=max_detections,
             ).to(device=get_working_device())
-
+            
         f = Path(str(self.file).replace(self.file.suffix, "_mct_model.onnx"))  # js dir
         mct.exporter.pytorch_export_model(model=quant_model, save_model_path=f, repr_dataset=representative_dataset_gen)
 
@@ -1181,6 +1164,19 @@ class Exporter:
             meta.key, meta.value = k, str(v)
 
         onnx.save(model_onnx, f)
+        
+        if not LINUX:
+            LOGGER.warning(f"{prefix} WARNING ⚠️ MCT imx500-converter is only supported on Linux.")
+        else:
+            check_requirements("imx500-converter[pt]==3.14.1")
+            try:
+                subprocess.run(["java", "--version"], check=True)
+            except FileNotFoundError:
+                LOGGER.error("Java 17 is required for the imx500 conversion. \n Please install Java with: \n sudo apt install openjdk-17-jdk openjdk-17-jre")
+                return None
+        
+            subprocess.run(["imxconv-pt", "-i", "yolov8n_mct_model.onnx", "-o", "yolov8n_imx500_model"], check=True)
+        
         return f, None
 
     def _add_tflite_metadata(self, file):
diff --git a/ultralytics/nn/autobackend.py b/ultralytics/nn/autobackend.py
index 8e36d54450..755df474a8 100644
--- a/ultralytics/nn/autobackend.py
+++ b/ultralytics/nn/autobackend.py
@@ -178,7 +178,7 @@ class AutoBackend(nn.Module):
 
             providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] if cuda else ["CPUExecutionProvider"]
             if mct:
-                check_requirements(["mct-nightly", "sony-custom-layers[torch]"])
+                check_requirements(["model_compression_toolkit==2.1.0", "sony-custom-layers[torch]", "onnxruntime-extensions"])
                 LOGGER.info(f"Loading {w} for ONNX MCT quantization inference...")
                 import mct_quantizers as mctq
                 from sony_custom_layers.pytorch.object_detection import nms_ort  # noqa