make YOLO work with pyAIO

Author: Tomasz Stanislawski

computer_vision/object_detection/yolo_v8/run.py

@@ -97,30 +97,40 @@ def run_pytorch_fp(model_path, batch_size, num_runs, timeout, images_path, anno_
     # Ultralytics sets it to True by default. This way we suppress the logging by default while still allowing the user
     # to set it to True if needed
     from utils.pytorch import PyTorchRunner
-    from ultralytics.utils.nms import non_max_suppression
 
     def run_single_pass(pytorch_runner, coco):
         shape = (640, 640)
-        inp = torch.stack(coco.get_input_array(shape))
-        output = pytorch_runner.run(batch_size, inp)
-        output = non_max_suppression(output)
+        dset = coco.get_input_array(shape)
+        outputs = []
+        for inp in dset:
+            output, *_ = pytorch_runner.run(1, inp)
+            outputs.append(output)
+        assert len(outputs) == batch_size
 
         for i in range(batch_size):
-            for d in range(output[i].shape[0]):
+            for b in range(len(outputs[i].boxes)):
+                bbox = outputs[i].boxes.xywh[b]
+                cls = int(outputs[i].boxes.cls[b])
                 coco.submit_bbox_prediction(
                     i,
-                    coco.convert_bbox_to_coco_order(output[i][d][:4].tolist()),
-                    output[i][d][4].item(),
-                    coco.translate_cat_id_to_coco(output[i][d][5].item())
+                    coco.convert_bbox_to_coco_order(bbox),
+                    cls,
+                    coco.translate_cat_id_to_coco(cls)
                 )
 
     dataset = COCODataset(batch_size, "RGB", "COCO_val2014_000000000000", images_path,
                           anno_path, pre_processing="PyTorch_objdet", sort_ascending=True, order="NCHW")
 
     from ultralytics import YOLO
+    import numpy as np
     model = YOLO(model_path)
 
-    runner = PyTorchRunner(model.model,
+    # make sure that model.predictor exists
+    dummy_input = np.zeros((640, 640, 3), dtype=np.uint8)
+    model.predict(dummy_input)
+    assert model.predictor is not None
+
+    runner = PyTorchRunner(model,
                            disable_jit_freeze=disable_jit_freeze,
                            example_inputs=torch.stack(dataset.get_input_array((640, 640))))
 

utils/pytorch.py

@@ -72,8 +72,17 @@ def __init__(self,
                 # More natural comparison to version.parse("2.0") returns False for 2.0.0a0+git07156c4.dev, which is
                 # wrong. There was never a PyTorch 1.14, so this comparison acts like comparing to 2.0, but works
                 # correctly for such edge cases.
-                self._frozen_script = torch.compile(self._model, backend="aio-2" if AIO else "inductor",
-                                                    options={"modelname": self._model._get_name()} if AIO else {})
+                backend = "aio-2" if AIO else "inductor"
+                options = {"modelname": self._model._get_name()} if AIO else {}
+
+                if self._model._get_name() == "YOLO":
+                    # use compiled model for a prodictor model
+                    frozen_script = torch.compile(self._model.model, backend=backend, options=options)
+                    self._model.predictor.model.model = frozen_script
+                    self._frozen_script = self._model
+                else:
+                    self._frozen_script = torch.compile(self._model, backend=backend, options=options)
+
             elif os.environ.get("TORCH_COMPILE") == "1" and not version.parse(
                     pkg_resources.get_distribution("torch").version) >= version.parse("1.14"):
                 utils.print_goodbye_message_and_die(