Customize YOLOv8 image_size & device + Allow Saving Slices (#929)

Co-authored-by: fatih <[email protected]>
Co-authored-by: fcakyon <[email protected]>

README.md

@@ -267,4 +267,7 @@ python -m scripts.run_code_style format
 
 <a align="left" href="https://github.com/pranavdurai10" target="_blank">Pranav Durai</a>
 
+<a align="left" href="https://github.com/lakshaymehra" target="_blank">Lakshay Mehra</a>
+
 </div>
+

sahi/models/yolov8.py

@@ -58,7 +58,14 @@ def perform_inference(self, image: np.ndarray):
         # Confirm model is loaded
         if self.model is None:
             raise ValueError("Model is not loaded, load it by calling .load_model()")
-        prediction_result = self.model(image[:, :, ::-1], verbose=False)  # YOLOv8 expects numpy arrays to have BGR
+        if self.image_size is not None:  # ADDED IMAGE SIZE OPTION FOR YOLOV8 MODELS:
+            prediction_result = self.model(
+                image[:, :, ::-1], imgsz=self.image_size, verbose=False, device=self.device
+            )  # YOLOv8 expects numpy arrays to have BGR
+        else:
+            prediction_result = self.model(
+                image[:, :, ::-1], verbose=False, device=self.device
+            )  # YOLOv8 expects numpy arrays to have BGR
         prediction_result = [
             result.boxes.data[result.boxes.data[:, 4] >= self.confidence_threshold] for result in prediction_result
         ]

sahi/predict.py

@@ -125,6 +125,8 @@ def get_prediction(
 def get_sliced_prediction(
     image,
     detection_model=None,
+    output_file_name=None,  # ADDED OUTPUT FILE NAME TO (OPTIONALLY) SAVE SLICES
+    interim_dir="slices/",  # ADDED INTERIM DIRECTORY TO (OPTIONALLY) SAVE SLICES
     slice_height: int = None,
     slice_width: int = None,
     overlap_height_ratio: float = 0.2,
@@ -199,6 +201,8 @@ def get_sliced_prediction(
     time_start = time.time()
     slice_image_result = slice_image(
         image=image,
+        output_file_name=output_file_name,  # ADDED OUTPUT FILE NAME TO (OPTIONALLY) SAVE SLICES
+        output_dir=interim_dir,  # ADDED INTERIM DIRECTORY TO (OPTIONALLY) SAVE SLICES
         slice_height=slice_height,
         slice_width=slice_width,
         overlap_height_ratio=overlap_height_ratio,