Skip to content

Commit caac6e9

Browse files
Improve yolov8 config (#988)
Co-authored-by: fatih cagatay akyon <34196005+fcakyon@users.noreply.github.com>
1 parent f75f04a commit caac6e9

File tree

1 file changed

+9
-12
lines changed

1 file changed

+9
-12
lines changed

sahi/models/yolov8.py

Lines changed: 9 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -59,18 +59,15 @@ def perform_inference(self, image: np.ndarray):
5959
if self.model is None:
6060
raise ValueError("Model is not loaded, load it by calling .load_model()")
6161

62-
if self.image_size is not None: # ADDED IMAGE SIZE OPTION FOR YOLOV8 MODELS:
63-
prediction_result = self.model(
64-
image[:, :, ::-1], imgsz=self.image_size, verbose=False, device=self.device
65-
) # YOLOv8 expects numpy arrays to have BGR
66-
else:
67-
prediction_result = self.model(
68-
image[:, :, ::-1], verbose=False, device=self.device
69-
) # YOLOv8 expects numpy arrays to have BGR
70-
71-
prediction_result = [
72-
result.boxes.data[result.boxes.data[:, 4] >= self.confidence_threshold] for result in prediction_result
73-
]
62+
kwargs = {"cfg": self.config_path, "verbose": False, "conf": self.confidence_threshold, "device": self.device}
63+
64+
if self.image_size is not None:
65+
kwargs = {"imgsz": self.image_size, **kwargs}
66+
67+
prediction_result = self.model(image[:, :, ::-1], **kwargs) # YOLOv8 expects numpy arrays to have BGR
68+
69+
# We do not filter results again as confidence threshold is already applied above
70+
prediction_result = [result.boxes.data for result in prediction_result]
7471

7572
self._original_predictions = prediction_result
7673

0 commit comments

Comments
 (0)