From 42fd93b57aa6ec266a481350cb3f364f59084aff Mon Sep 17 00:00:00 2001
From: Wenxin Yang <595026238@qq.com>
Date: Thu, 26 Dec 2024 22:22:47 +0800
Subject: [PATCH] ultralytics.py

---
 sahi/models/ultralytics.py | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/sahi/models/ultralytics.py b/sahi/models/ultralytics.py
index dd421f5c..a9434b80 100644
--- a/sahi/models/ultralytics.py
+++ b/sahi/models/ultralytics.py
@@ -67,9 +67,11 @@ def perform_inference(self, image: np.ndarray):
 
         if self.image_size is not None:
             kwargs = {"imgsz": self.image_size, **kwargs}
+        if type(image) is list:
 
-        prediction_result = self.model(image[:, :, ::-1], **kwargs)  # YOLOv8 expects numpy arrays to have BGR
-
+            prediction_result = self.model(image, **kwargs)  # YOLOv8 expects numpy arrays to have BGR
+        else :
+            prediction_result = self.model(image[:, :, ::-1], **kwargs)
         if self.has_mask:
             if not prediction_result[0].masks:
                 prediction_result[0].masks = Masks(
@@ -109,7 +111,10 @@ def perform_inference(self, image: np.ndarray):
             prediction_result = [result.boxes.data for result in prediction_result]
 
         self._original_predictions = prediction_result
-        self._original_shape = image.shape
+        if type(image) == list:
+            self._original_shape = image[0].shape
+        else:
+            self._original_shape = image.shape
 
     @property
     def category_names(self):