单目3D初始代码
This commit is contained in:
zhao.zhu
63
examples/YOLOv8-SAHI-Inference-Video/README.md
Executable file
63
examples/YOLOv8-SAHI-Inference-Video/README.md
Executable file
@@ -0,0 +1,63 @@
|
||||
# YOLO11 with SAHI for Video Inference
|
||||
|
||||
[Slicing Aided Hyper Inference (SAHI)](https://github.com/obss/sahi) is a powerful technique designed to optimize [object detection](https://en.wikipedia.org/wiki/Object_detection) algorithms, particularly for large-scale and [high-resolution imagery](https://en.wikipedia.org/wiki/Image_resolution). It works by partitioning images or video frames into manageable slices, performing detection on each slice using models like [Ultralytics YOLO11](https://docs.ultralytics.com/models/yolo11/), and then intelligently merging the results. This approach significantly improves detection accuracy for small objects and maintains performance on high-resolution inputs.
|
||||
|
||||
This tutorial guides you through running Ultralytics YOLO11 inference on video files using the SAHI library for enhanced detection capabilities. For a detailed guide on using SAHI with Ultralytics models, see the [SAHI Tiled Inference guide](https://docs.ultralytics.com/guides/sahi-tiled-inference/).
|
||||
|
||||
## 📋 Table of Contents
|
||||
|
||||
- [Step 1: Install Required Libraries](#-step-1-install-required-libraries)
|
||||
- [Step 2: Run Inference with SAHI using Ultralytics YOLO11](#-step-2-run-inference-with-sahi-using-ultralytics-yolo11)
|
||||
- [Usage Options](#-usage-options)
|
||||
- [Contribute](#-contribute)
|
||||
|
||||
## ⚙️ Step 1: Install Required Libraries
|
||||
|
||||
First, clone the [Ultralytics repository](https://github.com/ultralytics/ultralytics) to access the example script. Then, install the necessary [Python](https://www.python.org/) packages, including `sahi` and `ultralytics`, using [pip](https://pip.pypa.io/en/stable/). Finally, navigate into the example directory.
|
||||
|
||||
```bash
|
||||
# Clone the ultralytics repository
|
||||
git clone https://github.com/ultralytics/ultralytics
|
||||
|
||||
# Install dependencies
|
||||
# Ensure you have Python 3.8 or later installed
|
||||
pip install -U sahi ultralytics opencv-python
|
||||
|
||||
# Change directory to the example folder
|
||||
cd ultralytics/examples/YOLOv8-SAHI-Inference-Video
|
||||
```
|
||||
|
||||
## 🚀 Step 2: Run Inference with SAHI using Ultralytics YOLO11
|
||||
|
||||
Once the setup is complete, you can run inference on your video file. The provided script `yolov8_sahi.py` leverages SAHI for tiled inference with a specified YOLO11 model.
|
||||
|
||||
Execute the script using the command line, specifying the path to your video file. You can also choose different YOLO11 model weights.
|
||||
|
||||
```bash
|
||||
# Run inference and save the output video with bounding boxes
|
||||
python yolov8_sahi.py --source "path/to/your/video.mp4" --save-img
|
||||
|
||||
# Run inference using a specific YOLO11 model (e.g., yolo11n.pt) and save results
|
||||
python yolov8_sahi.py --source "path/to/your/video.mp4" --save-img --weights "yolo11n.pt"
|
||||
|
||||
# Run inference with smaller slices for potentially better small object detection
|
||||
python yolov8_sahi.py --source "path/to/your/video.mp4" --save-img --slice-height 512 --slice-width 512
|
||||
```
|
||||
|
||||
This script processes the video frame by frame, applying SAHI's slicing and inference logic. When saving is enabled, it exports annotated frames to `runs/detect/predict`. Learn more about prediction with Ultralytics models in the [Predict mode documentation](https://docs.ultralytics.com/modes/predict/).
|
||||
|
||||
## 🛠️ Usage Options
|
||||
|
||||
The script `yolov8_sahi.py` accepts several command-line arguments to customize the inference process:
|
||||
|
||||
- `--source`: **Required**. Path to the input video file (e.g., `"../path/to/video.mp4"`).
|
||||
- `--weights`: Optional. Path to the YOLO11 model weights file (e.g., `"yolo11n.pt"`, `"yolo11s.pt"`). Defaults to `"yolo11n.pt"`. You can download various models or use your custom-trained ones. See [Ultralytics YOLO models](https://docs.ultralytics.com/models/) for more options.
|
||||
- `--save-img`: Optional. Flag to export annotated frames. Saved to `runs/detect/predict`.
|
||||
- `--slice-height`: Optional. Height of each image slice for SAHI. Defaults to `1024`.
|
||||
- `--slice-width`: Optional. Width of each image slice for SAHI. Defaults to `1024`.
|
||||
|
||||
Experiment with these options, especially slice dimensions, to optimize detection performance for your specific [video processing](https://en.wikipedia.org/wiki/Video_processing) task and target object sizes. Using appropriate [datasets](https://docs.ultralytics.com/datasets/) for training can also significantly impact performance.
|
||||
|
||||
## ✨ Contribute
|
||||
|
||||
Contributions to enhance this example or add new features are welcome! If you encounter issues or have suggestions, please open an issue or submit a pull request in the [Ultralytics GitHub repository](https://github.com/ultralytics/ultralytics). Check out our [contribution guide](https://docs.ultralytics.com/help/contributing/) for more details.
|
||||
149
examples/YOLOv8-SAHI-Inference-Video/yolov8_sahi.py
Executable file
149
examples/YOLOv8-SAHI-Inference-Video/yolov8_sahi.py
Executable file
@@ -0,0 +1,149 @@
|
||||
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
import argparse
|
||||
import os
|
||||
|
||||
import cv2
|
||||
from sahi import AutoDetectionModel
|
||||
from sahi.predict import get_sliced_prediction
|
||||
from sahi.utils.ultralytics import download_model_weights
|
||||
|
||||
from ultralytics.utils.files import increment_path
|
||||
|
||||
|
||||
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: Load a YOLO11 model with specified weights for object detection using SAHI.
|
||||
inference: Run object detection on a video using YOLO11 and SAHI.
|
||||
parse_opt: Parse command line arguments for the inference process.
|
||||
|
||||
Examples:
|
||||
Initialize and run SAHI inference on a video
|
||||
>>> sahi_inference = SAHIInference()
|
||||
>>> sahi_inference.inference(weights="yolo11n.pt", source="video.mp4", view_img=True)
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize the SAHIInference class for performing sliced inference using SAHI with YOLO11 models."""
|
||||
self.detection_model = None
|
||||
|
||||
def load_model(self, weights: str, device: str) -> None:
|
||||
"""Load a YOLO11 model with specified weights for object detection using SAHI.
|
||||
|
||||
Args:
|
||||
weights (str): Path to the model weights file.
|
||||
device (str): CUDA device, i.e., '0' or '0,1,2,3' or 'cpu'.
|
||||
"""
|
||||
from ultralytics.utils.torch_utils import select_device
|
||||
|
||||
if weights and os.path.exists(weights):
|
||||
yolo11_model_path = weights
|
||||
else:
|
||||
yolo11_model_path = f"models/{weights}"
|
||||
download_model_weights(yolo11_model_path) # Download model if not present
|
||||
self.detection_model = AutoDetectionModel.from_pretrained(
|
||||
model_type="ultralytics", model_path=yolo11_model_path, device=select_device(device)
|
||||
)
|
||||
|
||||
def inference(
|
||||
self,
|
||||
weights: str = "yolo11n.pt",
|
||||
source: str = "test.mp4",
|
||||
view_img: bool = False,
|
||||
save_img: bool = False,
|
||||
exist_ok: bool = False,
|
||||
device: str = "",
|
||||
hide_conf: bool = False,
|
||||
slice_width: int = 512,
|
||||
slice_height: int = 512,
|
||||
) -> 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.
|
||||
device (str, optional): CUDA device, i.e., '0' or '0,1,2,3' or 'cpu'.
|
||||
hide_conf (bool, optional): Whether to hide confidence scores in the output.
|
||||
slice_width (int, optional): Slice width for inference.
|
||||
slice_height (int, optional): Slice height for inference.
|
||||
"""
|
||||
# Video setup
|
||||
cap = cv2.VideoCapture(source)
|
||||
if not cap.isOpened():
|
||||
raise FileNotFoundError(f"Unable to open video source: '{source}'")
|
||||
|
||||
save_dir = None
|
||||
if save_img:
|
||||
save_dir = increment_path("runs/detect/predict", exist_ok)
|
||||
save_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Load model
|
||||
self.load_model(weights, device)
|
||||
idx = 0 # Index for image frame writing
|
||||
while cap.isOpened():
|
||||
success, frame = cap.read()
|
||||
if not success:
|
||||
break
|
||||
|
||||
# Perform sliced prediction using SAHI
|
||||
results = get_sliced_prediction(
|
||||
frame[..., ::-1], # Convert BGR to RGB
|
||||
self.detection_model,
|
||||
slice_height=slice_height,
|
||||
slice_width=slice_width,
|
||||
)
|
||||
|
||||
# Display results if requested
|
||||
if view_img:
|
||||
cv2.imshow("Ultralytics YOLO Inference", frame)
|
||||
|
||||
# Save results if requested
|
||||
if save_img and save_dir is not None:
|
||||
idx += 1
|
||||
results.export_visuals(export_dir=save_dir, file_name=f"img_{idx}", hide_conf=hide_conf)
|
||||
|
||||
# Break loop if 'q' is pressed
|
||||
if cv2.waitKey(1) & 0xFF == ord("q"):
|
||||
break
|
||||
|
||||
# Clean up resources
|
||||
cap.release()
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
@staticmethod
|
||||
def parse_opt() -> 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")
|
||||
parser.add_argument("--device", default="", help="cuda device, i.e. 0 or 0,1,2,3 or cpu")
|
||||
parser.add_argument("--hide-conf", default=False, action="store_true", help="display or hide confidences")
|
||||
parser.add_argument("--slice-width", default=512, type=int, help="Slice width for inference")
|
||||
parser.add_argument("--slice-height", default=512, type=int, help="Slice height for inference")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
inference = SAHIInference()
|
||||
inference.inference(**vars(inference.parse_opt()))
|
||||
Reference in New Issue
Block a user