Update SAHI example from YOLOv8 to YOLO11 (#18276)

Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
This commit is contained in:
Muhammad Rizwan Munawar 2024-12-17 21:29:36 +05:00 committed by GitHub
parent a3d807be13
commit ce50e33fa4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 39 additions and 32 deletions

View file

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