From 354ea825e844264e6d9dc9d5e741b408e3be1f28 Mon Sep 17 00:00:00 2001
From: Francesco Mattioli <francesco.mttl@gmail.com>
Date: Wed, 6 Nov 2024 11:26:18 +0100
Subject: [PATCH] removed nms=False support

---
 ultralytics/engine/exporter.py | 95 ++++++++++++++++------------------
 1 file changed, 46 insertions(+), 49 deletions(-)

diff --git a/ultralytics/engine/exporter.py b/ultralytics/engine/exporter.py
index 010acc27a..b5878ff84 100644
--- a/ultralytics/engine/exporter.py
+++ b/ultralytics/engine/exporter.py
@@ -1106,7 +1106,8 @@ class Exporter:
             raise ValueError("MCT export is not supported for end2end models.")
         if "C2f" not in self.model.__str__():
             raise ValueError("MCT export is only supported for YOLOv8 detection models")
-        check_requirements("model-compression-toolkit==2.1.1")
+        check_requirements(("model-compression-toolkit==2.1.1", "sony-custom-layers[torch]"))
+        from sony_custom_layers.pytorch.object_detection.nms import multiclass_nms
         import model_compression_toolkit as mct
 
         def representative_dataset_gen(dataloader=self.get_int8_calibration_dataloader(prefix)):
@@ -1145,54 +1146,50 @@ class Exporter:
             )[0]
         )
 
-        if self.args.nms:
-            check_requirements("sony-custom-layers[torch]")
-            from sony_custom_layers.pytorch.object_detection.nms import multiclass_nms
-
-            class NMSWrapper(torch.nn.Module):
-                def __init__(
-                    self,
-                    model: torch.nn.Module,
-                    score_threshold: float = 0.001,
-                    iou_threshold: float = 0.7,
-                    max_detections: int = 300,
-                ):
-                    """
-                    Wrapping PyTorch Module with multiclass_nms layer from sony_custom_layers.
-
-                    Args:
-                        model (nn.Module): Model instance.
-                        score_threshold (float): Score threshold for non-maximum suppression.
-                        iou_threshold (float): Intersection over union threshold for non-maximum suppression.
-                        max_detections (float): The number of detections to return.
-                    """
-                    super().__init__()
-                    self.model = model
-                    self.score_threshold = score_threshold
-                    self.iou_threshold = iou_threshold
-                    self.max_detections = max_detections
-
-                def forward(self, images):
-                    # model inference
-                    outputs = self.model(images)
-
-                    boxes = outputs[0]
-                    scores = outputs[1]
-                    nms = multiclass_nms(
-                        boxes=boxes,
-                        scores=scores,
-                        score_threshold=self.score_threshold,
-                        iou_threshold=self.iou_threshold,
-                        max_detections=self.max_detections,
-                    )
-                    return nms
-
-            quant_model = NMSWrapper(
-                model=quant_model,
-                score_threshold=self.args.conf or 0.001,
-                iou_threshold=self.args.iou,
-                max_detections=self.args.max_det,
-            ).to(self.device)
+        class NMSWrapper(torch.nn.Module):
+            def __init__(
+                self,
+                model: torch.nn.Module,
+                score_threshold: float = 0.001,
+                iou_threshold: float = 0.7,
+                max_detections: int = 300,
+            ):
+                """
+                Wrapping PyTorch Module with multiclass_nms layer from sony_custom_layers.
+
+                Args:
+                    model (nn.Module): Model instance.
+                    score_threshold (float): Score threshold for non-maximum suppression.
+                    iou_threshold (float): Intersection over union threshold for non-maximum suppression.
+                    max_detections (float): The number of detections to return.
+                """
+                super().__init__()
+                self.model = model
+                self.score_threshold = score_threshold
+                self.iou_threshold = iou_threshold
+                self.max_detections = max_detections
+
+            def forward(self, images):
+                # model inference
+                outputs = self.model(images)
+
+                boxes = outputs[0]
+                scores = outputs[1]
+                nms = multiclass_nms(
+                    boxes=boxes,
+                    scores=scores,
+                    score_threshold=self.score_threshold,
+                    iou_threshold=self.iou_threshold,
+                    max_detections=self.max_detections,
+                )
+                return nms
+
+        quant_model = NMSWrapper(
+            model=quant_model,
+            score_threshold=self.args.conf or 0.001,
+            iou_threshold=self.args.iou,
+            max_detections=self.args.max_det,
+        ).to(self.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)