Update SAHI example from YOLOv8 to YOLO11 (#18276)
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
This commit is contained in:
parent
a3d807be13
commit
ce50e33fa4
4 changed files with 39 additions and 32 deletions
|
|
@ -6,32 +6,32 @@ from pathlib import Path
|
|||
import cv2
|
||||
from sahi import AutoDetectionModel
|
||||
from sahi.predict import get_sliced_prediction
|
||||
from sahi.utils.yolov8 import download_yolov8s_model
|
||||
from sahi.utils.ultralytics import download_yolo11n_model
|
||||
|
||||
from ultralytics.utils.files import increment_path
|
||||
from ultralytics.utils.plotting import Annotator, colors
|
||||
|
||||
|
||||
class SAHIInference:
|
||||
"""Runs YOLOv8 and SAHI for object detection on video with options to view, save, and track results."""
|
||||
"""Runs Ultralytics YOLO11 and SAHI for object detection on video with options to view, save, and track results."""
|
||||
|
||||
def __init__(self):
|
||||
"""Initializes the SAHIInference class for performing sliced inference using SAHI with YOLOv8 models."""
|
||||
"""Initializes the SAHIInference class for performing sliced inference using SAHI with YOLO11 models."""
|
||||
self.detection_model = None
|
||||
|
||||
def load_model(self, weights):
|
||||
"""Loads a YOLOv8 model with specified weights for object detection using SAHI."""
|
||||
yolov8_model_path = f"models/{weights}"
|
||||
download_yolov8s_model(yolov8_model_path)
|
||||
"""Loads a YOLO11 model with specified weights for object detection using SAHI."""
|
||||
yolo11_model_path = f"models/{weights}"
|
||||
download_yolo11n_model(yolo11_model_path)
|
||||
self.detection_model = AutoDetectionModel.from_pretrained(
|
||||
model_type="yolov8", model_path=yolov8_model_path, confidence_threshold=0.3, device="cpu"
|
||||
model_type="ultralytics", model_path=yolo11_model_path, confidence_threshold=0.3, device="cpu"
|
||||
)
|
||||
|
||||
def inference(
|
||||
self, weights="yolov8n.pt", source="test.mp4", view_img=False, save_img=False, exist_ok=False, track=False
|
||||
self, weights="yolo11n.pt", source="test.mp4", view_img=False, save_img=False, exist_ok=False, track=False
|
||||
):
|
||||
"""
|
||||
Run object detection on a video using YOLOv8 and SAHI.
|
||||
Run object detection on a video using YOLO11 and SAHI.
|
||||
|
||||
Args:
|
||||
weights (str): Model weights path.
|
||||
|
|
@ -93,7 +93,7 @@ class SAHIInference:
|
|||
def parse_opt(self):
|
||||
"""Parse command line arguments."""
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--weights", type=str, default="yolov8n.pt", help="initial weights path")
|
||||
parser.add_argument("--weights", type=str, default="yolo11n.pt", help="initial weights path")
|
||||
parser.add_argument("--source", type=str, required=True, help="video file path")
|
||||
parser.add_argument("--view-img", action="store_true", help="show results")
|
||||
parser.add_argument("--save-img", action="store_true", help="save results")
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue