# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
import argparse
from pathlib import Path
import cv2
from sahi import AutoDetectionModel
from sahi.predict import get_sliced_prediction
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 Ultralytics YOLO11 and SAHI for object detection on video with options to view, save, and track results.
This class integrates SAHI (Slicing Aided Hyper Inference) with YOLO11 models to perform efficient object detection
on large images by slicing them into smaller pieces, running inference on each slice, and then merging the results.
Attributes:
detection_model (AutoDetectionModel): The loaded YOLO11 model wrapped with SAHI functionality.
Methods:
load_model: Loads a YOLO11 model with specified weights.
inference: Runs object detection on a video using the loaded model.
parse_opt: Parses command line arguments for the inference process.
"""
def __init__(self):
"""Initializes the SAHIInference class for performing sliced inference using SAHI with YOLO11 models."""
self.detection_model = None
def load_model(self, weights: str) -> None:
"""
Load a YOLO11 model with specified weights for object detection using SAHI.
Args:
weights (str): Path to the model weights file.
"""
yolo11_model_path = f"models/{weights}"
download_yolo11n_model(yolo11_model_path) # Download model if not present
self.detection_model = AutoDetectionModel.from_pretrained(
model_type="ultralytics", model_path=yolo11_model_path, device="cpu"
)
def inference(
self,
weights: str = "yolo11n.pt",
source: str = "test.mp4",
view_img: bool = False,
save_img: bool = False,
exist_ok: bool = False,
) -> None:
"""
Run object detection on a video using YOLO11 and SAHI.
The function processes each frame of the video, applies sliced inference using SAHI,
and optionally displays and/or saves the results with bounding boxes and labels.
Args:
weights (str): Model weights path.
source (str): Video file path.
view_img (bool): Whether to display results in a window.
save_img (bool): Whether to save results to a video file.
exist_ok (bool): Whether to overwrite existing output files.
"""
# Video setup
cap = cv2.VideoCapture(source)
assert cap.isOpened(), "Error reading video file"
frame_width, frame_height = int(cap.get(3)), int(cap.get(4))
# Output setup
save_dir = increment_path(Path("ultralytics_results_with_sahi") / "exp", exist_ok)
save_dir.mkdir(parents=True, exist_ok=True)
video_writer = cv2.VideoWriter(
str(save_dir / f"{Path(source).stem}.avi"),
cv2.VideoWriter_fourcc(*"MJPG"),
int(cap.get(5)),
(frame_width, frame_height),
)
# Load model
self.load_model(weights)
while cap.isOpened():
success, frame = cap.read()
if not success:
break
annotator = Annotator(frame) # Initialize annotator for plotting detection results
# Perform sliced prediction using SAHI
results = get_sliced_prediction(
frame[..., ::-1], # Convert BGR to RGB
self.detection_model,
slice_height=512,
slice_width=512,
)
# Extract detection data from results
detection_data = [
(det.category.name, det.category.id, (det.bbox.minx, det.bbox.miny, det.bbox.maxx, det.bbox.maxy))
for det in results.object_prediction_list
]
# Annotate frame with detection results
for det in detection_data:
annotator.box_label(det[2], label=str(det[0]), color=colors(int(det[1]), True))
# Display results if requested
if view_img:
cv2.imshow(Path(source).stem, frame)
# Save results if requested
if save_img:
video_writer.write(frame)
# Break loop if 'q' is pressed
if cv2.waitKey(1) & 0xFF == ord("q"):
break
# Clean up resources
video_writer.release()
cap.release()
cv2.destroyAllWindows()
def parse_opt(self) -> argparse.Namespace:
"""
Parse command line arguments for the inference process.
Returns:
(argparse.Namespace): Parsed command line arguments.
"""
parser = argparse.ArgumentParser()
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")
parser.add_argument("--exist-ok", action="store_true", help="existing project/name ok, do not increment")
return parser.parse_args()
if __name__ == "__main__":
inference = SAHIInference()
inference.inference(**vars(inference.parse_opt()))
注释翻译成中文
博客提及在Python、Matlab和R中进行多元函数的绘制，这属于信息技术领域中数据分析与可视化相关内容，借助这些工具可对多元函数进行图形展示。
扫一扫
4076
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
