@@ -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